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      Buffer 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          if (tail.pos == tail.limit) {
76            // We allocated a tail segment, but didn't end up needing it. Recycle!
77            sink.head = tail.pop();
78            SegmentPool.recycle(tail);
79          }
80          return -1;
81        }
82        if (sourceExhausted) throw new EOFException("source exhausted prematurely");
83      } catch (DataFormatException e) {
84        throw new IOException(e);
85      }
86    }
87  }
88
89  /**
90   * Refills the inflater with compressed data if it needs input. (And only if
91   * it needs input). Returns true if the inflater required input but the source
92   * was exhausted.
93   */
94  public boolean refill() throws IOException {
95    if (!inflater.needsInput()) return false;
96
97    releaseInflatedBytes();
98    if (inflater.getRemaining() != 0) throw new IllegalStateException("?"); // TODO: possible?
99
100    // If there are compressed bytes in the source, assign them to the inflater.
101    if (source.exhausted()) return true;
102
103    // Assign buffer bytes to the inflater.
104    Segment head = source.buffer().head;
105    bufferBytesHeldByInflater = head.limit - head.pos;
106    inflater.setInput(head.data, head.pos, bufferBytesHeldByInflater);
107    return false;
108  }
109
110  /** When the inflater has processed compressed data, remove it from the buffer. */
111  private void releaseInflatedBytes() throws IOException {
112    if (bufferBytesHeldByInflater == 0) return;
113    int toRelease = bufferBytesHeldByInflater - inflater.getRemaining();
114    bufferBytesHeldByInflater -= toRelease;
115    source.skip(toRelease);
116  }
117
118  @Override public Timeout timeout() {
119    return source.timeout();
120  }
121
122  @Override public void close() throws IOException {
123    if (closed) return;
124    inflater.end();
125    closed = true;
126    source.close();
127  }
128}
129