Clean up cipher codec for maximum speed and minimal memory copy.

This commit is contained in:
md_5 2013-03-14 19:39:03 +11:00
parent d54f2462a8
commit b71d253de2

View File

@ -4,6 +4,7 @@ import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.ByteToByteCodec; import io.netty.handler.codec.ByteToByteCodec;
import javax.crypto.Cipher; import javax.crypto.Cipher;
import javax.crypto.ShortBufferException;
/** /**
* This class is a complete solution for encrypting and decoding bytes in a * This class is a complete solution for encrypting and decoding bytes in a
@ -15,8 +16,8 @@ public class CipherCodec extends ByteToByteCodec
private Cipher encrypt; private Cipher encrypt;
private Cipher decrypt; private Cipher decrypt;
private ByteBuf heapIn; private byte[] heapIn = new byte[ 0 ];
private ByteBuf heapOut; private byte[] heapOut = new byte[ 0 ];
public CipherCodec(Cipher encrypt, Cipher decrypt) public CipherCodec(Cipher encrypt, Cipher decrypt)
{ {
@ -27,86 +28,29 @@ public class CipherCodec extends ByteToByteCodec
@Override @Override
public void encode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception public void encode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception
{ {
cipher( ctx, in, out, encrypt ); cipher( in, out, encrypt );
} }
@Override @Override
public void decode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception public void decode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception
{ {
cipher( ctx, in, out, decrypt ); cipher( in, out, decrypt );
} }
@Override private void cipher(ByteBuf in, ByteBuf out, Cipher cipher) throws ShortBufferException
public void freeInboundBuffer(ChannelHandlerContext ctx) throws Exception
{ {
super.freeInboundBuffer( ctx );
free();
}
@Override
public void freeOutboundBuffer(ChannelHandlerContext ctx) throws Exception
{
super.freeOutboundBuffer( ctx );
free();
}
private void free()
{
if ( heapIn != null )
{
heapIn.release();
heapIn = null;
}
if ( heapOut != null )
{
heapOut.release();
heapOut = null;
}
}
private void cipher(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out, Cipher cipher) throws Exception
{
try
{
// Allocate input buffer
if ( heapIn == null )
{
heapIn = ctx.alloc().heapBuffer();
}
// Allocate correct amount of space
int readableBytes = in.readableBytes(); int readableBytes = in.readableBytes();
heapIn.capacity( heapIn.writerIndex() + readableBytes ); if ( heapIn.length < readableBytes )
// Read into buffer
in.readBytes( heapIn );
// Allocate output buffer
if ( heapOut == null )
{ {
heapOut = ctx.alloc().heapBuffer(); heapIn = new byte[ readableBytes ];
} }
in.readBytes( heapIn, 0, readableBytes );
// Get output size
int outputSize = cipher.getOutputSize( readableBytes ); int outputSize = cipher.getOutputSize( readableBytes );
// Check we have enough space if ( heapOut.length < outputSize )
if ( heapOut.writableBytes() < outputSize )
{ {
heapOut.capacity( heapOut.writerIndex() + outputSize ); heapOut = new byte[ outputSize ];
}
// Do the processing
int processed = cipher.update( heapIn.array(), heapIn.arrayOffset() + heapIn.readerIndex(), readableBytes, heapOut.array(), heapOut.arrayOffset() + heapOut.writerIndex() );
// Tell the out buffer we read some from it
heapOut.writerIndex( heapOut.writerIndex() + processed );
out.writeBytes( heapOut );
heapIn.clear();
heapOut.clear();
} catch ( Throwable ex )
{
// TODO: Remove this once we are stable
ex.printStackTrace();
} }
out.writeBytes( heapOut, 0, cipher.update( heapIn, 0, readableBytes, heapOut ) );
} }
} }