From 804e1fa9ccf426603e6d478f63c63bc00baae0fa Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill@us.ibm.com>
Date: Tue, 13 Nov 2018 11:56:09 -0800
Subject: [PATCH] Fix ref-counting when CompositeByteBuf is used with
 retainedSlice() (#8497)

Motivation:

ByteBuf.retainedSlice() and similar methods produce sliced buffers with
an independent refcount to the buffer that they wrap.


One of the optimizations in 10539f4dc738c48e8d63c46b72ca32906d7f40ec was
to use the ref to the unwrapped buffer object for added slices, but this
did not take into account the above special case when later releasing.

Thanks to @rkapsi for discovering this via #8495.

Modifications:

Since a reference to the slice is still kept in the Component class,
just changed Component.freeIfNecessary() to release the slice in
preference to the unwrapped buf.

Also added a unit test which reproduces the bug.

Result:

Fixes #8495
---
 .../io/netty/buffer/CompositeByteBuf.java     | 13 +++++--
 .../buffer/AbstractCompositeByteBufTest.java  | 34 +++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java b/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
index 77ba13fcae..c9212d5c6e 100644
--- a/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
+++ b/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
@@ -513,11 +513,11 @@ public class CompositeByteBuf extends AbstractReferenceCountedByteBuf implements
     public CompositeByteBuf removeComponent(int cIndex) {
         checkComponentIndex(cIndex);
         Component comp = components[cIndex];
-        removeComp(cIndex);
         if (lastAccessed == comp) {
             lastAccessed = null;
         }
         comp.freeIfNecessary();
+        removeComp(cIndex);
         if (comp.length() > 0) {
             // Only need to call updateComponentOffsets if the length was > 0
             updateComponentOffsets(cIndex);
@@ -1815,7 +1815,15 @@ public class CompositeByteBuf extends AbstractReferenceCountedByteBuf implements
         }
 
         void freeIfNecessary() {
-            buf.release(); // We should not get a NPE here. If so, it must be a bug.
+            // Release the slice if present since it may have a different
+            // refcount to the unwrapped buf if it is a PooledSlicedByteBuf
+            ByteBuf buffer = slice;
+            if (buffer != null) {
+                buffer.release();
+            } else {
+                buf.release();
+            }
+            // null out in either case since it could be racy
             slice = null;
         }
     }
@@ -2181,7 +2189,6 @@ public class CompositeByteBuf extends AbstractReferenceCountedByteBuf implements
         }
         int newSize = size - to + from;
         for (int i = newSize; i < size; i++) {
-            components[i].slice = null;
             components[i] = null;
         }
         componentCount = newSize;
diff --git a/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java b/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
index 363b25d8c7..8e0f148cdf 100644
--- a/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
+++ b/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
@@ -1144,6 +1144,40 @@ public abstract class AbstractCompositeByteBufTest extends AbstractByteBufTest {
         assertEquals(0, buffer.refCnt());
     }
 
+    @Test
+    public void testReleasesItsComponents2() {
+        // It is important to use a pooled allocator here to ensure
+        // the slices returned by readRetainedSlice are of type
+        // PooledSlicedByteBuf, which maintains an independent refcount
+        // (so that we can be sure to cover this case)
+        ByteBuf buffer = PooledByteBufAllocator.DEFAULT.buffer(); // 1
+
+        buffer.writeBytes(new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10});
+
+        // use readRetainedSlice this time - produces different kind of slices
+        ByteBuf s1 = buffer.readRetainedSlice(2); // 2
+        ByteBuf s2 = s1.readRetainedSlice(2); // 3
+        ByteBuf s3 = s2.readRetainedSlice(2); // 4
+        ByteBuf s4 = s3.readRetainedSlice(2); // 5
+
+        ByteBuf composite = Unpooled.compositeBuffer()
+            .addComponent(s1)
+            .addComponents(s2, s3, s4)
+            .order(ByteOrder.LITTLE_ENDIAN);
+
+        assertEquals(1, composite.refCnt());
+        assertEquals(2, buffer.refCnt());
+
+        // releasing composite should release the 4 components
+        composite.release();
+        assertEquals(0, composite.refCnt());
+        assertEquals(1, buffer.refCnt());
+
+        // last remaining ref to buffer
+        buffer.release();
+        assertEquals(0, buffer.refCnt());
+    }
+
     @Test
     public void testReleasesOnShrink() {