From a71e33ae35a5e1ce7a0cfe14aafd33c62aa6a075 Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill@us.ibm.com>
Date: Fri, 22 Mar 2019 03:18:10 -0700
Subject: [PATCH] Fix possible ByteBuf leak when CompositeByteBuf is resized
 (#8946)

Motivation:

The special case fixed in #8497 also requires that we keep a derived slice when trimming components in place, as done by the capacity(int) and discardReadBytes() methods.

Modifications:

Ensure that we keep a ref to trimmed components' original retained slice in capacity(int) and discardReadBytes() methods, so that it is released properly when the they are later freed. Add unit test which fails prior to the fix.

Result:

Edge case leak is eliminated.
---
 .../io/netty/buffer/CompositeByteBuf.java     | 15 ++++++--
 .../buffer/AbstractCompositeByteBufTest.java  | 34 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java b/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
index d7a1c6acd8..40c4ad2da0 100644
--- a/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
+++ b/buffer/src/main/java/io/netty/buffer/CompositeByteBuf.java
@@ -752,7 +752,12 @@ public class CompositeByteBuf extends AbstractReferenceCountedByteBuf implements
                 if (bytesToTrim < cLength) {
                     // Trim the last component
                     c.endOffset -= bytesToTrim;
-                    c.slice = null;
+                    ByteBuf slice = c.slice;
+                    if (slice != null) {
+                        // We must replace the cached slice with a derived one to ensure that
+                        // it can later be released properly in the case of PooledSlicedByteBuf.
+                        c.slice = slice.slice(0, c.length());
+                    }
                     break;
                 }
                 c.free();
@@ -1728,10 +1733,16 @@ public class CompositeByteBuf extends AbstractReferenceCountedByteBuf implements
             }
             firstComponentId++;
         } else {
+            int trimmedBytes = readerIndex - c.offset;
             c.offset = 0;
             c.endOffset -= readerIndex;
             c.adjustment += readerIndex;
-            c.slice = null;
+            ByteBuf slice = c.slice;
+            if (slice != null) {
+                // We must replace the cached slice with a derived one to ensure that
+                // it can later be released properly in the case of PooledSlicedByteBuf.
+                c.slice = slice.slice(trimmedBytes, c.length());
+            }
         }
 
         removeCompRange(0, firstComponentId);
diff --git a/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java b/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
index e25662dbe1..ee2a66ee12 100644
--- a/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
+++ b/buffer/src/test/java/io/netty/buffer/AbstractCompositeByteBufTest.java
@@ -1242,6 +1242,40 @@ public abstract class AbstractCompositeByteBufTest extends AbstractByteBufTest {
         assertEquals(0, b2.refCnt());
     }
 
+    @Test
+    public void testReleasesOnShrink2() {
+        // 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();
+
+        buffer.writeShort(1).writeShort(2);
+
+        ByteBuf b1 = buffer.readRetainedSlice(2);
+        ByteBuf b2 = b1.retainedSlice(b1.readerIndex(), 2);
+
+        // composite takes ownership of b1 and b2
+        ByteBuf composite = Unpooled.compositeBuffer()
+            .addComponents(b1, b2);
+
+        assertEquals(4, composite.capacity());
+
+        // reduce capacity down to two, will drop the second component
+        composite.capacity(2);
+        assertEquals(2, composite.capacity());
+
+        // releasing composite should release the components
+        composite.release();
+        assertEquals(0, composite.refCnt());
+        assertEquals(0, b1.refCnt());
+        assertEquals(0, b2.refCnt());
+
+        // release last remaining ref to buffer
+        buffer.release();
+        assertEquals(0, buffer.refCnt());
+    }
+
     @Test
     public void testAllocatorIsSameWhenCopy() {
         testAllocatorIsSameWhenCopy(false);