fix: harden arithmetic and bounds checks

extensions-contrib/aliyun-oss-extensions/src/test/java/org/apache/druid/data/input/aliyun/OssInputSourceTest.java

@@ -551,7 +551,7 @@ public void testReader() throws IOException
       Assert.assertEquals("world", nextRow.getDimension("dim2").get(0));
     }
 
-    Assert.assertEquals(2 * CONTENT.length, inputStats.getProcessedBytes());
+    Assert.assertEquals(2L * CONTENT.length, inputStats.getProcessedBytes());
     EasyMock.verify(OSSCLIENT);
   }
 
@@ -598,7 +598,7 @@ public void testCompressedReader() throws IOException
       Assert.assertEquals("world", nextRow.getDimension("dim2").get(0));
     }
 
-    Assert.assertEquals(2 * CONTENT.length, inputStats.getProcessedBytes());
+    Assert.assertEquals(2L * CONTENT.length, inputStats.getProcessedBytes());
     EasyMock.verify(OSSCLIENT);
   }
 

extensions-contrib/aliyun-oss-extensions/src/test/java/org/apache/druid/storage/aliyun/OssTaskLogsTest.java

@@ -347,7 +347,8 @@ public void test_taskLog_fetch_withNegativeRange() throws IOException
     EasyMock.replay(ossClient);
 
     OssTaskLogs ossTaskLogs = getOssTaskLogs();
-    Optional<InputStream> inputStreamOptional = ossTaskLogs.streamTaskLog(KEY_1, -1 * (LOG_CONTENTS.length() - 1));
+    final Optional<InputStream> inputStreamOptional =
+        ossTaskLogs.streamTaskLog(KEY_1, 1L - LOG_CONTENTS.length());
     String taskLogs = new BufferedReader(
         new InputStreamReader(inputStreamOptional.get(), StandardCharsets.UTF_8))
         .lines()

extensions-contrib/compressed-bigdecimal/src/main/java/org/apache/druid/compressedbigdecimal/ByteBufferCompressedBigDecimal.java

@@ -119,9 +119,9 @@ public int getArraySize()
    * @return the entry
    */
   @Override
-  public int getArrayEntry(int idx)
+  public int getArrayEntry(final int idx)
   {
-    return buf.getInt(position + idx * Integer.BYTES);
+    return buf.getInt(bufferPosition(idx));
   }
 
   /**
@@ -131,9 +131,15 @@ public int getArrayEntry(int idx)
    * @param val value to set
    */
   @Override
-  protected void setArrayEntry(int idx, int val)
+  protected void setArrayEntry(final int idx, final int val)
   {
-    buf.putInt(position + idx * Integer.BYTES, val);
+    buf.putInt(bufferPosition(idx), val);
+  }
+
+  private int bufferPosition(final int idx)
+  {
+    Preconditions.checkElementIndex(idx, size);
+    return Math.addExact(position, Math.multiplyExact(idx, Integer.BYTES));
   }
 
   @Override

extensions-contrib/compressed-bigdecimal/src/main/java/org/apache/druid/compressedbigdecimal/aggregator/CompressedBigDecimalAggregatorFactoryBase.java

