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,8 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufUtil;
-import io.netty.buffer.CompositeByteBuf;
 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 +32,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;
 
@@ -56,7 +55,20 @@ public abstract class DigestManager {
 
     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) {
+        MutableInt digestRef = new MutableInt(digest);
+        ByteBufVisitor.visitBuffers(buffer, offset, len,
+                (ByteBuf childBuffer, int childIndex, int childLength) -> {
+                    if (childLength > 0) {
+                        // recursively visit the sub buffer and update the digest
+                        int updatedDigest = internalUpdate(digestRef.intValue(), childBuffer, childIndex, childLength);
+                        digestRef.setValue(updatedDigest);
+                    }
+                });
+        return digestRef.intValue();
+    }
 
     abstract void populateValueAndReset(int digest, ByteBuf buffer);
 
@@ -136,34 +148,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 +173,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();