package net.md_5.bungee.netty;

import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.ByteToByteCodec;
import javax.crypto.Cipher;

/**
 * This class is a complete solution for encrypting and decoding bytes in a
 * Netty stream. It takes two {@link Cipher} instances, used for encryption and
 * decryption respectively.
 */
public class CipherCodec extends ByteToByteCodec
{

    private Cipher encrypt;
    private Cipher decrypt;
    private ByteBuf heapIn;
    private ByteBuf heapOut;

    public CipherCodec(Cipher encrypt, Cipher decrypt)
    {
        this.encrypt = encrypt;
        this.decrypt = decrypt;
    }

    @Override
    public void encode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception
    {
        cipher( ctx, in, out, encrypt );
    }

    @Override
    public void decode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception
    {
        cipher( ctx, in, out, decrypt );
    }

    @Override
    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();
            heapIn.capacity( heapIn.writerIndex() + readableBytes );
            // Read into buffer
            in.readBytes( heapIn );

            // Allocate output buffer
            if ( heapOut == null )
            {
                heapOut = ctx.alloc().heapBuffer();
            }

            // Get output size
            int outputSize = cipher.getOutputSize( readableBytes );
            // Check we have enough space
            if ( heapOut.writableBytes() < outputSize )
            {
                heapOut.capacity( heapOut.writerIndex() + 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();
        }
    }
}

@@ -178,16 +178,58 @@
 .Cipher;
+%0Aimport javax.crypto.ShortBufferException;
 %0A%0A/**%0A *
@@ -538,22 +538,37 @@
 ate 
-ByteBuf
+byte%5B%5D
  heapIn
+ = new byte%5B 0 %5D
 ;%0A  
@@ -581,23 +581,22 @@
 ate 
-ByteBuf
+byte%5B%5D
  heapOut
 ;%0A%0A 
@@ -591,16 +591,32 @@
  heapOut
+ = new byte%5B 0 %5D
 ;%0A%0A    p
@@ -861,37 +861,32 @@
 %0A        cipher(
- ctx,
  in, out, encryp
@@ -1024,21 +1024,16 @@
  cipher(
- ctx,
  in, out
@@ -1060,921 +1060,102 @@
     
-@Override%0A    public void freeInboundBuffer(ChannelHandlerContext ctx) throws Exception%0A    %7B%0A        super.freeInboundBuffer( ctx );%0A        free();%0A    %7D%0A%0A    @Override%0A    public void freeOutboundBuffer(ChannelHandlerContext ctx) throws Exception%0A    %7B%0A        super.freeOutboundBuffer( ctx );%0A        free();%0A    %7D%0A%0A    private void free()%0A    %7B%0A        if ( heapIn != null )%0A        %7B%0A            heapIn.release();%0A            heapIn = null;%0A        %7D%0A        if ( heapOut != null )%0A        %7B%0A            heapOut.release();%0A            heapOut = null;%0A        %7D%0A    %7D%0A%0A    private void cipher(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out, Cipher cipher) throws Exception%0A    %7B%0A        try%0A        %7B%0A            // Allocate input buffer%0A            if ( heapIn == null )%0A            %7B%0A                heapIn = ctx.alloc().heapBuffer();%0A            %7D%0A%0A            // Allocate correct amount of space
+private void cipher(ByteBuf in, ByteBuf out, Cipher cipher) throws ShortBufferException
 %0A    
+%7B%0A
     
@@ -1210,51 +1210,28 @@
     
-    heapIn.capacity( heapIn.writerIndex() +
+if ( heapIn.length %3C
  rea
@@ -1234,33 +1234,32 @@
  readableBytes )
-;
 %0A            // 
@@ -1255,404 +1255,210 @@
     
+%7B%0A
     
-// Read into buffer%0A            in.readBytes( heapIn );%0A%0A            // Allocate output buffer%0A            if ( heapOut == null )%0A            %7B%0A                heapOut = ctx.alloc().heapBuffer();%0A            %7D%0A%0A            // Get output size%0A            int outputSize = cipher.getOutputSize( readableBytes );%0A            // Check we have enough space%0A            if ( heapOut.writableBytes()
+        heapIn = new byte%5B readableBytes %5D;%0A        %7D%0A        in.readBytes( heapIn, 0, readableBytes );%0A%0A        int outputSize = cipher.getOutputSize( readableBytes );%0A        if ( heapOut.length
  %3C o
@@ -1473,32 +1473,26 @@
 e )%0A        
-    %7B%0A  
+%7B%0A
             
@@ -1495,51 +1495,27 @@
     
-  
 heapOut
-.capacity( heapOut.writerIndex() +
+ = new byte%5B
  out
@@ -1522,17 +1522,17 @@
 putSize 
-)
+%5D
 ;%0A      
@@ -1537,174 +1537,41 @@
     
-    %7D%0A%0A            // Do the processing%0A            int processed = cipher.update( heapIn.array(), heapIn.arrayOffset() + heapIn.readerIndex(), readabl
+%7D%0A        out.writ
 eBytes
-,
+(
  heapOut
 .arr
@@ -1570,222 +1570,53 @@
 pOut
-.array(), heapOut.arrayOffset() + heapOut.writerIndex() );%0A%0A            // Tell the out buffer we read some from it%0A            heapOut.writerIndex( heapOut.writerIndex() + processed );%0A%0A            out.writ
+, 0, cipher.update( heapIn, 0, readabl
 eBytes
-(
+,
  hea
@@ -1625,205 +1625,11 @@
 ut )
-;%0A            heapIn.clear();%0A            heapOut.clear();%0A        %7D catch ( Throwable ex )%0A        %7B%0A            // TODO: Remove this once we are stable%0A            ex.printStackTrace();%0A        %7D
+ );
 %0A   

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