@@ -86,9 +86,9 @@ protected CompressedBigDecimalAggregatorFactoryBase(
     this.scale = scale == null ? DEFAULT_SCALE : scale;
     this.strictNumberParsing = strictNumberParsing == null ? DEFAULT_STRICT_NUMBER_PARSING : strictNumberParsing;
 
-    byte[] fieldNameBytes = StringUtils.toUtf8(fieldName);
+    final byte[] fieldNameBytes = StringUtils.toUtf8(fieldName);
 
-    cacheKey = ByteBuffer.allocate(1 + fieldNameBytes.length + 2 * Integer.BYTES + 1)
+    cacheKey = ByteBuffer.allocate(Math.addExact(fieldNameBytes.length, 2 * Integer.BYTES + 2))
                          .put(cacheTypeId)
                          .put(fieldNameBytes)
                          .putInt(this.size)

extensions-contrib/ddsketch/src/main/java/org/apache/druid/query/aggregation/ddsketch/DDSketchObjectStrategy.java

@@ -22,6 +22,7 @@
 import com.datadoghq.sketch.ddsketch.DDSketch;
 import com.datadoghq.sketch.ddsketch.DDSketchProtoBinding;
 import com.datadoghq.sketch.ddsketch.store.CollapsingLowestDenseStore;
+import com.google.common.base.Preconditions;
 import com.google.protobuf.InvalidProtocolBufferException;
 import org.apache.druid.segment.data.ObjectStrategy;
 
@@ -44,11 +45,18 @@ public DDSketch fromByteBuffer(ByteBuffer buffer, int numBytes)
     if (numBytes == 0) {
       return null;
     }
-    ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
-    readOnlyBuffer.limit(buffer.position() + numBytes);
+    Preconditions.checkArgument(
+        numBytes >= 0 && numBytes <= buffer.remaining(),
+        "numBytes[%s] exceeds buffer remaining[%s]",
+        numBytes,
+        buffer.remaining()
+    );
+    final ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
+    readOnlyBuffer.limit(Math.addExact(buffer.position(), numBytes));
     try {
-      com.datadoghq.sketch.ddsketch.proto.DDSketch proto = com.datadoghq.sketch.ddsketch.proto.DDSketch.parseFrom(readOnlyBuffer);
-      DDSketch recovered = DDSketchProtoBinding.fromProto(() -> new CollapsingLowestDenseStore(1000), proto);
+      final com.datadoghq.sketch.ddsketch.proto.DDSketch proto =
+          com.datadoghq.sketch.ddsketch.proto.DDSketch.parseFrom(readOnlyBuffer);
+      final DDSketch recovered = DDSketchProtoBinding.fromProto(() -> new CollapsingLowestDenseStore(1000), proto);
       return recovered;
     } 
     catch (InvalidProtocolBufferException e) {

extensions-contrib/momentsketch/src/main/java/org/apache/druid/query/aggregation/momentsketch/MomentSketchObjectStrategy.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.query.aggregation.momentsketch;
 
+import com.google.common.base.Preconditions;
 import org.apache.druid.query.aggregation.momentsketch.aggregator.MomentSketchAggregatorFactory;
 import org.apache.druid.segment.data.ObjectStrategy;
 
@@ -41,7 +42,13 @@ public MomentSketchWrapper fromByteBuffer(ByteBuffer buffer, int numBytes)
     if (numBytes == 0) {
       return null;
     }
-    buffer.limit(buffer.position() + numBytes);
+    Preconditions.checkArgument(
+        numBytes > 0 && numBytes <= buffer.remaining(),
+        "numBytes[%s] exceeds buffer remaining[%s]",
+        numBytes,
+        buffer.remaining()
+    );
+    buffer.limit(Math.addExact(buffer.position(), numBytes));
     return MomentSketchWrapper.fromBytes(buffer);
   }
 

extensions-contrib/spectator-histogram/src/test/java/org/apache/druid/spectator/histogram/SpectatorHistogramTest.java

@@ -57,7 +57,7 @@ public void testToBytesSmallValues()
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
     int valSize = 0;
-    Assert.assertEquals("Should compact small values within key bytes", 5 * (keySize + valSize), bytes.length);
+    Assert.assertEquals("Should compact small values within key bytes", 5L * (keySize + valSize), bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(1L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -86,7 +86,7 @@ public void testToBytesSmallishValues()
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
     int valSize = Byte.BYTES;
-    Assert.assertEquals("Should compact small values to a byte", 5 * (keySize + valSize), bytes.length);
+    Assert.assertEquals("Should compact small values to a byte", 5L * (keySize + valSize), bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(64L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -115,7 +115,7 @@ public void testToBytesMedValues()
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
     int valSize = Short.BYTES;
-    Assert.assertEquals("Should compact medium values to short", 5 * (keySize + valSize), bytes.length);
+    Assert.assertEquals("Should compact medium values to short", 5L * (keySize + valSize), bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(512L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -144,7 +144,7 @@ public void testToBytesLargerValues()
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
     int valSize = Integer.BYTES;
-    Assert.assertEquals("Should compact larger values to integer", 5 * (keySize + valSize), bytes.length);
+    Assert.assertEquals("Should compact larger values to integer", 5L * (keySize + valSize), bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(100000L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -173,7 +173,7 @@ public void testToBytesBiggestValues()
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
     int valSize = Long.BYTES;
-    Assert.assertEquals("Should not compact larger values", 5 * (keySize + valSize), bytes.length);
+    Assert.assertEquals("Should not compact larger values", 5L * (keySize + valSize), bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(10000000000L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -201,7 +201,7 @@ public void testToBytesMixedValues()
 
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
-    Assert.assertEquals("Should not compact larger values", (5 * keySize) + 0 + 2 + 4 + 8 + 8, bytes.length);
+    Assert.assertEquals("Should not compact larger values", (5L * keySize) + 0 + 2 + 4 + 8 + 8, bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(1L, deserialized.get(PercentileBuckets.indexOf(10)));
@@ -232,7 +232,7 @@ public void testToBytesBoundaryValues()
 
     byte[] bytes = histogram.toBytes();
     int keySize = Short.BYTES;
-    Assert.assertEquals("Should compact", (8 * keySize) + 0 + 1 + 1 + 2 + 2 + 4 + 4 + 8, bytes.length);
+    Assert.assertEquals("Should compact", (8L * keySize) + 0 + 1 + 1 + 2 + 2 + 4 + 4 + 8, bytes.length);
 
     SpectatorHistogram deserialized = SpectatorHistogram.deserialize(bytes);
     Assert.assertEquals(63L, deserialized.get(6));

extensions-contrib/tdigestsketch/src/main/java/org/apache/druid/query/aggregation/tdigestsketch/TDigestSketchObjectStrategy.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.query.aggregation.tdigestsketch;
 
+import com.google.common.base.Preconditions;
 import com.tdunning.math.stats.MergingDigest;
 import org.apache.druid.segment.data.ObjectStrategy;
 
@@ -41,8 +42,14 @@ public MergingDigest fromByteBuffer(ByteBuffer buffer, int numBytes)
     if (numBytes == 0) {
       return null;
     }
-    ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
-    readOnlyBuffer.limit(buffer.position() + numBytes);
+    Preconditions.checkArgument(
+        numBytes > 0 && numBytes <= buffer.remaining(),
+        "numBytes[%s] exceeds buffer remaining[%s]",
+        numBytes,
+        buffer.remaining()
+    );
+    final ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
+    readOnlyBuffer.limit(Math.addExact(buffer.position(), numBytes));
     return MergingDigest.fromBytes(readOnlyBuffer);
   }
 

extensions-core/google-extensions/src/test/java/org/apache/druid/data/input/google/GoogleCloudStorageInputSourceTest.java

@@ -419,7 +419,7 @@ public void testReader() throws IOException
       Assert.assertEquals("hello", nextRow.getDimension("dim1").get(0));
       Assert.assertEquals("world", nextRow.getDimension("dim2").get(0));
     }
-    Assert.assertEquals(2 * CONTENT.length, inputStats.getProcessedBytes());
+    Assert.assertEquals(2L * CONTENT.length, inputStats.getProcessedBytes());
   }
 
   @Test
@@ -466,7 +466,7 @@ public void testCompressedReader() throws IOException
       Assert.assertEquals("hello", nextRow.getDimension("dim1").get(0));
       Assert.assertEquals("world", nextRow.getDimension("dim2").get(0));
     }
-    Assert.assertEquals(2 * CONTENT.length, inputStats.getProcessedBytes());
+    Assert.assertEquals(2L * CONTENT.length, inputStats.getProcessedBytes());
   }
 
   @Test

extensions-core/histogram/src/main/java/org/apache/druid/query/aggregation/histogram/ApproximateHistogram.java

@@ -1261,7 +1261,7 @@ public void toBytesCompact(ByteBuffer buf)
       for (int i = 0; i < binCount; ++i) {
         // repeat each value bins[i] times for approximate bins
         if ((bins[i] & APPROX_FLAG_BIT) != 0) {
-          for (int k = 0; k < (bins[i] & COUNT_BITS); ++k) {
+          for (long k = 0; k < (bins[i] & COUNT_BITS); ++k) {
             buf.putFloat(positions[i]);
           }
         }
@@ -1277,7 +1277,7 @@ public void toBytesCompact(ByteBuffer buf)
     for (int i = 0; i < binCount; ++i) {
       // repeat each value bins[i] times for exact bins
       if ((bins[i] & APPROX_FLAG_BIT) == 0) {
-        for (int k = 0; k < (bins[i] & COUNT_BITS); ++k) {
+        for (long k = 0; k < (bins[i] & COUNT_BITS); ++k) {
           buf.putFloat(positions[i]);
         }
       }
@@ -1371,7 +1371,7 @@ public static ApproximateHistogram fromBytesCompact(ByteBuffer buf)
       }
       return histogram;
     } else {
-      byte approxCount = (byte) (-1 * count);
+      final int approxCount = -count;
 
       Map<Float, Long> approx = new HashMap<>();
 

extensions-core/histogram/src/main/java/org/apache/druid/query/aggregation/histogram/ApproximateHistogramFoldingSerde.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.query.aggregation.histogram;
 
+import com.google.common.base.Preconditions;
 import it.unimi.dsi.fastutil.bytes.ByteArrays;
 import org.apache.druid.data.input.InputRow;
 import org.apache.druid.data.input.Rows;
@@ -93,7 +94,13 @@ public Class<ApproximateHistogram> getClazz()
       @Override
       public ApproximateHistogram fromByteBuffer(ByteBuffer buffer, int numBytes)
       {
-        buffer.limit(buffer.position() + numBytes);
+        Preconditions.checkArgument(
+            numBytes >= 0 && numBytes <= buffer.remaining(),
+            "numBytes[%s] exceeds buffer remaining[%s]",
+            numBytes,
+            buffer.remaining()
+        );
+        buffer.limit(Math.addExact(buffer.position(), numBytes));
         return ApproximateHistogram.fromBytes(buffer);
       }
 

extensions-core/histogram/src/main/java/org/apache/druid/query/aggregation/histogram/FixedBucketsHistogram.java

@@ -993,7 +993,7 @@ protected static FixedBucketsHistogram fromByteBufferFullNoSerdeHeader(ByteBuffe
     double lowerLimit = buf.getDouble();
     double upperLimit = buf.getDouble();
     int numBuckets = buf.getInt();
-    OutlierHandlingMode outlierHandlingMode = OutlierHandlingMode.values()[buf.get()];
+    final OutlierHandlingMode outlierHandlingMode = readOutlierHandlingMode(buf);
 
     long count = buf.getLong();
     long lowerOutlierCount = buf.getLong();
@@ -1034,7 +1034,7 @@ private static FixedBucketsHistogram fromBytesSparse(ByteBuffer buf)
     double lowerLimit = buf.getDouble();
     double upperLimit = buf.getDouble();
     int numBuckets = buf.getInt();
-    OutlierHandlingMode outlierHandlingMode = OutlierHandlingMode.values()[buf.get()];
+    final OutlierHandlingMode outlierHandlingMode = readOutlierHandlingMode(buf);
 
     long count = buf.getLong();
     long lowerOutlierCount = buf.getLong();
@@ -1067,6 +1067,21 @@ private static FixedBucketsHistogram fromBytesSparse(ByteBuffer buf)
     );
   }
 
+  private static OutlierHandlingMode readOutlierHandlingMode(final ByteBuffer buf)
+  {
+    final byte mode = buf.get();
+    switch (mode) {
+      case 0:
+        return OutlierHandlingMode.IGNORE;
+      case 1:
+        return OutlierHandlingMode.OVERFLOW;
+      case 2:
+        return OutlierHandlingMode.CLIP;
+      default:
+        throw new ISE("Invalid outlier handling mode [%s].", mode);
+    }
+  }
+
   /**
    * Compute the size in bytes of a full-encoding serialized histogram, without the serialization header
    *

extensions-core/histogram/src/main/java/org/apache/druid/query/aggregation/histogram/FixedBucketsHistogramSerde.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.query.aggregation.histogram;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.Ordering;
 import org.apache.druid.data.input.InputRow;
 import org.apache.druid.data.input.Rows;
@@ -134,8 +135,14 @@ public Class<? extends FixedBucketsHistogram> getClazz()
       @Override
       public FixedBucketsHistogram fromByteBuffer(ByteBuffer buffer, int numBytes)
       {
-        buffer.limit(buffer.position() + numBytes);
-        FixedBucketsHistogram fbh = FixedBucketsHistogram.fromByteBuffer(buffer);
+        Preconditions.checkArgument(
+            numBytes >= 0 && numBytes <= buffer.remaining(),
+            "numBytes[%s] exceeds buffer remaining[%s]",
+            numBytes,
+            buffer.remaining()
+        );
+        buffer.limit(Math.addExact(buffer.position(), numBytes));
+        final FixedBucketsHistogram fbh = FixedBucketsHistogram.fromByteBuffer(buffer);
         return fbh;
       }
 

extensions-core/lookups-cached-global/src/test/java/org/apache/druid/data/input/MapPopulatorTest.java

@@ -39,7 +39,7 @@ public class MapPopulatorTest
   public void test_getByteLengthOfObject_string_stringLength()
   {
     String o = "string";
-    Assert.assertEquals((o.length() * Character.BYTES) + 40, MapPopulator.getByteLengthOfObject(o));
+    Assert.assertEquals(((long) o.length() * Character.BYTES) + 40, MapPopulator.getByteLengthOfObject(o));
   }
 
   @Test

extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/S3TaskLogsTest.java

@@ -511,7 +511,8 @@ public void test_taskLog_fetch_withNegativeRange() throws IOException
 
     S3TaskLogs s3TaskLogs = getS3TaskLogs();
 
-    Optional<InputStream> inputStreamOptional = s3TaskLogs.streamTaskLog(KEY_1, -1 * (LOG_CONTENTS.length() - 1));
+    final Optional<InputStream> inputStreamOptional =
+        s3TaskLogs.streamTaskLog(KEY_1, 1L - LOG_CONTENTS.length());
     String taskLogs = new BufferedReader(
         new InputStreamReader(inputStreamOptional.get(), StandardCharsets.UTF_8))
         .lines()

extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/output/RetryableS3OutputStreamTest.java

@@ -126,7 +126,7 @@ public void testWriteAndHappy() throws IOException
     }
     // each chunk is 10 bytes, so there should be 10 chunks.
     Assert.assertEquals(10, s3.partRequests.size());
-    s3.assertCompleted(chunkSize, Integer.BYTES * 25);
+    s3.assertCompleted(chunkSize, Integer.BYTES * 25L);
   }
 
   @Test
@@ -144,7 +144,7 @@ public void testWriteSizeLargerThanConfiguredMaxChunkSizeShouldSucceed() throws
     }
     // each chunk 10 bytes, so there should be 2 chunks.
     Assert.assertEquals(2, s3.partRequests.size());
-    s3.assertCompleted(chunkSize, Integer.BYTES * 3);
+    s3.assertCompleted(chunkSize, Integer.BYTES * 3L);
   }
 
   @Test
@@ -195,7 +195,7 @@ public void testSuccessToUploadAfterRetry() throws IOException
     }
     // each chunk is 10 bytes, so there should be 10 chunks.
     Assert.assertEquals(10, s3.partRequests.size());
-    s3.assertCompleted(chunkSize, Integer.BYTES * 25);
+    s3.assertCompleted(chunkSize, Integer.BYTES * 25L);
   }
 
   @Test