prevent packer's buffer from getting into unmarked state

richardstartin · richardstartin · commit 0df0579b26e4 · 2020-07-27T21:30:00.000+01:00
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/serialization/msgpack/Packer.java b/dd-trace-core/src/main/java/datadog/trace/core/serialization/msgpack/Packer.java
@@ -60,14 +60,18 @@ public Packer(Codec codec, ByteBufferConsumer sink, ByteBuffer buffer) {
     this.codec = codec;
     this.sink = sink;
     this.buffer = buffer;
-    this.buffer.position(MAX_ARRAY_HEADER_SIZE);
-    buffer.mark();
+    initBuffer();
   }
 
   public Packer(ByteBufferConsumer sink, ByteBuffer buffer) {
     this(Codec.INSTANCE, sink, buffer);
   }
 
+  private void initBuffer() {
+    this.buffer.position(MAX_ARRAY_HEADER_SIZE);
+    buffer.mark();
+  }
+
   @Override
   public <T> void format(T message, Mapper<T> mapper) {
     try {
@@ -118,7 +122,7 @@ public void flush() {
     writeArrayHeader(messageCount);
     buffer.position(pos);
     sink.accept(messageCount, buffer.slice());
-    buffer.position(MAX_ARRAY_HEADER_SIZE);
+    initBuffer();
     buffer.limit(buffer.capacity());
     messageCount = 0;
   }
diff --git a/dd-trace-core/src/test/java/datadog/trace/core/serialization/msgpack/VerifyThatPackerDoesNotThrowInvalidMarkException.java b/dd-trace-core/src/test/java/datadog/trace/core/serialization/msgpack/VerifyThatPackerDoesNotThrowInvalidMarkException.java
@@ -0,0 +1,58 @@
+package datadog.trace.core.serialization.msgpack;
+
+import java.nio.BufferOverflowException;
+import java.nio.ByteBuffer;
+import java.nio.InvalidMarkException;
+import java.util.concurrent.ThreadLocalRandom;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class VerifyThatPackerDoesNotThrowInvalidMarkException {
+
+  @Parameterized.Parameters
+  public static Object[][] bufferSizes() {
+    return new Object[][] {{10}, {16}, {100}, {128}, {1000}, {1024}};
+  }
+
+  private final int bufferSize;
+
+  public VerifyThatPackerDoesNotThrowInvalidMarkException(int bufferSize) {
+    this.bufferSize = bufferSize;
+  }
+
+  @Test
+  public void provokeInvalidMarkException() {
+    ByteBuffer buffer = ByteBuffer.allocate(bufferSize);
+    Packer packer =
+        new Packer(
+            new ByteBufferConsumer() {
+              @Override
+              public void accept(int messageCount, ByteBuffer buffer) {}
+            },
+            buffer);
+    Mapper<byte[]> mapper =
+        new Mapper<byte[]>() {
+          @Override
+          public void map(byte[] data, Writable packer) {
+            packer.writeBinary(data, 0, data.length);
+          }
+        };
+    int position = buffer.position();
+    int length = 0;
+    for (int i = 0; i < 10_000; ++i) {
+      try {
+        byte[] message = new byte[ThreadLocalRandom.current().nextInt(0, bufferSize * 2)];
+        position = buffer.position();
+        length = message.length;
+        packer.format(message, mapper);
+      } catch (BufferOverflowException e) {
+        // acceptable, need to be able to handle data larger than the limit
+      } catch (InvalidMarkException e) {
+        Assert.fail(e.getMessage() + " " + position + " " + length);
+      }
+    }
+  }
+}
