From 4606890513f72a620aa29fb7b446cd8e2d3ae5fc Mon Sep 17 00:00:00 2001
From: Xiaoyan Lin <linxiaoyan18@gmail.com>
Date: Thu, 28 Jan 2016 13:31:52 -0800
Subject: [PATCH] Retain AbstractBinaryMemcacheDecoder.currentMessage when
 sending it out and release it when it's not used.

Motivation:

AbstractBinaryMemcacheDecoder.currentMessage is not retained after sending it out. Hence, if a message contains `extras`, `io.netty.util.IllegalReferenceCountException` will be thrown in `channelInactive`.

Modifications:

Retain AbstractBinaryMemcacheDecoder.currentMessage After putting it to `out` and release it when it's not used.

Result:

No IllegalReferenceCountException or leak.
---
 .../binary/AbstractBinaryMemcacheDecoder.java | 17 ++++++++-------
 .../binary/BinaryMemcacheDecoderTest.java     | 21 +++++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/codec-memcache/src/main/java/io/netty/handler/codec/memcache/binary/AbstractBinaryMemcacheDecoder.java b/codec-memcache/src/main/java/io/netty/handler/codec/memcache/binary/AbstractBinaryMemcacheDecoder.java
index eb5c5c9c7c..47e77eef4d 100644
--- a/codec-memcache/src/main/java/io/netty/handler/codec/memcache/binary/AbstractBinaryMemcacheDecoder.java
+++ b/codec-memcache/src/main/java/io/netty/handler/codec/memcache/binary/AbstractBinaryMemcacheDecoder.java
@@ -79,6 +79,7 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
                 currentMessage = decodeHeader(in);
                 state = State.READ_EXTRAS;
             } catch (Exception e) {
+                resetDecoder();
                 out.add(invalidMessage(e));
                 return;
             }
@@ -94,6 +95,7 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
 
                 state = State.READ_KEY;
             } catch (Exception e) {
+                resetDecoder();
                 out.add(invalidMessage(e));
                 return;
             }
@@ -107,10 +109,10 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
                     currentMessage.setKey(in.toString(in.readerIndex(), keyLength, CharsetUtil.UTF_8));
                     in.skipBytes(keyLength);
                 }
-
-                out.add(currentMessage);
+                out.add(currentMessage.retain());
                 state = State.READ_CONTENT;
             } catch (Exception e) {
+                resetDecoder();
                 out.add(invalidMessage(e));
                 return;
             }
@@ -150,9 +152,11 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
                     out.add(LastMemcacheContent.EMPTY_LAST_CONTENT);
                 }
 
+                resetDecoder();
                 state = State.READ_HEADER;
                 return;
             } catch (Exception e) {
+                resetDecoder();
                 out.add(invalidChunk(e));
                 return;
             }
@@ -200,10 +204,6 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
     public void channelInactive(ChannelHandlerContext ctx) throws Exception {
         super.channelInactive(ctx);
 
-        if (currentMessage != null) {
-            currentMessage.release();
-        }
-
         resetDecoder();
     }
 
@@ -211,7 +211,10 @@ public abstract class AbstractBinaryMemcacheDecoder<M extends BinaryMemcacheMess
      * Prepare for next decoding iteration.
      */
     protected void resetDecoder() {
-        currentMessage = null;
+        if (currentMessage != null) {
+            currentMessage.release();
+            currentMessage = null;
+        }
         alreadyReadChunkSize = 0;
     }
 
diff --git a/codec-memcache/src/test/java/io/netty/handler/codec/memcache/binary/BinaryMemcacheDecoderTest.java b/codec-memcache/src/test/java/io/netty/handler/codec/memcache/binary/BinaryMemcacheDecoderTest.java
index 1a27f43952..d7f15f31ac 100644
--- a/codec-memcache/src/test/java/io/netty/handler/codec/memcache/binary/BinaryMemcacheDecoderTest.java
+++ b/codec-memcache/src/test/java/io/netty/handler/codec/memcache/binary/BinaryMemcacheDecoderTest.java
@@ -30,6 +30,7 @@ import static org.hamcrest.CoreMatchers.*;
 import static org.hamcrest.MatcherAssert.*;
 import static org.hamcrest.core.IsNull.notNullValue;
 import static org.hamcrest.core.IsNull.nullValue;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Verifies the correct functionality of the {@link AbstractBinaryMemcacheDecoder}.
@@ -248,4 +249,24 @@ public class BinaryMemcacheDecoderTest {
         assertThat(content.content().toString(CharsetUtil.UTF_8), is(msgBody));
         content.release();
     }
+
+    @Test
+    public void shouldRetainCurrentMessageWhenSendingItOut() {
+        channel = new EmbeddedChannel(
+                new BinaryMemcacheRequestEncoder(),
+                new BinaryMemcacheRequestDecoder());
+
+        String key = "Netty";
+        ByteBuf extras = Unpooled.copiedBuffer("extras", CharsetUtil.UTF_8);
+        BinaryMemcacheRequest request = new DefaultBinaryMemcacheRequest(key, extras);
+        request.setKeyLength((short) key.length());
+        request.setExtrasLength((byte) extras.readableBytes());
+
+        assertTrue(channel.writeOutbound(request));
+        assertTrue(channel.writeInbound(channel.outboundMessages().toArray()));
+
+        BinaryMemcacheRequest read = channel.readInbound();
+        read.release();
+        // tearDown will call "channel.finish()"
+    }
 }