001/**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.commons.crypto.stream;
019
020import java.io.IOException;
021import java.nio.ByteBuffer;
022import java.security.GeneralSecurityException;
023import java.util.Properties;
024import java.util.Queue;
025import java.util.concurrent.ConcurrentLinkedQueue;
026
027import javax.crypto.Cipher;
028import javax.crypto.spec.IvParameterSpec;
029
030import org.apache.commons.crypto.cipher.CryptoCipher;
031import org.apache.commons.crypto.cipher.CryptoCipherFactory;
032import org.apache.commons.crypto.stream.input.Input;
033import org.apache.commons.crypto.utils.IoUtils;
034import org.apache.commons.crypto.utils.Utils;
035
036/**
037 * PositionedCryptoInputStream provides the capability to decrypt the stream
038 * starting at random position as well as provides the foundation for positioned
039 * read for decrypting. This needs a stream cipher mode such as AES CTR mode.
040 */
041public class PositionedCryptoInputStream extends CtrCryptoInputStream {
042
043    /**
044     * DirectBuffer pool
045     */
046    private final Queue<ByteBuffer> bufferPool = new ConcurrentLinkedQueue<>();
047
048    /**
049     * CryptoCipher pool
050     */
051    private final Queue<CipherState> cipherPool = new ConcurrentLinkedQueue<>();
052
053    /**
054     * properties for constructing a CryptoCipher
055     */
056    private final Properties properties;
057
058    /**
059     * Constructs a {@link PositionedCryptoInputStream}.
060     *
061     * @param properties The {@code Properties} class represents a set of
062     *        properties.
063     * @param in the input data.
064     * @param key crypto key for the cipher.
065     * @param iv Initialization vector for the cipher.
066     * @param streamOffset the start offset in the data.
067     * @throws IOException if an I/O error occurs.
068     */
069    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by PositionedCryptoInputStream.
070    public PositionedCryptoInputStream(final Properties properties, final Input in, final byte[] key,
071            final byte[] iv, final long streamOffset) throws IOException {
072        this(properties, in, Utils.getCipherInstance("AES/CTR/NoPadding", properties),
073                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
074    }
075
076    /**
077     * Constructs a {@link PositionedCryptoInputStream}.
078     *
079     * @param properties the properties of stream
080     * @param input the input data.
081     * @param cipher the CryptoCipher instance.
082     * @param bufferSize the bufferSize.
083     * @param key crypto key for the cipher.
084     * @param iv Initialization vector for the cipher.
085     * @param streamOffset the start offset in the data.
086     * @throws IOException if an I/O error occurs.
087     */
088    protected PositionedCryptoInputStream(final Properties properties, final Input input, final CryptoCipher cipher,
089            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
090            throws IOException {
091        super(input, cipher, bufferSize, key, iv, streamOffset);
092        this.properties = properties;
093    }
094
095    /**
096     * Reads up to the specified number of bytes from a given position within a
097     * stream and return the number of bytes read. This does not change the
098     * current offset of the stream, and is thread-safe.
099     *
100     * @param buffer the buffer into which the data is read.
101     * @param length the maximum number of bytes to read.
102     * @param offset the start offset in the data.
103     * @param position the offset from the start of the stream.
104     * @throws IOException if an I/O error occurs.
105     * @return int the total number of decrypted data bytes read into the
106     *         buffer.
107     */
108    public int read(final long position, final byte[] buffer, final int offset, final int length)
109            throws IOException {
110        checkStream();
111        final int n = input.read(position, buffer, offset, length);
112        if (n > 0) {
113            // This operation does not change the current offset of the file
114            decrypt(position, buffer, offset, n);
115        }
116        return n;
117    }
118
119    /**
120     * Reads the specified number of bytes from a given position within a
121     * stream. This does not change the current offset of the stream and is
122     * thread-safe.
123     *
124     * @param buffer the buffer into which the data is read.
125     * @param length the maximum number of bytes to read.
126     * @param offset the start offset in the data.
127     * @param position the offset from the start of the stream.
128     * @throws IOException if an I/O error occurs.
129     */
130    public void readFully(final long position, final byte[] buffer, final int offset, final int length)
131            throws IOException {
132        checkStream();
133        IoUtils.readFully(input, position, buffer, offset, length);
134        if (length > 0) {
135            // This operation does not change the current offset of the file
136            decrypt(position, buffer, offset, length);
137        }
138    }
139
140    /**
141     * Reads the specified number of bytes from a given position within a
142     * stream. This does not change the current offset of the stream and is
143     * thread-safe.
144     *
145     * @param position the offset from the start of the stream.
146     * @param buffer the buffer into which the data is read.
147     * @throws IOException if an I/O error occurs.
148     */
149    public void readFully(final long position, final byte[] buffer) throws IOException {
150        readFully(position, buffer, 0, buffer.length);
151    }
152
153    /**
154     * Decrypts length bytes in buffer starting at offset. Output is also put
155     * into buffer starting at offset. It is thread-safe.
156     *
157     * @param buffer the buffer into which the data is read.
158     * @param offset the start offset in the data.
159     * @param position the offset from the start of the stream.
160     * @param length the maximum number of bytes to read.
161     * @throws IOException if an I/O error occurs.
162     */
163    protected void decrypt(final long position, final byte[] buffer, final int offset, final int length)
164            throws IOException {
165        final ByteBuffer inByteBuffer = getBuffer();
166        final ByteBuffer outByteBuffer = getBuffer();
167        CipherState state = null;
168        try {
169            state = getCipherState();
170            final byte[] iv = getInitIV().clone();
171            resetCipher(state, position, iv);
172            byte padding = getPadding(position);
173            inByteBuffer.position(padding); // Set proper position for input data.
174
175            int n = 0;
176            while (n < length) {
177                final int toDecrypt = Math.min(length - n, inByteBuffer.remaining());
178                inByteBuffer.put(buffer, offset + n, toDecrypt);
179
180                // Do decryption
181                decrypt(state, inByteBuffer, outByteBuffer, padding);
182
183                outByteBuffer.get(buffer, offset + n, toDecrypt);
184                n += toDecrypt;
185                padding = postDecryption(state, inByteBuffer, position + n, iv);
186            }
187        } finally {
188            returnBuffer(inByteBuffer);
189            returnBuffer(outByteBuffer);
190            returnCipherState(state);
191        }
192    }
193
194    /**
195     * Does the decryption using inBuffer as input and outBuffer as output. Upon
196     * return, inBuffer is cleared; the decrypted data starts at
197     * outBuffer.position() and ends at outBuffer.limit().
198     *
199     * @param state the CipherState instance.
200     * @param inByteBuffer the input buffer.
201     * @param outByteBuffer the output buffer.
202     * @param padding the padding.
203     * @throws IOException if an I/O error occurs.
204     */
205    private void decrypt(final CipherState state, final ByteBuffer inByteBuffer,
206            final ByteBuffer outByteBuffer, final byte padding) throws IOException {
207        Utils.checkState(inByteBuffer.position() >= padding);
208        if (inByteBuffer.position() == padding) {
209            // There is no real data in inBuffer.
210            return;
211        }
212        inByteBuffer.flip();
213        outByteBuffer.clear();
214        decryptBuffer(state, inByteBuffer, outByteBuffer);
215        inByteBuffer.clear();
216        outByteBuffer.flip();
217        if (padding > 0) {
218            /*
219             * The plain text and cipher text have a 1:1 mapping, they start at
220             * the same position.
221             */
222            outByteBuffer.position(padding);
223        }
224    }
225
226    /**
227     * Does the decryption using inBuffer as input and outBuffer as output.
228     *
229     * @param state the CipherState instance.
230     * @param inByteBuffer the input buffer.
231     * @param outByteBuffer the output buffer.
232     * @throws IOException if an I/O error occurs.
233     */
234    private void decryptBuffer(final CipherState state, final ByteBuffer inByteBuffer,
235            final ByteBuffer outByteBuffer) throws IOException {
236        final int inputSize = inByteBuffer.remaining();
237        try {
238            final int n = state.getCryptoCipher().update(inByteBuffer, outByteBuffer);
239            if (n < inputSize) {
240                /**
241                 * Typically code will not get here. CryptoCipher#update will
242                 * consume all input data and put result in outBuffer.
243                 * CryptoCipher#doFinal will reset the cipher context.
244                 */
245                state.getCryptoCipher().doFinal(inByteBuffer, outByteBuffer);
246                state.reset(true);
247            }
248        } catch (final GeneralSecurityException e) {
249            throw new IOException(e);
250        }
251    }
252
253    /**
254     * This method is executed immediately after decryption. Check whether
255     * cipher should be updated and recalculate padding if needed.
256     *
257     * @param state the CipherState instance.
258     * @param inByteBuffer the input buffer.
259     * @param position the offset from the start of the stream.
260     * @param iv the iv.
261     * @return the padding.
262     * @throws IOException if an I/O error occurs.
263     */
264    private byte postDecryption(final CipherState state, final ByteBuffer inByteBuffer,
265            final long position, final byte[] iv) throws IOException {
266        byte padding = 0;
267        if (state.isReset()) {
268            /*
269             * This code is generally not executed since the cipher usually
270             * maintains cipher context (e.g. the counter) internally. However,
271             * some implementations can't maintain context so a re-init is
272             * necessary after each decryption call.
273             */
274            resetCipher(state, position, iv);
275            padding = getPadding(position);
276            inByteBuffer.position(padding);
277        }
278        return padding;
279    }
280
281    /**
282     * Calculates the counter and iv, reset the cipher.
283     *
284     * @param state the CipherState instance.
285     * @param position the offset from the start of the stream.
286     * @param iv the iv.
287     * @throws IOException if an I/O error occurs.
288     */
289    private void resetCipher(final CipherState state, final long position, final byte[] iv)
290            throws IOException {
291        final long counter = getCounter(position);
292        CtrCryptoInputStream.calculateIV(getInitIV(), counter, iv);
293        try {
294            state.getCryptoCipher().init(Cipher.DECRYPT_MODE, key,
295                    new IvParameterSpec(iv));
296        } catch (final GeneralSecurityException e) {
297        }
298        state.reset(false);
299    }
300
301    /**
302     * Gets CryptoCipher from pool.
303     *
304     * @return the CipherState instance.
305     * @throws IOException if an I/O error occurs.
306     */
307    private CipherState getCipherState() throws IOException {
308        CipherState state = cipherPool.poll();
309        if (state == null) {
310            CryptoCipher cryptoCipher;
311            try {
312                cryptoCipher = CryptoCipherFactory.getCryptoCipher("AES/CTR/NoPadding", properties);
313            } catch (final GeneralSecurityException e) {
314                throw new IOException(e);
315            }
316            state = new CipherState(cryptoCipher);
317        }
318
319        return state;
320    }
321
322    /**
323     * Returns CryptoCipher to pool.
324     *
325     * @param state the CipherState instance.
326     */
327    private void returnCipherState(final CipherState state) {
328        if (state != null) {
329            cipherPool.add(state);
330        }
331    }
332
333    /**
334     * Gets direct buffer from pool.
335     *
336     * @return the buffer.
337     */
338    private ByteBuffer getBuffer() {
339        ByteBuffer buffer = bufferPool.poll();
340        if (buffer == null) {
341            buffer = ByteBuffer.allocateDirect(getBufferSize());
342        }
343
344        return buffer;
345    }
346
347    /**
348     * Returns direct buffer to pool.
349     *
350     * @param buf the buffer.
351     */
352    private void returnBuffer(final ByteBuffer buf) {
353        if (buf != null) {
354            buf.clear();
355            bufferPool.add(buf);
356        }
357    }
358
359    /**
360     * Overrides the {@link CryptoInputStream#close()}. Closes this input stream
361     * and releases any system resources associated with the stream.
362     *
363     * @throws IOException if an I/O error occurs.
364     */
365    @Override
366    public void close() throws IOException {
367        if (!isOpen()) {
368            return;
369        }
370
371        cleanBufferPool();
372        super.close();
373    }
374
375    /** Clean direct buffer pool */
376    private void cleanBufferPool() {
377        ByteBuffer buf;
378        while ((buf = bufferPool.poll()) != null) {
379            CryptoInputStream.freeDirectBuffer(buf);
380        }
381    }
382
383    private static class CipherState {
384
385        private final CryptoCipher cryptoCipher;
386        private boolean reset;
387
388        /**
389         * The constructor of {@link CipherState}.
390         *
391         * @param cipher the CryptoCipher instance.
392         */
393        public CipherState(final CryptoCipher cipher) {
394            this.cryptoCipher = cipher;
395            this.reset = false;
396        }
397
398        /**
399         * Gets the CryptoCipher instance.
400         *
401         * @return the cipher.
402         */
403        public CryptoCipher getCryptoCipher() {
404            return cryptoCipher;
405        }
406
407        /**
408         * Gets the reset.
409         *
410         * @return the value of reset.
411         */
412        public boolean isReset() {
413            return reset;
414        }
415
416        /**
417         * Sets the value of reset.
418         *
419         * @param reset the reset.
420         */
421        public void reset(final boolean reset) {
422            this.reset = reset;
423        }
424    }
425}