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 */
018
019package org.apache.commons.crypto.stream;
020
021import java.io.IOException;
022import java.io.OutputStream;
023import java.nio.ByteBuffer;
024import java.nio.channels.WritableByteChannel;
025import java.security.GeneralSecurityException;
026import java.security.Key;
027import java.security.spec.AlgorithmParameterSpec;
028import java.util.Objects;
029import java.util.Properties;
030
031import javax.crypto.Cipher;
032import javax.crypto.ShortBufferException;
033import javax.crypto.spec.IvParameterSpec;
034
035import org.apache.commons.crypto.cipher.CryptoCipher;
036import org.apache.commons.crypto.stream.output.ChannelOutput;
037import org.apache.commons.crypto.stream.output.Output;
038import org.apache.commons.crypto.stream.output.StreamOutput;
039import org.apache.commons.crypto.utils.Utils;
040
041/**
042 * {@link CryptoOutputStream} encrypts data and writes to the under layer
043 * output. It supports any mode of operations such as AES CBC/CTR/GCM mode in
044 * concept. It is not thread-safe.
045 * <p>
046 * This class should only be used with blocking sinks. Using this class to wrap
047 * a non-blocking sink may lead to high CPU usage.
048 * </p>
049 */
050
051public class CryptoOutputStream extends OutputStream implements
052        WritableByteChannel {
053    private final byte[] oneByteBuf = new byte[1];
054
055    /** The output. */
056    final Output output; // package protected for access by rypto classes; do not expose further
057
058    /** the CryptoCipher instance */
059    final CryptoCipher cipher; // package protected for access by crypto classes; do not expose further
060
061    /** The buffer size. */
062    private final int bufferSize;
063
064    /** Crypto key for the cipher. */
065    final Key key; // package protected for access by crypto classes; do not expose further
066
067    /** the algorithm parameters */
068    private final AlgorithmParameterSpec params;
069
070    /** Flag to mark whether the output stream is closed. */
071    private boolean closed;
072
073    /**
074     * Input data buffer. The data starts at inBuffer.position() and ends at
075     * inBuffer.limit().
076     */
077    ByteBuffer inBuffer; // package protected for access by crypto classes; do not expose further
078
079    /**
080     * Encrypted data buffer. The data starts at outBuffer.position() and ends
081     * at outBuffer.limit().
082     */
083    ByteBuffer outBuffer; // package protected for access by crypto classes; do not expose further
084
085    /**
086     * Constructs a {@link CryptoOutputStream}.
087     *
088     * @param transformation the name of the transformation, e.g.,
089     * <i>AES/CBC/PKCS5Padding</i>.
090     * See the Java Cryptography Architecture Standard Algorithm Name Documentation
091     * for information about standard transformation names.
092     * @param properties The {@code Properties} class represents a set of
093     *        properties.
094     * @param outputStream the output stream.
095     * @param key crypto key for the cipher.
096     * @param params the algorithm parameters.
097     * @throws IOException if an I/O error occurs.
098     */
099    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CryptoOutputStream.
100    public CryptoOutputStream(final String transformation,
101            final Properties properties, final OutputStream outputStream, final Key key,
102            final AlgorithmParameterSpec params) throws IOException {
103        this(outputStream, Utils.getCipherInstance(transformation, properties),
104                CryptoInputStream.getBufferSize(properties), key, params);
105
106    }
107
108    /**
109     * Constructs a {@link CryptoOutputStream}.
110     *
111     * @param transformation the name of the transformation, e.g.,
112     * <i>AES/CBC/PKCS5Padding</i>.
113     * See the Java Cryptography Architecture Standard Algorithm Name Documentation
114     * for information about standard transformation names.
115     * @param properties The {@code Properties} class represents a set of
116     *        properties.
117     * @param out the WritableByteChannel instance.
118     * @param key crypto key for the cipher.
119     * @param params the algorithm parameters.
120     * @throws IOException if an I/O error occurs.
121     */
122    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CryptoOutputStream.
123    public CryptoOutputStream(final String transformation,
124            final Properties properties, final WritableByteChannel out, final Key key,
125            final AlgorithmParameterSpec params) throws IOException {
126        this(out, Utils.getCipherInstance(transformation, properties), CryptoInputStream
127                .getBufferSize(properties), key, params);
128
129    }
130
131    /**
132     * Constructs a {@link CryptoOutputStream}.
133     *
134     * @param outputStream the output stream.
135     * @param cipher the CryptoCipher instance.
136     * @param bufferSize the bufferSize.
137     * @param key crypto key for the cipher.
138     * @param params the algorithm parameters.
139     * @throws IOException if an I/O error occurs.
140     */
141    protected CryptoOutputStream(final OutputStream outputStream, final CryptoCipher cipher,
142            final int bufferSize, final Key key, final AlgorithmParameterSpec params)
143            throws IOException {
144        this(new StreamOutput(outputStream, bufferSize), cipher, bufferSize, key, params);
145    }
146
147    /**
148     * Constructs a {@link CryptoOutputStream}.
149     *
150     * @param channel the WritableByteChannel instance.
151     * @param cipher the cipher instance.
152     * @param bufferSize the bufferSize.
153     * @param key crypto key for the cipher.
154     * @param params the algorithm parameters.
155     * @throws IOException if an I/O error occurs.
156     */
157    protected CryptoOutputStream(final WritableByteChannel channel, final CryptoCipher cipher,
158            final int bufferSize, final Key key, final AlgorithmParameterSpec params)
159            throws IOException {
160        this(new ChannelOutput(channel), cipher, bufferSize, key, params);
161    }
162
163    /**
164     * Constructs a {@link CryptoOutputStream}.
165     *
166     * @param output the output stream.
167     * @param cipher the CryptoCipher instance.
168     * @param bufferSize the bufferSize.
169     * @param key crypto key for the cipher.
170     * @param params the algorithm parameters.
171     * @throws IOException if an I/O error occurs.
172     */
173    protected CryptoOutputStream(final Output output, final CryptoCipher cipher,
174            final int bufferSize, final Key key, final AlgorithmParameterSpec params)
175            throws IOException {
176
177        this.output = output;
178        this.bufferSize = CryptoInputStream.checkBufferSize(cipher, bufferSize);
179        this.cipher = cipher;
180
181        this.key = key;
182        this.params = params;
183
184        if (!(params instanceof IvParameterSpec)) {
185            // other AlgorithmParameterSpec such as GCMParameterSpec is not
186            // supported now.
187            throw new IOException("Illegal parameters");
188        }
189
190        inBuffer = ByteBuffer.allocateDirect(this.bufferSize);
191        outBuffer = ByteBuffer.allocateDirect(this.bufferSize
192                + cipher.getBlockSize());
193
194        initCipher();
195    }
196
197    /**
198     * Overrides the {@link java.io.OutputStream#write(byte[])}. Writes the
199     * specified byte to this output stream.
200     *
201     * @param b the data.
202     * @throws IOException if an I/O error occurs.
203     */
204    @Override
205    public void write(final int b) throws IOException {
206        oneByteBuf[0] = (byte) (b & 0xff);
207        write(oneByteBuf, 0, oneByteBuf.length);
208    }
209
210    /**
211     * Overrides the {@link java.io.OutputStream#write(byte[], int, int)}.
212     * Encryption is buffer based. If there is enough room in {@link #inBuffer},
213     * then write to this buffer. If {@link #inBuffer} is full, then do
214     * encryption and write data to the underlying stream.
215     *
216     * @param array the data.
217     * @param off the start offset in the data.
218     * @param len the number of bytes to write.
219     * @throws IOException if an I/O error occurs.
220     */
221    @Override
222    public void write(final byte[] array, int off, int len) throws IOException {
223        checkStream();
224        Objects.requireNonNull(array, "array");
225        if (off < 0 || len < 0 || off > array.length || len > array.length - off) {
226            throw new IndexOutOfBoundsException();
227        }
228
229        while (len > 0) {
230            final int remaining = inBuffer.remaining();
231            if (len < remaining) {
232                inBuffer.put(array, off, len);
233                len = 0;
234            } else {
235                inBuffer.put(array, off, remaining);
236                off += remaining;
237                len -= remaining;
238                encrypt();
239            }
240        }
241    }
242
243    /**
244     * Overrides the {@link OutputStream#flush()}. To flush, we need to encrypt
245     * the data in the buffer and write to the underlying stream, then do the
246     * flush.
247     *
248     * @throws IOException if an I/O error occurs.
249     */
250    @Override
251    public void flush() throws IOException {
252        checkStream();
253        encrypt();
254        output.flush();
255        super.flush();
256    }
257
258    /**
259     * Overrides the {@link OutputStream#close()}. Closes this output stream and
260     * releases any system resources associated with this stream.
261     *
262     * @throws IOException if an I/O error occurs.
263     */
264    @Override
265    public void close() throws IOException {
266        if (closed) {
267            return;
268        }
269
270        try {
271            encryptFinal();
272            output.close();
273            freeBuffers();
274            cipher.close();
275            super.close();
276        } finally {
277            closed = true;
278        }
279    }
280
281    /**
282     * Overrides the {@link java.nio.channels.Channel#isOpen()}. Tells whether or not this channel
283     * is open.
284     *
285     * @return {@code true} if, and only if, this channel is open
286     */
287    @Override
288    public boolean isOpen() {
289        return !closed;
290    }
291
292    /**
293     * Overrides the
294     * {@link java.nio.channels.WritableByteChannel#write(ByteBuffer)}. Writes a
295     * sequence of bytes to this channel from the given buffer.
296     *
297     * @param src The buffer from which bytes are to be retrieved.
298     * @return The number of bytes written, possibly zero.
299     * @throws IOException if an I/O error occurs.
300     */
301    @Override
302    public int write(final ByteBuffer src) throws IOException {
303        checkStream();
304        final int len = src.remaining();
305        int remaining = len;
306        while (remaining > 0) {
307            final int space = inBuffer.remaining();
308            if (remaining < space) {
309                inBuffer.put(src);
310                remaining = 0;
311            } else {
312                // to void copy twice, we set the limit to copy directly
313                final int oldLimit = src.limit();
314                final int newLimit = src.position() + space;
315                src.limit(newLimit);
316
317                inBuffer.put(src);
318
319                // restore the old limit
320                src.limit(oldLimit);
321
322                remaining -= space;
323                encrypt();
324            }
325        }
326
327        return len;
328    }
329
330    /**
331     * Initializes the cipher.
332     *
333     * @throws IOException if an I/O error occurs.
334     */
335    protected void initCipher() throws IOException {
336        try {
337            cipher.init(Cipher.ENCRYPT_MODE, key, params);
338        } catch (final GeneralSecurityException e) {
339            throw new IOException(e);
340        }
341    }
342
343    /**
344     * Does the encryption, input is {@link #inBuffer} and output is
345     * {@link #outBuffer}.
346     *
347     * @throws IOException if an I/O error occurs.
348     */
349    protected void encrypt() throws IOException {
350
351        inBuffer.flip();
352        outBuffer.clear();
353
354        try {
355            cipher.update(inBuffer, outBuffer);
356        } catch (final ShortBufferException e) {
357            throw new IOException(e);
358        }
359
360        inBuffer.clear();
361        outBuffer.flip();
362
363        // write to output
364        while (outBuffer.hasRemaining()) {
365            output.write(outBuffer);
366        }
367    }
368
369    /**
370     * Does final encryption of the last data.
371     *
372     * @throws IOException if an I/O error occurs.
373     */
374    protected void encryptFinal() throws IOException {
375        inBuffer.flip();
376        outBuffer.clear();
377
378        try {
379            cipher.doFinal(inBuffer, outBuffer);
380        } catch (final GeneralSecurityException e) {
381            throw new IOException(e);
382        }
383
384        inBuffer.clear();
385        outBuffer.flip();
386
387        // write to output
388        while (outBuffer.hasRemaining()) {
389            output.write(outBuffer);
390        }
391    }
392
393    /**
394     * Checks whether the stream is closed.
395     *
396     * @throws IOException if an I/O error occurs.
397     */
398    protected void checkStream() throws IOException {
399        if (closed) {
400            throw new IOException("Stream closed");
401        }
402    }
403
404    /** Forcibly free the direct buffers. */
405    protected void freeBuffers() {
406        CryptoInputStream.freeDirectBuffer(inBuffer);
407        CryptoInputStream.freeDirectBuffer(outBuffer);
408    }
409
410    /**
411     * Gets the outBuffer.
412     *
413     * @return the outBuffer.
414     */
415    protected ByteBuffer getOutBuffer() {
416        return outBuffer;
417    }
418
419    /**
420     * Gets the internal Cipher.
421     *
422     * @return the cipher instance.
423     */
424    protected CryptoCipher getCipher() {
425        return cipher;
426    }
427
428    /**
429     * Gets the buffer size.
430     *
431     * @return the buffer size.
432     */
433    protected int getBufferSize() {
434        return bufferSize;
435    }
436
437    /**
438     * Gets the inBuffer.
439     *
440     * @return the inBuffer.
441     */
442    protected ByteBuffer getInBuffer() {
443        return inBuffer;
444    }
445}