1 // Copyright 2016 Google Inc. All rights reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 package com.google.archivepatcher.shared;
16 
17 import java.io.Closeable;
18 import java.io.IOException;
19 import java.io.InputStream;
20 import java.io.OutputStream;
21 
22 /**
23  * A pipe that moves data from an {@link InputStream} to an {@link OutputStream}, optionally
24  * uncompressing the input data on-the-fly.
25  */
26 public class PartiallyUncompressingPipe implements Closeable {
27   /**
28    * The uncompressor used to uncompress compressed input streams.
29    */
30   private final DeflateUncompressor uncompressor;
31 
32   /**
33    * The output stream to write to.
34    */
35   private final CountingOutputStream out;
36 
37   /**
38    * A buffer used when copying bytes.
39    */
40   private final byte[] copyBuffer;
41 
42   /**
43    * Modes available for {@link PartiallyUncompressingPipe#pipe(InputStream, Mode)}.
44    */
45   public static enum Mode {
46     /**
47      * Copy bytes form the {@link InputStream} to the {@link OutputStream} without modification.
48      */
49     COPY,
50 
51     /**
52      * Treat the {@link InputStream} as a deflate stream with nowrap=false, uncompress the bytes
53      * on-the-fly and write the uncompressed data to the {@link OutputStream}.
54      */
55     UNCOMPRESS_WRAPPED,
56 
57     /**
58      * Treat the {@link InputStream} as a deflate stream with nowrap=true, uncompress the bytes
59      * on-the-fly and write the uncompressed data to the {@link OutputStream}.
60      */
61     UNCOMPRESS_NOWRAP,
62   }
63 
64   /**
65    * Constructs a new stream.
66    * @param out the stream, to write to
67    * @param copyBufferSize the size of the buffer to use when copying instead of uncompressing
68    */
PartiallyUncompressingPipe(OutputStream out, int copyBufferSize)69   public PartiallyUncompressingPipe(OutputStream out, int copyBufferSize) {
70     this.out = new CountingOutputStream(out);
71     uncompressor = new DeflateUncompressor();
72     uncompressor.setCaching(true);
73     copyBuffer = new byte[copyBufferSize];
74   }
75 
76   /**
77    * Pipes the entire contents of the specified {@link InputStream} to the configured
78    * {@link OutputStream}, optionally uncompressing on-the-fly.
79    * @param in the stream to read from
80    * @param mode the mode to use for reading and writing
81    * @return the number of bytes written to the output stream
82    * @throws IOException if anything goes wrong
83    */
pipe(InputStream in, Mode mode)84   public long pipe(InputStream in, Mode mode) throws IOException {
85     long bytesWrittenBefore = out.getNumBytesWritten();
86     if (mode == Mode.COPY) {
87       int numRead = 0;
88       while ((numRead = in.read(copyBuffer)) >= 0) {
89         out.write(copyBuffer, 0, numRead);
90       }
91     } else {
92       uncompressor.setNowrap(mode == Mode.UNCOMPRESS_NOWRAP);
93       uncompressor.uncompress(in, out);
94     }
95     out.flush();
96     return out.getNumBytesWritten() - bytesWrittenBefore;
97   }
98 
99   /**
100    * Returns the number of bytes written to the stream so far.
101    * @return as described
102    */
getNumBytesWritten()103   public long getNumBytesWritten() {
104     return out.getNumBytesWritten();
105   }
106 
107   @Override
close()108   public void close() throws IOException {
109     uncompressor.release();
110     out.close();
111   }
112 }
113