[NC-1772] Avoid allocating and returning a ByteBuf from MessageFramer.frame. (#116)

We can write directly to the target instead and avoid the caller needing to release the buffer.
Fixes memory leak when clients are repeatedly connecting and disconnecting.
Signed-off-by: Adrian Sutton <adrian.sutton@consensys.net>

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/netty/AbstractHandshakeHandler.java

@@ -144,7 +144,7 @@ abstract class AbstractHandshakeHandler extends SimpleChannelInboundHandler<Byte
           || outboundMessage.getData().getCode() != WireMessageCodes.HELLO) {
         throw new IllegalStateException("First wire message sent wasn't a HELLO.");
       }
-      out.writeBytes(framer.frame(outboundMessage.getData()));
+      framer.frame(outboundMessage.getData(), out);
       context.pipeline().remove(this);
     }
   }

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/netty/MessageFramer.java

@@ -32,6 +32,6 @@ final class MessageFramer extends MessageToByteEncoder<OutboundMessage> {
   @Override
   protected void encode(
       final ChannelHandlerContext ctx, final OutboundMessage msg, final ByteBuf out) {
-    out.writeBytes(framer.frame(multiplexer.multiplex(msg.getCapability(), msg.getData())));
+    framer.frame(multiplexer.multiplex(msg.getCapability(), msg.getData()), out);
   }
 }

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/rlpx/framing/Framer.java

@@ -288,9 +288,9 @@ public class Framer {
    * MACs.
    *
    * @param message The message to frame.
-   * @return The framed message, as byte buffer.
+   * @param output The {@link ByteBuf} to write framed data to.
    */
-  public synchronized ByteBuf frame(final MessageData message) {
+  public synchronized void frame(final MessageData message, final ByteBuf output) {
     Preconditions.checkArgument(
         message.getSize() < LENGTH_MAX_MESSAGE_FRAME, "Message size in excess of maximum length.");
     // Compress message
@@ -306,24 +306,22 @@ public class Framer {
         // Construct new, compressed message
         final ByteBuf compressedBuf = NetworkMemoryPool.allocate(compressed.length);
         compressedBuf.writeBytes(compressed);
-        return frameAndReleaseMessage(new RawMessage(message.getCode(), compressedBuf));
+        frameAndReleaseMessage(new RawMessage(message.getCode(), compressedBuf), output);
       } finally {
         // We have to release the original message because frameAndRelease only released the
         // compressed copy.
         message.release();
       }
     } else {
-      return frameAndReleaseMessage(message);
+      frameAndReleaseMessage(message, output);
     }
   }
 
   @VisibleForTesting
-  ByteBuf frameAndReleaseMessage(final MessageData message) {
+  void frameAndReleaseMessage(final MessageData message, final ByteBuf buf) {
     try {
       final int frameSize = message.getSize() + LENGTH_MESSAGE_ID;
       final int pad = padding16(frameSize);
-      final int bufSize = LENGTH_FULL_HEADER + (frameSize + pad + LENGTH_MAC);
-      final ByteBuf buf = NetworkMemoryPool.allocate(bufSize);
 
       final byte id = (byte) message.getCode();
 
@@ -343,9 +341,6 @@ public class Framer {
       hMac = Arrays.copyOf(hMac, LENGTH_MAC);
       buf.writeBytes(h).writeBytes(hMac);
 
-      // Sanity check.
-      assert buf.writerIndex() == LENGTH_FULL_HEADER;
-
       // Encrypt payload.
       final byte[] f = new byte[frameSize + pad];
 
@@ -367,11 +362,6 @@ public class Framer {
       fMac = Arrays.copyOf(secrets.updateEgress(xor(fMac, fMacSeed)).getEgressMac(), LENGTH_MAC);
 
       buf.writeBytes(f).writeBytes(fMac);
-
-      // Sanity check: all expected bytes are written.
-      assert buf.writerIndex() == LENGTH_FULL_HEADER + (frameSize + pad + LENGTH_MAC);
-
-      return buf;
     } finally {
       message.release();
     }

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/rlpx/framing/FramerTest.java

@@ -59,7 +59,7 @@ public class FramerTest {
     final Framer framer = new Framer(secrets);
 
     assertThatExceptionOfType(IllegalArgumentException.class)
-        .isThrownBy(() -> framer.frame(ethMessage))
+        .isThrownBy(() -> framer.frame(ethMessage, Unpooled.buffer()))
         .withMessageContaining("Message size in excess of maximum length.");
   }
 
@@ -80,7 +80,8 @@ public class FramerTest {
     final ByteBuf buf = wrappedBuffer(byteArray);
 
     final MessageData ethMessage = new RawMessage(0x00, buf);
-    final ByteBuf framedMessage = framer.frameAndReleaseMessage(ethMessage);
+    final ByteBuf framedMessage = Unpooled.buffer();
+    framer.frameAndReleaseMessage(ethMessage, framedMessage);
 
     final HandshakeSecrets deframeSecrets = secretsFrom(td, true);
     final Framer deframer = new Framer(deframeSecrets);
@@ -186,7 +187,8 @@ public class FramerTest {
     framer = new Framer(secrets);
 
     for (int i = 0; i < decrypted.size(); i++) {
-      final ByteBuf b = framer.frame(decrypted.get(i));
+      final ByteBuf b = Unpooled.buffer();
+      framer.frame(decrypted.get(i), b);
       final byte[] enc = new byte[b.readableBytes()];
       b.readBytes(enc);
       final byte[] expected = decodeHexDump(messages.get(i).get("data").asText());