From 89d1c5722ea4d8609351efe14b0d4ea25a53be63 Mon Sep 17 00:00:00 2001
From: Martin Dindoffer <contact@dindoffer.eu>
Date: Wed, 4 May 2022 18:40:12 +0200
Subject: [PATCH] Fix integer overflow in computeSerializedSizeInWords

---
 .../src/main/java/org/capnproto/Serialize.java   | 11 +++++++----
 .../test/java/org/capnproto/SerializeTest.java   | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/runtime/src/main/java/org/capnproto/Serialize.java b/runtime/src/main/java/org/capnproto/Serialize.java
index 6dcccad..4329dfc 100644
--- a/runtime/src/main/java/org/capnproto/Serialize.java
+++ b/runtime/src/main/java/org/capnproto/Serialize.java
@@ -178,15 +178,18 @@ public final class Serialize {
     }
 
     public static long computeSerializedSizeInWords(MessageBuilder message) {
-        final ByteBuffer[] segments = message.getSegmentsForOutput();
+        return computeSerializedSizeInWords(message.getSegmentsForOutput());
+    }
 
-        // From the capnproto documentation:
+    //VisibleForTesting
+    static long computeSerializedSizeInWords(ByteBuffer[] segments) {
+        // From the capnproto documentation (https://capnproto.org/encoding.html#serialization-over-a-stream):
         // "When transmitting over a stream, the following should be sent..."
         long bytes = 0;
         // "(4 bytes) The number of segments, minus one..."
         bytes += 4;
         // "(N * 4 bytes) The size of each segment, in words."
-        bytes += segments.length * 4;
+        bytes += segments.length * 4L;
         // "(0 or 4 bytes) Padding up to the next word boundary."
         if (bytes % 8 != 0) {
             bytes += 4;
@@ -194,7 +197,7 @@ public final class Serialize {
 
         // The content of each segment, in order.
         for (int i = 0; i < segments.length; ++i) {
-            final ByteBuffer s = segments[i];
+            ByteBuffer s = segments[i];
             bytes += s.limit();
         }
 
diff --git a/runtime/src/test/java/org/capnproto/SerializeTest.java b/runtime/src/test/java/org/capnproto/SerializeTest.java
index 5ad9abe..e9051a5 100644
--- a/runtime/src/test/java/org/capnproto/SerializeTest.java
+++ b/runtime/src/test/java/org/capnproto/SerializeTest.java
@@ -30,8 +30,15 @@ import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 
 import org.junit.Assert;
+import org.junit.Ignore;
 import org.junit.Test;
 
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
 public class SerializeTest {
 
   /**
@@ -194,4 +201,13 @@ public class SerializeTest {
             java.nio.channels.Channels.newChannel(new java.io.ByteArrayInputStream(input));
         MessageReader message = Serialize.read(channel);
   }
+
+    @Test
+    @Ignore("Ignored by default because the huge array used in the test results in a long execution")
+    public void computeSerializedSizeInWordsShouldNotOverflowOnLargeSegmentCounts() {
+        ByteBuffer dummySegmentBuffer = ByteBuffer.allocate(0);
+        ByteBuffer[] segments = new ByteBuffer[Integer.MAX_VALUE / 2];
+        Arrays.fill(segments, dummySegmentBuffer);
+        assertThat(Serialize.computeSerializedSizeInWords(segments), is((segments.length * 4L + 4) / Constants.BYTES_PER_WORD));
+    }
 }