1/*
2 * Copyright (C) 2014 Square, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16package okio;
17
18import java.io.EOFException;
19import java.io.IOException;
20import java.util.zip.DataFormatException;
21import java.util.zip.Inflater;
22
23/**
24 * A source that uses <a href="http://tools.ietf.org/html/rfc1951">DEFLATE</a>
25 * to decompress data read from another source.
26 */
27public final class InflaterSource implements Source {
28  private final BufferedSource source;
29  private final Inflater inflater;
30
31  /**
32   * When we call Inflater.setInput(), the inflater keeps our byte array until
33   * it needs input again. This tracks how many bytes the inflater is currently
34   * holding on to.
35   */
36  private int bufferBytesHeldByInflater;
37  private boolean closed;
38
39  public InflaterSource(Source source, Inflater inflater) {
40    this(Okio.buffer(source), inflater);
41  }
42
43  /**
44   * This package-private constructor shares a buffer with its trusted caller.
45   * In general we can't share a BufferedSource because the inflater holds input
46   * bytes until they are inflated.
47   */
48  InflaterSource(BufferedSource source, Inflater inflater) {
49    if (source == null) throw new IllegalArgumentException("source == null");
50    if (inflater == null) throw new IllegalArgumentException("inflater == null");
51    this.source = source;
52    this.inflater = inflater;
53  }
54
55  @Override public long read(
56      OkBuffer sink, long byteCount) throws IOException {
57    if (byteCount < 0) throw new IllegalArgumentException("byteCount < 0: " + byteCount);
58    if (closed) throw new IllegalStateException("closed");
59    if (byteCount == 0) return 0;
60
61    while (true) {
62      boolean sourceExhausted = refill();
63
64      // Decompress the inflater's compressed data into the sink.
65      try {
66        Segment tail = sink.writableSegment(1);
67        int bytesInflated = inflater.inflate(tail.data, tail.limit, Segment.SIZE - tail.limit);
68        if (bytesInflated > 0) {
69          tail.limit += bytesInflated;
70          sink.size += bytesInflated;
71          return bytesInflated;
72        }
73        if (inflater.finished() || inflater.needsDictionary()) {
74          releaseInflatedBytes();
75          return -1;
76        }
77        if (sourceExhausted) throw new EOFException("source exhausted prematurely");
78      } catch (DataFormatException e) {
79        throw new IOException(e);
80      }
81    }
82  }
83
84  /**
85   * Refills the inflater with compressed data if it needs input. (And only if
86   * it needs input). Returns true if the inflater required input but the source
87   * was exhausted.
88   */
89  public boolean refill() throws IOException {
90    if (!inflater.needsInput()) return false;
91
92    releaseInflatedBytes();
93    if (inflater.getRemaining() != 0) throw new IllegalStateException("?"); // TODO: possible?
94
95    // If there are compressed bytes in the source, assign them to the inflater.
96    if (source.exhausted()) return true;
97
98    // Assign buffer bytes to the inflater.
99    Segment head = source.buffer().head;
100    bufferBytesHeldByInflater = head.limit - head.pos;
101    inflater.setInput(head.data, head.pos, bufferBytesHeldByInflater);
102    return false;
103  }
104
105  /** When the inflater has processed compressed data, remove it from the buffer. */
106  private void releaseInflatedBytes() throws IOException {
107    if (bufferBytesHeldByInflater == 0) return;
108    int toRelease = bufferBytesHeldByInflater - inflater.getRemaining();
109    bufferBytesHeldByInflater -= toRelease;
110    source.skip(toRelease);
111  }
112
113  @Override public Source deadline(Deadline deadline) {
114    source.deadline(deadline);
115    return this;
116  }
117
118  @Override public void close() throws IOException {
119    if (closed) return;
120    inflater.end();
121    closed = true;
122    source.close();
123  }
124}
125