Fix checksum calculation bug when the payload is a CompositeByteBuf with readerIndex > 0

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/checksum/CRC32CDigestManager.java

@@ -46,7 +46,7 @@ void populateValueAndReset(int digest, ByteBuf buf) {
     }
 
     @Override
-    int update(int digest, ByteBuf data, int offset, int len) {
+    int internalUpdate(int digest, ByteBuf data, int offset, int len) {
         return Crc32cIntChecksum.resumeChecksum(digest, data, offset, len);
     }
 }

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/checksum/CRC32DigestManager.java

@@ -62,7 +62,7 @@ void populateValueAndReset(int digest, ByteBuf buf) {
     }
 
     @Override
-    int update(int digest, ByteBuf data, int offset, int len) {
+    int internalUpdate(int digest, ByteBuf data, int offset, int len) {
         crc.get().update(data, offset, len);
         return 0;
     }

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/checksum/DigestManager.java

@@ -20,10 +20,9 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufUtil;
-import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.DuplicatedByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
 import io.netty.buffer.Unpooled;
-import io.netty.util.ReferenceCountUtil;
 import io.netty.util.ReferenceCounted;
 import io.netty.util.concurrent.FastThreadLocal;
 import java.security.GeneralSecurityException;
@@ -34,6 +33,7 @@
 import org.apache.bookkeeper.proto.BookieProtocol;
 import org.apache.bookkeeper.proto.DataFormats.LedgerMetadataFormat.DigestType;
 import org.apache.bookkeeper.util.ByteBufList;
+import org.apache.commons.lang3.mutable.MutableInt;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -49,14 +49,89 @@ public abstract class DigestManager {
 
     public static final int METADATA_LENGTH = 32;
     public static final int LAC_METADATA_LENGTH = 16;
+    private static final int MAX_SUB_BUFFER_VISIT_RECURSION_DEPTH = 10;
 
     final long ledgerId;
     final boolean useV2Protocol;
     private final ByteBufAllocator allocator;
 
     abstract int getMacCodeLength();
 
-    abstract int update(int digest, ByteBuf buffer, int offset, int len);
+    abstract int internalUpdate(int digest, ByteBuf buffer, int offset, int len);
+
+    final int update(int digest, ByteBuf buffer, int offset, int len) {
+        return recursiveSubBufferVisitForDigestUpdate(digest, buffer, offset, len, 0);
+    }
+
+    private int recursiveSubBufferVisitForDigestUpdate(int digest, ByteBuf buffer, int offset, int len, int depth) {
+        if (len == 0) {
+            return digest;
+        }
+        if (depth < MAX_SUB_BUFFER_VISIT_RECURSION_DEPTH && !buffer.hasMemoryAddress() && !buffer.hasArray()) {
+            return visitWrappedBuffersAndCallInternalUpdateForEach(digest, buffer, offset, len, depth);
+        }
+        return internalUpdate(digest, buffer, offset, len);
+    }
+
+    /**
+     * This method is used to visit the wrapped buffers and call internalUpdate for each of them.
+     * CompositeByteBuf is one of the wrapped buffer types that will be visited. It can contain multiple
+     * sub buffers. The sub buffers can be wrapped buffers as well.
+     *
+     * Netty doesn't provide an API to visit the wrapped buffers, so we have to use the a hack here.
+     *
+     * @param digest current digest value
+     * @param buffer the buffer to visit
+     * @param offset the offset in the buffer
+     * @param len the length in the buffer
+     * @param depth the recursion depth of the wrapped buffer visit
+     * @return updated digest value
+     */
+    private int visitWrappedBuffersAndCallInternalUpdateForEach(int digest, ByteBuf buffer, int offset, int len,
+                                                                int depth) {
+        // hold the digest in a MutableInt so that it can be updated in the wrapped buffer visit
+        MutableInt digestRef = new MutableInt(digest);
+        ByteBuf sliced = buffer.slice(offset, len);
+        // call getBytes to trigger the wrapped buffer visit
+        sliced.getBytes(0, new DuplicatedByteBuf(Unpooled.EMPTY_BUFFER) {
+            @Override
+            public ByteBuf setBytes(int index, ByteBuf src, int srcIndex, int length) {
+                if (length > 0) {
+                    // recursively visit the sub buffer and update the digest
+                    int updatedDigest = recursiveSubBufferVisitForDigestUpdate(digestRef.intValue(), src, srcIndex,
+                            length, depth + 1);
+                    digestRef.setValue(updatedDigest);
+                }
+                return this;
+            }
+
+            @Override
+            public boolean hasArray() {
+                // return false so that the wrapped buffer is visited
+                return false;
+            }
+
+            @Override
+            public boolean hasMemoryAddress() {
+                // return false so that the wrapped buffer is visited
+                return false;
+            }
+
+            @Override
+            public int nioBufferCount() {
+                // return 0 so that the wrapped buffer is visited
+                return 0;
+            }
+
+            @Override
+            public int capacity() {
+                // should return sufficient capacity for the total length
+                return len;
+            }
+        });
+        // return updated digest value
+        return digestRef.intValue();
+    }
 
     abstract void populateValueAndReset(int digest, ByteBuf buffer);
 
@@ -136,34 +211,19 @@ private ReferenceCounted computeDigestAndPackageForSendingV2(long entryId, long
 
         // Compute checksum over the headers
         int digest = update(0, buf, buf.readerIndex(), buf.readableBytes());
-
-        // don't unwrap slices
-        final ByteBuf unwrapped = data.unwrap() != null && data.unwrap() instanceof CompositeByteBuf
-                ? data.unwrap() : data;
-        ReferenceCountUtil.retain(unwrapped);
-        ReferenceCountUtil.safeRelease(data);
-
-        if (unwrapped instanceof CompositeByteBuf) {
-            CompositeByteBuf cbb = (CompositeByteBuf) unwrapped;
-            for (int i = 0; i < cbb.numComponents(); i++) {
-                ByteBuf b = cbb.component(i);
-                digest = update(digest, b, b.readerIndex(), b.readableBytes());
-            }
-        } else {
-            digest = update(digest, unwrapped, unwrapped.readerIndex(), unwrapped.readableBytes());
-        }
+        digest = update(digest, data, data.readerIndex(), data.readableBytes());
 
         populateValueAndReset(digest, buf);
 
         // Reset the reader index to the beginning
         buf.readerIndex(0);
 
         if (isSmallEntry) {
-            buf.writeBytes(unwrapped, unwrapped.readerIndex(), unwrapped.readableBytes());
-            unwrapped.release();
+            buf.writeBytes(data, data.readerIndex(), data.readableBytes());
+            data.release();
             return buf;
         } else {
-            return ByteBufList.get(buf, unwrapped);
+            return ByteBufList.get(buf, data);
         }
     }
 
@@ -176,25 +236,9 @@ private ByteBufList computeDigestAndPackageForSendingV3(long entryId, long lastA
         headersBuffer.writeLong(length);
 
         int digest = update(0, headersBuffer, 0, METADATA_LENGTH);
-
-        // don't unwrap slices
-        final ByteBuf unwrapped = data.unwrap() != null && data.unwrap() instanceof CompositeByteBuf
-                ? data.unwrap() : data;
-        ReferenceCountUtil.retain(unwrapped);
-        ReferenceCountUtil.release(data);
-
-        if (unwrapped instanceof CompositeByteBuf) {
-            CompositeByteBuf cbb = ((CompositeByteBuf) unwrapped);
-            for (int i = 0; i < cbb.numComponents(); i++) {
-                ByteBuf b = cbb.component(i);
-                digest = update(digest, b, b.readerIndex(), b.readableBytes());
-            }
-        } else {
-            digest = update(digest, unwrapped, unwrapped.readerIndex(), unwrapped.readableBytes());
-        }
+        digest = update(digest, data, data.readerIndex(), data.readableBytes());
         populateValueAndReset(digest, headersBuffer);
-
-        return ByteBufList.get(headersBuffer, unwrapped);
+        return ByteBufList.get(headersBuffer, data);
     }
 
     /**

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/checksum/DummyDigestManager.java

@@ -38,7 +38,7 @@ int getMacCodeLength() {
     }
 
     @Override
-    int update(int digest, ByteBuf buffer, int offset, int len) {
+    int internalUpdate(int digest, ByteBuf buffer, int offset, int len) {
         return 0;
     }
 

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/checksum/MacDigestManager.java

@@ -96,7 +96,7 @@ void populateValueAndReset(int digest, ByteBuf buffer) {
     }
 
     @Override
-    int update(int digest, ByteBuf data, int offset, int len) {
+    int internalUpdate(int digest, ByteBuf data, int offset, int len) {
         mac.get().update(data.slice(offset, len).nioBuffer());
         return 0;
     }

bookkeeper-server/src/main/java/org/apache/bookkeeper/util/ByteBufList.java

@@ -22,7 +22,6 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.CompositeByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandler.Sharable;
 import io.netty.channel.ChannelHandlerContext;
@@ -133,43 +132,14 @@ public static ByteBufList get() {
      * Append a {@link ByteBuf} at the end of this {@link ByteBufList}.
      */
     public void add(ByteBuf buf) {
-        final ByteBuf unwrapped = buf.unwrap() != null && buf.unwrap() instanceof CompositeByteBuf
-                ? buf.unwrap() : buf;
-        ReferenceCountUtil.retain(unwrapped);
-        ReferenceCountUtil.release(buf);
-
-        if (unwrapped instanceof CompositeByteBuf) {
-            ((CompositeByteBuf) unwrapped).forEach(b -> {
-                ReferenceCountUtil.retain(b);
-                buffers.add(b);
-            });
-            ReferenceCountUtil.release(unwrapped);
-        } else {
-            buffers.add(unwrapped);
-        }
+        buffers.add(buf);
     }
 
     /**
      * Prepend a {@link ByteBuf} at the beginning of this {@link ByteBufList}.
      */
     public void prepend(ByteBuf buf) {
-        // don't unwrap slices
-        final ByteBuf unwrapped = buf.unwrap() != null && buf.unwrap() instanceof CompositeByteBuf
-                ? buf.unwrap() : buf;
-        ReferenceCountUtil.retain(unwrapped);
-        ReferenceCountUtil.release(buf);
-
-        if (unwrapped instanceof CompositeByteBuf) {
-            CompositeByteBuf composite = (CompositeByteBuf) unwrapped;
-            for (int i = composite.numComponents() - 1; i >= 0; i--) {
-                ByteBuf b = composite.component(i);
-                ReferenceCountUtil.retain(b);
-                buffers.add(0, b);
-            }
-            ReferenceCountUtil.release(unwrapped);
-        } else {
-            buffers.add(0, unwrapped);
-        }
+        buffers.add(0, buf);
     }
 
     /**
@@ -285,7 +255,7 @@ public ByteBufList retain() {
     @Override
     protected void deallocate() {
         for (int i = 0; i < buffers.size(); i++) {
-            ReferenceCountUtil.release(buffers.get(i));
+            buffers.get(i).release();
         }
 
         buffers.clear();