From 927e28cb20854827a769d4fa8e3f3d9b6814e607 Mon Sep 17 00:00:00 2001
From: Trustin Lee <t@motd.kr>
Date: Sun, 18 May 2014 05:09:22 +0900
Subject: [PATCH] Reduce memory usage of SslHandler when OpenSslEngine is in
 use

Motivation:

JDK's SSLEngine.wrap() requires the output buffer to be always as large as MAX_ENCRYPTED_PACKET_LENGTH even if the input buffer contains small number of bytes.  Our OpenSslEngine implementation does not have such wasteful behaviot.

Modifications:

If the current SSLEngine is OpenSslEngine, allocate as much as only needed.

Result:

Less peak memory usage.
---
 .../io/netty/handler/ssl/OpenSslEngine.java   |  2 ++
 .../java/io/netty/handler/ssl/SslHandler.java | 33 ++++++++++++++++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/handler/src/main/java/io/netty/handler/ssl/OpenSslEngine.java b/handler/src/main/java/io/netty/handler/ssl/OpenSslEngine.java
index c0b1e400e2..0fc8c4a904 100644
--- a/handler/src/main/java/io/netty/handler/ssl/OpenSslEngine.java
+++ b/handler/src/main/java/io/netty/handler/ssl/OpenSslEngine.java
@@ -66,6 +66,8 @@ public final class OpenSslEngine extends SSLEngine {
     // Header (5) + Data (2^14) + Compression (1024) + Encryption (1024) + MAC (20) + Padding (256)
     static final int MAX_ENCRYPTED_PACKET_LENGTH = MAX_CIPHERTEXT_LENGTH + 5 + 20 + 256;
 
+    static final int MAX_ENCRYPTION_OVERHEAD_LENGTH = MAX_ENCRYPTED_PACKET_LENGTH - MAX_PLAINTEXT_LENGTH;
+
     private static final AtomicIntegerFieldUpdater<OpenSslEngine> DESTROYED_UPDATER =
             AtomicIntegerFieldUpdater.newUpdater(OpenSslEngine.class, "destroyed");
 
diff --git a/handler/src/main/java/io/netty/handler/ssl/SslHandler.java b/handler/src/main/java/io/netty/handler/ssl/SslHandler.java
index dfe1238086..27696df399 100644
--- a/handler/src/main/java/io/netty/handler/ssl/SslHandler.java
+++ b/handler/src/main/java/io/netty/handler/ssl/SslHandler.java
@@ -180,6 +180,14 @@ public class SslHandler extends ByteToMessageDecoder implements ChannelOutboundH
     private final SSLEngine engine;
     private final int maxPacketBufferSize;
     private final Executor delegatedTaskExecutor;
+    /**
+     * {@code true} if and only if {@link SSLEngine#wrap(ByteBuffer, ByteBuffer)} requires the output buffer
+     * to be always as large as {@link #maxPacketBufferSize} even if the input buffer contains small amount of data.
+     * <p>
+     * If this flag is {@code false}, we allocate a smaller output buffer.
+     * </p>
+     */
+    private final boolean needsLargeOutNetBuf;
 
     private final boolean startTls;
     private boolean sentFirstMessage;
@@ -243,6 +251,7 @@ public class SslHandler extends ByteToMessageDecoder implements ChannelOutboundH
         this.delegatedTaskExecutor = delegatedTaskExecutor;
         this.startTls = startTls;
         maxPacketBufferSize = engine.getSession().getPacketBufferSize();
+        needsLargeOutNetBuf = !(engine instanceof OpenSslEngine);
     }
 
     public long getHandshakeTimeoutMillis() {
@@ -426,16 +435,18 @@ public class SslHandler extends ByteToMessageDecoder implements ChannelOutboundH
                 if (pending == null) {
                     break;
                 }
-                if (out == null) {
-                    out = allocate(ctx, maxPacketBufferSize);
-                }
 
                 if (!(pending.msg() instanceof ByteBuf)) {
                     ctx.write(pending.msg(), (ChannelPromise) pending.recycleAndGet());
                     pendingUnencryptedWrites.remove();
                     continue;
                 }
+
                 ByteBuf buf = (ByteBuf) pending.msg();
+                if (out == null) {
+                    out = allocateOutNetBuf(ctx, buf.readableBytes());
+                }
+
                 SSLEngineResult result = wrap(engine, buf, out);
 
                 if (!buf.isReadable()) {
@@ -513,7 +524,7 @@ public class SslHandler extends ByteToMessageDecoder implements ChannelOutboundH
         try {
             for (;;) {
                 if (out == null) {
-                    out = allocate(ctx, maxPacketBufferSize);
+                    out = allocateOutNetBuf(ctx, 0);
                 }
                 SSLEngineResult result = wrap(engine, Unpooled.EMPTY_BUFFER, out);
 
@@ -1228,6 +1239,20 @@ public class SslHandler extends ByteToMessageDecoder implements ChannelOutboundH
         }
     }
 
+    /**
+     * Allocates an outbound network buffer for {@link SSLEngine#wrap(ByteBuffer, ByteBuffer)} which can encrypt
+     * the specified amount of pending bytes.
+     */
+    private ByteBuf allocateOutNetBuf(ChannelHandlerContext ctx, int pendingBytes) {
+        if (needsLargeOutNetBuf) {
+            return allocate(ctx, maxPacketBufferSize);
+        } else {
+            return allocate(ctx, Math.min(
+                    pendingBytes + OpenSslEngine.MAX_ENCRYPTION_OVERHEAD_LENGTH,
+                    maxPacketBufferSize));
+        }
+    }
+
     private final class LazyChannelPromise extends DefaultPromise<Channel> {
 
         @Override