cz.d1x.dxcrypto.encryption.io.DecryptingOutputStream Maven / Gradle / Ivy

Go to download
package cz.d1x.dxcrypto.encryption.io;

import cz.d1x.dxcrypto.encryption.StreamingEncryptionEngine;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * This class is wrapper around given stream and handles first bytes written as IV to initialize cipher.
 *
 * @author d.richter
 */
public class DecryptingOutputStream extends FilterOutputStream {

    private final StreamingEncryptionEngine engine;
    private final byte[] ivBuffer;
    private final AtomicInteger ivWritePosition = new AtomicInteger();
    private OutputStream decryptingStream;

    public DecryptingOutputStream(OutputStream out, StreamingEncryptionEngine engine, int blockSize) {
        super(out);
        this.engine = engine;
        ivBuffer = new byte[blockSize];
    }

    @Override
    public void write(int b) throws IOException {
        final int currentWritePosition = ivWritePosition.get();
        if (currentWritePosition < ivBuffer.length) {
            ivBuffer[ivWritePosition.getAndIncrement()] = (byte) b;
            if (ivWritePosition.get() == ivBuffer.length) {
                decryptingStream = engine.decrypt(out, ivBuffer);
            }
        } else {
            decryptingStream.write(b);
        }
    }

    @Override
    public void close() throws IOException {
        if (decryptingStream != null) {
            decryptingStream.close();
        }
        super.close();
    }
}