Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.commons.crypto.stream;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.util.Properties;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;

import org.apache.commons.crypto.cipher.CryptoCipher;
import org.apache.commons.crypto.cipher.CryptoCipherFactory;
import org.apache.commons.crypto.stream.input.Input;
import org.apache.commons.crypto.utils.IoUtils;
import org.apache.commons.crypto.utils.Utils;

PositionedCryptoInputStream provides the capability to decrypt the stream
starting at random position as well as provides the foundation for positioned
read for decrypting. This needs a stream cipher mode such as AES CTR mode.
/**
 * PositionedCryptoInputStream provides the capability to decrypt the stream
 * starting at random position as well as provides the foundation for positioned
 * read for decrypting. This needs a stream cipher mode such as AES CTR mode.
 */
public class PositionedCryptoInputStream extends CtrCryptoInputStream {

    
DirectBuffer pool
/**
     * DirectBuffer pool
     */
    private final Queue<ByteBuffer> bufferPool = new ConcurrentLinkedQueue<>();

    
CryptoCipher pool
/**
     * CryptoCipher pool
     */
    private final Queue<CipherState> cipherPool = new ConcurrentLinkedQueue<>();

    
properties for constructing a CryptoCipher
/**
     * properties for constructing a CryptoCipher
     */
    private final Properties properties;

    
Constructs a PositionedCryptoInputStream. 
Params:
properties – The Properties class represents a set of properties.
in – the input data.
key – crypto key for the cipher.
iv – Initialization vector for the cipher.
streamOffset – the start offset in the data.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link PositionedCryptoInputStream}.
     *
     * @param properties The {@code Properties} class represents a set of
     *        properties.
     * @param in the input data.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the data.
     * @throws IOException if an I/O error occurs.
     */
    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by PositionedCryptoInputStream.
    public PositionedCryptoInputStream(final Properties properties, final Input in, final byte[] key,
            final byte[] iv, final long streamOffset) throws IOException {
        this(properties, in, Utils.getCipherInstance("AES/CTR/NoPadding", properties),
                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
    }

    
Constructs a PositionedCryptoInputStream. 
Params:
properties – the properties of stream
input – the input data.
cipher – the CryptoCipher instance.
bufferSize – the bufferSize.
key – crypto key for the cipher.
iv – Initialization vector for the cipher.
streamOffset – the start offset in the data.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link PositionedCryptoInputStream}.
     *
     * @param properties the properties of stream
     * @param input the input data.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the data.
     * @throws IOException if an I/O error occurs.
     */
    protected PositionedCryptoInputStream(final Properties properties, final Input input, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
            throws IOException {
        super(input, cipher, bufferSize, key, iv, streamOffset);
        this.properties = properties;
    }

    
Reads up to the specified number of bytes from a given position within a
stream and return the number of bytes read. This does not change the
current offset of the stream, and is thread-safe.
Params:
buffer – the buffer into which the data is read.
length – the maximum number of bytes to read.
offset – the start offset in the data.
position – the offset from the start of the stream.
Throws:
IOException – if an I/O error occurs.
Returns:
int the total number of decrypted data bytes read into the
        buffer./**
     * Reads up to the specified number of bytes from a given position within a
     * stream and return the number of bytes read. This does not change the
     * current offset of the stream, and is thread-safe.
     *
     * @param buffer the buffer into which the data is read.
     * @param length the maximum number of bytes to read.
     * @param offset the start offset in the data.
     * @param position the offset from the start of the stream.
     * @throws IOException if an I/O error occurs.
     * @return int the total number of decrypted data bytes read into the
     *         buffer.
     */
    public int read(final long position, final byte[] buffer, final int offset, final int length)
            throws IOException {
        checkStream();
        final int n = input.read(position, buffer, offset, length);
        if (n > 0) {
            // This operation does not change the current offset of the file
            decrypt(position, buffer, offset, n);
        }
        return n;
    }

    
Reads the specified number of bytes from a given position within a
stream. This does not change the current offset of the stream and is
thread-safe.
Params:
buffer – the buffer into which the data is read.
length – the maximum number of bytes to read.
offset – the start offset in the data.
position – the offset from the start of the stream.
Throws:
IOException – if an I/O error occurs./**
     * Reads the specified number of bytes from a given position within a
     * stream. This does not change the current offset of the stream and is
     * thread-safe.
     *
     * @param buffer the buffer into which the data is read.
     * @param length the maximum number of bytes to read.
     * @param offset the start offset in the data.
     * @param position the offset from the start of the stream.
     * @throws IOException if an I/O error occurs.
     */
    public void readFully(final long position, final byte[] buffer, final int offset, final int length)
            throws IOException {
        checkStream();
        IoUtils.readFully(input, position, buffer, offset, length);
        if (length > 0) {
            // This operation does not change the current offset of the file
            decrypt(position, buffer, offset, length);
        }
    }

    
Reads the specified number of bytes from a given position within a
stream. This does not change the current offset of the stream and is
thread-safe.
Params:
position – the offset from the start of the stream.
buffer – the buffer into which the data is read.
Throws:
IOException – if an I/O error occurs./**
     * Reads the specified number of bytes from a given position within a
     * stream. This does not change the current offset of the stream and is
     * thread-safe.
     *
     * @param position the offset from the start of the stream.
     * @param buffer the buffer into which the data is read.
     * @throws IOException if an I/O error occurs.
     */
    public void readFully(final long position, final byte[] buffer) throws IOException {
        readFully(position, buffer, 0, buffer.length);
    }

    
Decrypts length bytes in buffer starting at offset. Output is also put
into buffer starting at offset. It is thread-safe.
Params:
buffer – the buffer into which the data is read.
offset – the start offset in the data.
position – the offset from the start of the stream.
length – the maximum number of bytes to read.
Throws:
IOException – if an I/O error occurs./**
     * Decrypts length bytes in buffer starting at offset. Output is also put
     * into buffer starting at offset. It is thread-safe.
     *
     * @param buffer the buffer into which the data is read.
     * @param offset the start offset in the data.
     * @param position the offset from the start of the stream.
     * @param length the maximum number of bytes to read.
     * @throws IOException if an I/O error occurs.
     */
    protected void decrypt(final long position, final byte[] buffer, final int offset, final int length)
            throws IOException {
        final ByteBuffer inByteBuffer = getBuffer();
        final ByteBuffer outByteBuffer = getBuffer();
        CipherState state = null;
        try {
            state = getCipherState();
            final byte[] iv = getInitIV().clone();
            resetCipher(state, position, iv);
            byte padding = getPadding(position);
            inByteBuffer.position(padding); // Set proper position for input data.

            int n = 0;
            while (n < length) {
                final int toDecrypt = Math.min(length - n, inByteBuffer.remaining());
                inByteBuffer.put(buffer, offset + n, toDecrypt);

                // Do decryption
                decrypt(state, inByteBuffer, outByteBuffer, padding);

                outByteBuffer.get(buffer, offset + n, toDecrypt);
                n += toDecrypt;
                padding = postDecryption(state, inByteBuffer, position + n, iv);
            }
        } finally {
            returnBuffer(inByteBuffer);
            returnBuffer(outByteBuffer);
            returnCipherState(state);
        }
    }

    
Does the decryption using inBuffer as input and outBuffer as output. Upon
return, inBuffer is cleared; the decrypted data starts at
outBuffer.position() and ends at outBuffer.limit().
Params:
state – the CipherState instance.
inByteBuffer – the input buffer.
outByteBuffer – the output buffer.
padding – the padding.
Throws:
IOException – if an I/O error occurs./**
     * Does the decryption using inBuffer as input and outBuffer as output. Upon
     * return, inBuffer is cleared; the decrypted data starts at
     * outBuffer.position() and ends at outBuffer.limit().
     *
     * @param state the CipherState instance.
     * @param inByteBuffer the input buffer.
     * @param outByteBuffer the output buffer.
     * @param padding the padding.
     * @throws IOException if an I/O error occurs.
     */
    private void decrypt(final CipherState state, final ByteBuffer inByteBuffer,
            final ByteBuffer outByteBuffer, final byte padding) throws IOException {
        Utils.checkState(inByteBuffer.position() >= padding);
        if (inByteBuffer.position() == padding) {
            // There is no real data in inBuffer.
            return;
        }
        inByteBuffer.flip();
        outByteBuffer.clear();
        decryptBuffer(state, inByteBuffer, outByteBuffer);
        inByteBuffer.clear();
        outByteBuffer.flip();
        if (padding > 0) {
            /*
             * The plain text and cipher text have a 1:1 mapping, they start at
             * the same position.
             */
            outByteBuffer.position(padding);
        }
    }

    
Does the decryption using inBuffer as input and outBuffer as output.
Params:
state – the CipherState instance.
inByteBuffer – the input buffer.
outByteBuffer – the output buffer.
Throws:
IOException – if an I/O error occurs./**
     * Does the decryption using inBuffer as input and outBuffer as output.
     *
     * @param state the CipherState instance.
     * @param inByteBuffer the input buffer.
     * @param outByteBuffer the output buffer.
     * @throws IOException if an I/O error occurs.
     */
    private void decryptBuffer(final CipherState state, final ByteBuffer inByteBuffer,
            final ByteBuffer outByteBuffer) throws IOException {
        final int inputSize = inByteBuffer.remaining();
        try {
            final int n = state.getCryptoCipher().update(inByteBuffer, outByteBuffer);
            if (n < inputSize) {
                /**
                 * Typically code will not get here. CryptoCipher#update will
                 * consume all input data and put result in outBuffer.
                 * CryptoCipher#doFinal will reset the cipher context.
                 */
                state.getCryptoCipher().doFinal(inByteBuffer, outByteBuffer);
                state.reset(true);
            }
        } catch (final GeneralSecurityException e) {
            throw new IOException(e);
        }
    }

    
This method is executed immediately after decryption. Check whether
cipher should be updated and recalculate padding if needed.
Params:
state – the CipherState instance.
inByteBuffer – the input buffer.
position – the offset from the start of the stream.
iv – the iv.
Throws:
IOException – if an I/O error occurs.
Returns:
the padding./**
     * This method is executed immediately after decryption. Check whether
     * cipher should be updated and recalculate padding if needed.
     *
     * @param state the CipherState instance.
     * @param inByteBuffer the input buffer.
     * @param position the offset from the start of the stream.
     * @param iv the iv.
     * @return the padding.
     * @throws IOException if an I/O error occurs.
     */
    private byte postDecryption(final CipherState state, final ByteBuffer inByteBuffer,
            final long position, final byte[] iv) throws IOException {
        byte padding = 0;
        if (state.isReset()) {
            /*
             * This code is generally not executed since the cipher usually
             * maintains cipher context (e.g. the counter) internally. However,
             * some implementations can't maintain context so a re-init is
             * necessary after each decryption call.
             */
            resetCipher(state, position, iv);
            padding = getPadding(position);
            inByteBuffer.position(padding);
        }
        return padding;
    }

    
Calculates the counter and iv, reset the cipher.
Params:
state – the CipherState instance.
position – the offset from the start of the stream.
iv – the iv.
Throws:
IOException – if an I/O error occurs./**
     * Calculates the counter and iv, reset the cipher.
     *
     * @param state the CipherState instance.
     * @param position the offset from the start of the stream.
     * @param iv the iv.
     * @throws IOException if an I/O error occurs.
     */
    private void resetCipher(final CipherState state, final long position, final byte[] iv)
            throws IOException {
        final long counter = getCounter(position);
        CtrCryptoInputStream.calculateIV(getInitIV(), counter, iv);
        try {
            state.getCryptoCipher().init(Cipher.DECRYPT_MODE, key,
                    new IvParameterSpec(iv));
        } catch (final GeneralSecurityException e) {
        }
        state.reset(false);
    }

    
Gets CryptoCipher from pool.
Throws:
IOException – if an I/O error occurs.
Returns:
the CipherState instance./**
     * Gets CryptoCipher from pool.
     *
     * @return the CipherState instance.
     * @throws IOException if an I/O error occurs.
     */
    private CipherState getCipherState() throws IOException {
        CipherState state = cipherPool.poll();
        if (state == null) {
            CryptoCipher cryptoCipher;
            try {
                cryptoCipher = CryptoCipherFactory.getCryptoCipher("AES/CTR/NoPadding", properties);
            } catch (final GeneralSecurityException e) {
                throw new IOException(e);
            }
            state = new CipherState(cryptoCipher);
        }

        return state;
    }

    
Returns CryptoCipher to pool.
Params:
state – the CipherState instance./**
     * Returns CryptoCipher to pool.
     *
     * @param state the CipherState instance.
     */
    private void returnCipherState(final CipherState state) {
        if (state != null) {
            cipherPool.add(state);
        }
    }

    
Gets direct buffer from pool.
Returns:
the buffer./**
     * Gets direct buffer from pool.
     *
     * @return the buffer.
     */
    private ByteBuffer getBuffer() {
        ByteBuffer buffer = bufferPool.poll();
        if (buffer == null) {
            buffer = ByteBuffer.allocateDirect(getBufferSize());
        }

        return buffer;
    }

    
Returns direct buffer to pool.
Params:
buf – the buffer./**
     * Returns direct buffer to pool.
     *
     * @param buf the buffer.
     */
    private void returnBuffer(final ByteBuffer buf) {
        if (buf != null) {
            buf.clear();
            bufferPool.add(buf);
        }
    }

    
Overrides the CryptoInputStream.close(). Closes this input stream and releases any system resources associated with the stream. 
Throws:
IOException – if an I/O error occurs./**
     * Overrides the {@link CryptoInputStream#close()}. Closes this input stream
     * and releases any system resources associated with the stream.
     *
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public void close() throws IOException {
        if (!isOpen()) {
            return;
        }

        cleanBufferPool();
        super.close();
    }

    
Clean direct buffer pool /** Clean direct buffer pool */
    private void cleanBufferPool() {
        ByteBuffer buf;
        while ((buf = bufferPool.poll()) != null) {
            CryptoInputStream.freeDirectBuffer(buf);
        }
    }

    private static class CipherState {

        private final CryptoCipher cryptoCipher;
        private boolean reset;

        
The constructor of CipherState. 
Params:
cipher – the CryptoCipher instance./**
         * The constructor of {@link CipherState}.
         *
         * @param cipher the CryptoCipher instance.
         */
        public CipherState(final CryptoCipher cipher) {
            this.cryptoCipher = cipher;
            this.reset = false;
        }

        
Gets the CryptoCipher instance.
Returns:
the cipher./**
         * Gets the CryptoCipher instance.
         *
         * @return the cipher.
         */
        public CryptoCipher getCryptoCipher() {
            return cryptoCipher;
        }

        
Gets the reset.
Returns:
the value of reset./**
         * Gets the reset.
         *
         * @return the value of reset.
         */
        public boolean isReset() {
            return reset;
        }

        
Sets the value of reset.
Params:
reset – the reset./**
         * Sets the value of reset.
         *
         * @param reset the reset.
         */
        public void reset(final boolean reset) {
            this.reset = reset;
        }
    }
}