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.io.OutputStream;
022import java.nio.ByteBuffer;
023import java.nio.channels.WritableByteChannel;
024import java.security.GeneralSecurityException;
025import java.util.Properties;
026
027import javax.crypto.Cipher;
028import javax.crypto.spec.IvParameterSpec;
029import javax.crypto.spec.SecretKeySpec;
030
031import org.apache.commons.crypto.cipher.CryptoCipher;
032import org.apache.commons.crypto.stream.output.ChannelOutput;
033import org.apache.commons.crypto.stream.output.Output;
034import org.apache.commons.crypto.stream.output.StreamOutput;
035import org.apache.commons.crypto.utils.Utils;
036
037/**
038 * <p>
039 * CtrCryptoOutputStream encrypts data. It is not thread-safe. AES CTR mode is
040 * required in order to ensure that the plain text and cipher text have a 1:1
041 * mapping. The encryption is buffer based. The key points of the encryption are
042 * (1) calculating counter and (2) padding through stream position.
043 * </p>
044 * <p>
045 * counter = base + pos/(algorithm blocksize); padding = pos%(algorithm
046 * blocksize);
047 * </p>
048 * <p>
049 * The underlying stream offset is maintained as state.
050 * </p>
051 * <p>
052 * This class should only be used with blocking sinks. Using this class to wrap
053 * a non-blocking sink may lead to high CPU usage.
054 * </p>
055 */
056public class CtrCryptoOutputStream extends CryptoOutputStream {
057    /**
058     * Underlying stream offset.
059     */
060    private long streamOffset = 0;
061
062    /**
063     * The initial IV.
064     */
065    private final byte[] initIV;
066
067    /**
068     * Initialization vector for the cipher.
069     */
070    private final byte[] iv;
071
072    /**
073     * Padding = pos%(algorithm blocksize); Padding is put into
074     * {@link #inBuffer} before any other data goes in. The purpose of padding
075     * is to put input data at proper position.
076     */
077    private byte padding;
078
079    /**
080     * Flag to mark whether the cipher has been reset
081     */
082    private boolean cipherReset = false;
083
084    /**
085     * Constructs a {@link CtrCryptoOutputStream}.
086     *
087     * @param props The {@code Properties} class represents a set of
088     *        properties.
089     * @param out the output stream.
090     * @param key crypto key for the cipher.
091     * @param iv Initialization vector for the cipher.
092     * @throws IOException if an I/O error occurs.
093     */
094    public CtrCryptoOutputStream(final Properties props, final OutputStream out,
095            final byte[] key, final byte[] iv) throws IOException {
096        this(props, out, key, iv, 0);
097    }
098
099    /**
100     * Constructs a {@link CtrCryptoOutputStream}.
101     *
102     * @param props The {@code Properties} class represents a set of
103     *        properties.
104     * @param out the WritableByteChannel instance.
105     * @param key crypto key for the cipher.
106     * @param iv Initialization vector for the cipher.
107     * @throws IOException if an I/O error occurs.
108     */
109    public CtrCryptoOutputStream(final Properties props, final WritableByteChannel out,
110            final byte[] key, final byte[] iv) throws IOException {
111        this(props, out, key, iv, 0);
112    }
113
114    /**
115     * Constructs a {@link CtrCryptoOutputStream}.
116     *
117     * @param out the output stream.
118     * @param cipher the CryptoCipher instance.
119     * @param bufferSize the bufferSize.
120     * @param key crypto key for the cipher.
121     * @param iv Initialization vector for the cipher.
122     * @throws IOException if an I/O error occurs.
123     */
124    protected CtrCryptoOutputStream(final OutputStream out, final CryptoCipher cipher,
125            final int bufferSize, final byte[] key, final byte[] iv) throws IOException {
126        this(out, cipher, bufferSize, key, iv, 0);
127    }
128
129    /**
130     * Constructs a {@link CtrCryptoOutputStream}.
131     *
132     * @param channel the WritableByteChannel instance.
133     * @param cipher the CryptoCipher instance.
134     * @param bufferSize the bufferSize.
135     * @param key crypto key for the cipher.
136     * @param iv Initialization vector for the cipher.
137     * @throws IOException if an I/O error occurs.
138     */
139    protected CtrCryptoOutputStream(final WritableByteChannel channel,
140            final CryptoCipher cipher, final int bufferSize, final byte[] key, final byte[] iv)
141            throws IOException {
142        this(channel, cipher, bufferSize, key, iv, 0);
143    }
144
145    /**
146     * Constructs a {@link CtrCryptoOutputStream}.
147     *
148     * @param output the Output instance.
149     * @param cipher the CryptoCipher instance.
150     * @param bufferSize the bufferSize.
151     * @param key crypto key for the cipher.
152     * @param iv Initialization vector for the cipher.
153     * @throws IOException if an I/O error occurs.
154     */
155    protected CtrCryptoOutputStream(final Output output, final CryptoCipher cipher,
156            final int bufferSize, final byte[] key, final byte[] iv) throws IOException {
157        this(output, cipher, bufferSize, key, iv, 0);
158    }
159
160    /**
161     * Constructs a {@link CtrCryptoOutputStream}.
162     *
163     * @param properties The {@code Properties} class represents a set of
164     *        properties.
165     * @param outputStream the output stream.
166     * @param key crypto key for the cipher.
167     * @param iv Initialization vector for the cipher.
168     * @param streamOffset the start offset in the data.
169     * @throws IOException if an I/O error occurs.
170     */
171    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CtrCryptoOutputStream.
172    public CtrCryptoOutputStream(final Properties properties, final OutputStream outputStream,
173            final byte[] key, final byte[] iv, final long streamOffset) throws IOException {
174        this(outputStream, Utils.getCipherInstance(
175                "AES/CTR/NoPadding", properties),
176                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
177    }
178
179    /**
180     * Constructs a {@link CtrCryptoOutputStream}.
181     *
182     * @param properties The {@code Properties} class represents a set of
183     *        properties.
184     * @param channel the WritableByteChannel instance.
185     * @param key crypto key for the cipher.
186     * @param iv Initialization vector for the cipher.
187     * @param streamOffset the start offset in the data.
188     * @throws IOException if an I/O error occurs.
189     */
190    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CtrCryptoOutputStream.
191    public CtrCryptoOutputStream(final Properties properties, final WritableByteChannel channel,
192            final byte[] key, final byte[] iv, final long streamOffset) throws IOException {
193        this(channel, Utils.getCipherInstance(
194                "AES/CTR/NoPadding", properties),
195                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
196    }
197
198    /**
199     * Constructs a {@link CtrCryptoOutputStream}.
200     *
201     * @param outputStream the output stream.
202     * @param cipher the CryptoCipher instance.
203     * @param bufferSize the bufferSize.
204     * @param key crypto key for the cipher.
205     * @param iv Initialization vector for the cipher.
206     * @param streamOffset the start offset in the data.
207     * @throws IOException if an I/O error occurs.
208     */
209    protected CtrCryptoOutputStream(final OutputStream outputStream, final CryptoCipher cipher,
210            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
211            throws IOException {
212        this(new StreamOutput(outputStream, bufferSize), cipher, bufferSize, key, iv,
213                streamOffset);
214    }
215
216    /**
217     * Constructs a {@link CtrCryptoOutputStream}.
218     *
219     * @param channel the WritableByteChannel instance.
220     * @param cipher the CryptoCipher instance.
221     * @param bufferSize the bufferSize.
222     * @param key crypto key for the cipher.
223     * @param iv Initialization vector for the cipher.
224     * @param streamOffset the start offset in the data.
225     * @throws IOException if an I/O error occurs.
226     */
227    protected CtrCryptoOutputStream(final WritableByteChannel channel,
228            final CryptoCipher cipher, final int bufferSize, final byte[] key, final byte[] iv,
229            final long streamOffset) throws IOException {
230        this(new ChannelOutput(channel), cipher, bufferSize, key, iv,
231                streamOffset);
232    }
233
234    /**
235     * Constructs a {@link CtrCryptoOutputStream}.
236     *
237     * @param output the output stream.
238     * @param cipher the CryptoCipher instance.
239     * @param bufferSize the bufferSize.
240     * @param key crypto key for the cipher.
241     * @param iv Initialization vector for the cipher.
242     * @param streamOffset the start offset in the data.
243     * @throws IOException if an I/O error occurs.
244     */
245    protected CtrCryptoOutputStream(final Output output, final CryptoCipher cipher,
246            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
247            throws IOException {
248        super(output, cipher, bufferSize, new SecretKeySpec(key, "AES"),
249                new IvParameterSpec(iv));
250
251        CryptoInputStream.checkStreamCipher(cipher);
252        this.streamOffset = streamOffset;
253        this.initIV = iv.clone();
254        this.iv = iv.clone();
255
256        resetCipher();
257    }
258
259    /**
260     * Does the encryption, input is {@link #inBuffer} and output is
261     * {@link #outBuffer}.
262     *
263     * @throws IOException if an I/O error occurs.
264     */
265    @Override
266    protected void encrypt() throws IOException {
267        Utils.checkState(inBuffer.position() >= padding);
268        if (inBuffer.position() == padding) {
269            // There is no real data in the inBuffer.
270            return;
271        }
272
273        inBuffer.flip();
274        outBuffer.clear();
275        encryptBuffer(outBuffer);
276        inBuffer.clear();
277        outBuffer.flip();
278
279        if (padding > 0) {
280            /*
281             * The plain text and cipher text have a 1:1 mapping, they start at
282             * the same position.
283             */
284            outBuffer.position(padding);
285            padding = 0;
286        }
287
288        final int len = output.write(outBuffer);
289        streamOffset += len;
290        if (cipherReset) {
291            /*
292             * This code is generally not executed since the encryptor usually
293             * maintains encryption context (e.g. the counter) internally.
294             * However, some implementations can't maintain context so a re-init
295             * is necessary after each encryption call.
296             */
297            resetCipher();
298        }
299    }
300
301    /**
302     * Does final encryption of the last data.
303     *
304     * @throws IOException if an I/O error occurs.
305     */
306    @Override
307    protected void encryptFinal() throws IOException {
308        // The same as the normal encryption for Counter mode
309        encrypt();
310    }
311
312    /**
313     * Overrides the {@link CryptoOutputStream#initCipher()}. Initializes the
314     * cipher.
315     */
316    @Override
317    protected void initCipher() {
318        // Do nothing for initCipher
319        // Will reset the cipher considering the stream offset
320    }
321
322    /**
323     * Resets the {@link #cipher}: calculate counter and {@link #padding}.
324     *
325     * @throws IOException if an I/O error occurs.
326     */
327    private void resetCipher() throws IOException {
328        final long counter = streamOffset
329                / cipher.getBlockSize();
330        padding = (byte) (streamOffset % cipher.getBlockSize());
331        inBuffer.position(padding); // Set proper position for input data.
332
333        CtrCryptoInputStream.calculateIV(initIV, counter, iv);
334        try {
335            cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
336        } catch (final GeneralSecurityException e) {
337            throw new IOException(e);
338        }
339        cipherReset = false;
340    }
341
342    /**
343     * Does the encryption if the ByteBuffer data.
344     *
345     * @param out the output ByteBuffer.
346     * @throws IOException if an I/O error occurs.
347     */
348    private void encryptBuffer(final ByteBuffer out) throws IOException {
349        final int inputSize = inBuffer.remaining();
350        try {
351            final int n = cipher.update(inBuffer, out);
352            if (n < inputSize) {
353                /**
354                 * Typically code will not get here. CryptoCipher#update will
355                 * consume all input data and put result in outBuffer.
356                 * CryptoCipher#doFinal will reset the cipher context.
357                 */
358                cipher.doFinal(inBuffer, out);
359                cipherReset = true;
360            }
361        } catch (final GeneralSecurityException e) {
362            throw new IOException(e);
363        }
364    }
365
366    /**
367     * Get the underlying stream offset
368     *
369     * @return the underlying stream offset
370     */
371    protected long getStreamOffset() {
372        return streamOffset;
373    }
374
375    /**
376     * Set the underlying stream offset
377     *
378     * @param streamOffset the underlying stream offset
379     */
380    protected void setStreamOffset(final long streamOffset) {
381        this.streamOffset = streamOffset;
382    }
383}