From 8256f17ce110ff213a0c0a255a3838332f821b72 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:25:29 -0700 Subject: [PATCH 01/15] CountMin: check the image size in deserialize() - Check that the image holds the preamble. - For a non-empty image, check that it holds the total weight and the whole sketch array. An empty image still allocates the configured sketch array, as designed: a CountMin sketch does not grow, and is sized up front for the accuracy it needs. Adds AllocationTracker, a test helper that measures allocation by the current thread. Co-Authored-By: Claude Opus 5.5 --- .../datasketches/count/CountMinSketch.java | 8 +++- .../common/AllocationTracker.java | 46 +++++++++++++++++++ .../count/CountMinSketchTest.java | 33 +++++++++++++ 3 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 src/test/java/org/apache/datasketches/common/AllocationTracker.java diff --git a/src/main/java/org/apache/datasketches/count/CountMinSketch.java b/src/main/java/org/apache/datasketches/count/CountMinSketch.java index 9bf576e13..5354eb702 100644 --- a/src/main/java/org/apache/datasketches/count/CountMinSketch.java +++ b/src/main/java/org/apache/datasketches/count/CountMinSketch.java @@ -441,6 +441,8 @@ public byte[] toByteArray() { * @return The deserialized CountMinSketch. */ public static CountMinSketch deserialize(final byte[] b, final long seed) { + final int preambleBytes = Family.COUNTMIN.getMinPreLongs() * Long.BYTES; + Util.checkBounds(0, preambleBytes, b.length); final PositionalSegment posSeg = PositionalSegment.wrap(MemorySegment.ofArray(b)); final byte preambleLongs = posSeg.getByte(); @@ -476,8 +478,12 @@ public static CountMinSketch deserialize(final byte[] b, final long seed) { + Util.computeSeedHash(seed)); } - final CountMinSketch cms = new CountMinSketch(numHashes, numBuckets, seed); final boolean empty = (flagsByte & Flag.IS_EMPTY.mask()) > 0; + if (!empty) { //the image must hold the total weight and the whole sketch array + final long arrayBytes = (long) numHashes * numBuckets * Long.BYTES; + Util.checkBounds(preambleBytes, Long.BYTES + Math.max(arrayBytes, 0), b.length); + } + final CountMinSketch cms = new CountMinSketch(numHashes, numBuckets, seed); if (empty) { return cms; } diff --git a/src/test/java/org/apache/datasketches/common/AllocationTracker.java b/src/test/java/org/apache/datasketches/common/AllocationTracker.java new file mode 100644 index 000000000..7dd8f1b44 --- /dev/null +++ b/src/test/java/org/apache/datasketches/common/AllocationTracker.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.common; + +import java.lang.management.ManagementFactory; + +/** + * Test utility that measures how many bytes the current thread allocates while running a task. + * Used to check that deserializing a small, corrupt image does not cause a large allocation. + */ +public final class AllocationTracker { + private static final com.sun.management.ThreadMXBean TMX = + (com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean(); + + private AllocationTracker() { } + + /** + * Runs the given task and returns the number of bytes the current thread allocated while running it. + * Any exception thrown by the task is ignored. + * @param task the task to run + * @return the number of bytes allocated + */ + public static long allocatedBytes(final Runnable task) { + final long tid = Thread.currentThread().threadId(); + final long before = TMX.getThreadAllocatedBytes(tid); + try { task.run(); } catch (final RuntimeException e) { /* ignored */ } + return TMX.getThreadAllocatedBytes(tid) - before; + } +} diff --git a/src/test/java/org/apache/datasketches/count/CountMinSketchTest.java b/src/test/java/org/apache/datasketches/count/CountMinSketchTest.java index 0f37c1a71..59590d9a6 100644 --- a/src/test/java/org/apache/datasketches/count/CountMinSketchTest.java +++ b/src/test/java/org/apache/datasketches/count/CountMinSketchTest.java @@ -24,6 +24,8 @@ import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; import org.apache.datasketches.common.SketchesArgumentException; import org.apache.datasketches.common.SketchesException; import org.testng.annotations.Test; @@ -238,4 +240,35 @@ public void serializeDeserializeTest() { assertEquals(d.getEstimate(i), c.getEstimate(i)); } } + + @Test + public void emptyImageIsUsable() { + final long seed = 123L; + final CountMinSketch cms = CountMinSketch.deserialize(new CountMinSketch((byte) 3, 5, seed).toByteArray(), seed); + final CountMinSketch other = new CountMinSketch((byte) 3, 5, seed); + cms.merge(other); //both empty + assertTrue(cms.isEmpty()); + other.update(1L, 2); + cms.merge(other); + assertEquals(cms.getEstimate(1L), 2); + cms.update(1L, 3); + assertEquals(cms.getEstimate(1L), 5); + final CountMinSketch cms2 = CountMinSketch.deserialize(cms.toByteArray(), seed); + assertEquals(cms2.getEstimate(1L), 5); + } + + @Test + public void truncatedImageThrows() { + final long seed = 123L; + final CountMinSketch cms = new CountMinSketch((byte) 3, 5, seed); + cms.update(1L, 1); + final byte[] bytes = cms.toByteArray(); + ByteArrayUtil.putIntLE(bytes, 8, 1 << 28); //numBuckets larger than the image holds + final long allocated = AllocationTracker.allocatedBytes(() -> CountMinSketch.deserialize(bytes, seed)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + assertThrows(SketchesArgumentException.class, () -> CountMinSketch.deserialize(bytes, seed)); + final byte[] shortBytes = java.util.Arrays.copyOf(bytes, 8); + assertThrows(SketchesArgumentException.class, () -> CountMinSketch.deserialize(shortBytes, seed)); + } + } From c95fb7c9019b342a3406feafa4d2daf39791979f Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:28:46 -0700 Subject: [PATCH 02/15] Frequencies: check map sizes and item counts in binary images In FrequentLongsSketch.getInstance(MemorySegment) and FrequentItemsSketch.getInstance(MemorySegment, serDe): - Check that LG_MIN_MAP_SIZE <= lgCurMapSize <= lgMaxMapSize, and that activeItems is between 0 and the capacity of the current map, as C++ check_size() does. - Check, in long arithmetic, that the image holds the declared arrays. - Rebuild the sketch into a map just large enough for the active items. The map grows as needed. Co-Authored-By: Claude Opus 5.5 --- .../frequencies/FrequentItemsSketch.java | 15 +++--- .../frequencies/FrequentLongsSketch.java | 15 +++--- .../apache/datasketches/frequencies/Util.java | 36 +++++++++++++ .../frequencies/ItemsSketchTest.java | 33 ++++++++++++ .../frequencies/LongsSketchTest.java | 51 +++++++++++++++++++ 5 files changed, 138 insertions(+), 12 deletions(-) diff --git a/src/main/java/org/apache/datasketches/frequencies/FrequentItemsSketch.java b/src/main/java/org/apache/datasketches/frequencies/FrequentItemsSketch.java index 29967491f..a662a9e10 100644 --- a/src/main/java/org/apache/datasketches/frequencies/FrequentItemsSketch.java +++ b/src/main/java/org/apache/datasketches/frequencies/FrequentItemsSketch.java @@ -43,6 +43,7 @@ import static org.apache.datasketches.frequencies.PreambleUtil.insertSerVer; import static org.apache.datasketches.frequencies.Util.LG_MIN_MAP_SIZE; import static org.apache.datasketches.frequencies.Util.SAMPLE_SIZE; +import static org.apache.datasketches.frequencies.Util.checkMapSizes; import static org.apache.datasketches.frequencies.Util.checkStreamWeight; import java.lang.foreign.MemorySegment; @@ -275,17 +276,19 @@ public static FrequentItemsSketch getInstance(final MemorySegment srcSeg, MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, 0, preArr, 0, preLongs); checkStreamWeight(preArr[2]); - final FrequentItemsSketch fis = new FrequentItemsSketch<>(lgMaxMapSize, lgCurMapSize); - fis.streamWeight = 0; //update after - fis.offset = preArr[3]; - final int preBytes = preLongs << 3; final int activeItems = extractActiveItems(preArr[1]); + final int lgStartMapSize = + checkMapSizes(lgMaxMapSize, lgCurMapSize, activeItems, ReversePurgeItemHashMap.getLoadFactor()); + final long reqBytes = (preBytes + ((long)activeItems * Long.BYTES)); //count Arr only + checkBounds(0, reqBytes, srcSeg.byteSize()); //check MemorySegment capacity before allocating + + final FrequentItemsSketch fis = new FrequentItemsSketch<>(lgMaxMapSize, lgStartMapSize); + fis.streamWeight = 0; //update after + fis.offset = preArr[3]; //Get countArray final long[] countArray = new long[activeItems]; - final long reqBytes = (preBytes + ((long)activeItems * Long.BYTES)); //count Arr only - checkBounds(0, reqBytes, srcSeg.byteSize()); //check MemorySegment capacity MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, preBytes, countArray, 0, activeItems); diff --git a/src/main/java/org/apache/datasketches/frequencies/FrequentLongsSketch.java b/src/main/java/org/apache/datasketches/frequencies/FrequentLongsSketch.java index 4d5935d76..81da5c8b9 100644 --- a/src/main/java/org/apache/datasketches/frequencies/FrequentLongsSketch.java +++ b/src/main/java/org/apache/datasketches/frequencies/FrequentLongsSketch.java @@ -42,6 +42,7 @@ import static org.apache.datasketches.frequencies.PreambleUtil.insertSerVer; import static org.apache.datasketches.frequencies.Util.LG_MIN_MAP_SIZE; import static org.apache.datasketches.frequencies.Util.SAMPLE_SIZE; +import static org.apache.datasketches.frequencies.Util.checkMapSizes; import static org.apache.datasketches.frequencies.Util.checkStreamWeight; import java.lang.foreign.MemorySegment; @@ -278,17 +279,19 @@ public static FrequentLongsSketch getInstance(final MemorySegment srcSeg) { MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, 0, preArr, 0, preLongs); checkStreamWeight(preArr[2]); - final FrequentLongsSketch fls = new FrequentLongsSketch(lgMaxMapSize, lgCurMapSize); - fls.streamWeight = 0; //update after - fls.offset = preArr[3]; - final int preBytes = preLongs << 3; final int activeItems = extractActiveItems(preArr[1]); + final int lgStartMapSize = + checkMapSizes(lgMaxMapSize, lgCurMapSize, activeItems, ReversePurgeLongHashMap.getLoadFactor()); + final long reqBytes = preBytes + (2L * activeItems * Long.BYTES); //count Arr + Items Arr + checkBounds(0, reqBytes, srcSeg.byteSize()); //check MemorySegment capacity before allocating + + final FrequentLongsSketch fls = new FrequentLongsSketch(lgMaxMapSize, lgStartMapSize); + fls.streamWeight = 0; //update after + fls.offset = preArr[3]; //Get countArray final long[] countArray = new long[activeItems]; - final int reqBytes = preBytes + (2 * activeItems * Long.BYTES); //count Arr + Items Arr - checkBounds(0, reqBytes, srcSeg.byteSize()); //check MemorySegment capacity MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, preBytes, countArray, 0, activeItems); //Get itemArray diff --git a/src/main/java/org/apache/datasketches/frequencies/Util.java b/src/main/java/org/apache/datasketches/frequencies/Util.java index f17d98402..bbdadd59a 100644 --- a/src/main/java/org/apache/datasketches/frequencies/Util.java +++ b/src/main/java/org/apache/datasketches/frequencies/Util.java @@ -53,6 +53,42 @@ static long hash(long key) { return key; } + /** + * Checks the map sizes and the number of active items read from the serialized image of a + * non-empty sketch, and returns the log2 size of the map to rebuild the sketch into. + * + *

The returned size is just large enough to hold the active items, and no larger than + * lgCurMapSize. The map then grows as needed while the items are inserted, so the memory + * allocated is bounded by the size of the image, not by the declared map size.

+ * + * @param lgMaxMapSize log2 of the maximum map size + * @param lgCurMapSize log2 of the current map size + * @param activeItems the number of active items + * @param loadFactor the load factor of the hash map + * @return log2 of the map size to rebuild into + * @throws SketchesArgumentException if the values are not consistent + */ + static int checkMapSizes(final int lgMaxMapSize, final int lgCurMapSize, final int activeItems, + final double loadFactor) { + if (lgCurMapSize < LG_MIN_MAP_SIZE) { + throw new SketchesArgumentException( + "Possible Corruption: LgCurMapSize must be at least " + LG_MIN_MAP_SIZE + ": " + lgCurMapSize); + } + if (lgCurMapSize > lgMaxMapSize) { + throw new SketchesArgumentException( + "Possible Corruption: LgCurMapSize must not exceed LgMaxMapSize: " + + lgCurMapSize + " > " + lgMaxMapSize); + } + final long maxActive = (long) ((1L << Math.min(lgCurMapSize, 62)) * loadFactor); + if ((activeItems < 0) || (activeItems > maxActive)) { + throw new SketchesArgumentException( + "Possible Corruption: ActiveItems exceeds map capacity: " + activeItems); + } + int lgStart = LG_MIN_MAP_SIZE; + while ((lgStart < lgCurMapSize) && (((1L << lgStart) * loadFactor) < activeItems)) { lgStart++; } + return lgStart; + } + /** * Checks the stream weight read from the serialized image of a non-empty sketch. * @param streamWeight the stream weight diff --git a/src/test/java/org/apache/datasketches/frequencies/ItemsSketchTest.java b/src/test/java/org/apache/datasketches/frequencies/ItemsSketchTest.java index 73583a3ae..d0507f768 100644 --- a/src/test/java/org/apache/datasketches/frequencies/ItemsSketchTest.java +++ b/src/test/java/org/apache/datasketches/frequencies/ItemsSketchTest.java @@ -506,4 +506,37 @@ private static void tryBadBytes(final byte[] bytes) { static void println(final String s) { //System.out.println(s); //disable here } + + private static byte[] itemsImage() { + final FrequentItemsSketch fis = new FrequentItemsSketch<>(1 << 10); + for (long i = 0; i < 5; i++) { fis.update(i, i + 1); } + return fis.toByteArray(new ArrayOfLongsSerDe()); + } + + @Test + public void checkBinaryDeserLargeDeclaredMapDoesNotAllocate() { + final byte[] bytes = itemsImage(); + bytes[PreambleUtil.LG_MAX_MAP_SIZE_BYTE] = 30; + bytes[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = 30; + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = org.apache.datasketches.common.AllocationTracker.allocatedBytes( + () -> FrequentItemsSketch.getInstance(seg, new ArrayOfLongsSerDe())); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + final FrequentItemsSketch fis = FrequentItemsSketch.getInstance(seg, new ArrayOfLongsSerDe()); + assertEquals(fis.getNumActiveItems(), 5); + assertEquals(fis.getEstimate(4L), 5); + } + + @Test + public void checkBinaryDeserInconsistentSizes() { + final byte[] curTooBig = itemsImage(); + curTooBig[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = (byte) (curTooBig[PreambleUtil.LG_MAX_MAP_SIZE_BYTE] + 1); + final byte[] tooManyActive = itemsImage(); + org.apache.datasketches.common.ByteArrayUtil.putIntLE(tooManyActive, 8, 1 << 20); + for (final byte[] b : new byte[][] {curTooBig, tooManyActive}) { + try { FrequentItemsSketch.getInstance(MemorySegment.ofArray(b), new ArrayOfLongsSerDe()); fail(); } + catch (final SketchesArgumentException e) { } + } + } + } diff --git a/src/test/java/org/apache/datasketches/frequencies/LongsSketchTest.java b/src/test/java/org/apache/datasketches/frequencies/LongsSketchTest.java index b6527b178..6326f499e 100644 --- a/src/test/java/org/apache/datasketches/frequencies/LongsSketchTest.java +++ b/src/test/java/org/apache/datasketches/frequencies/LongsSketchTest.java @@ -796,4 +796,55 @@ private static FrequentLongsSketch newFrequencySketch(final double eps) { return new FrequentLongsSketch(maxMapSize); } + + private static byte[] longsImage() { + final FrequentLongsSketch fls = new FrequentLongsSketch(1 << 10); + for (int i = 0; i < 5; i++) { fls.update(i, i + 1); } + return fls.toByteArray(); + } + + @Test + public void checkBinaryDeserLargeDeclaredMapDoesNotAllocate() { + final byte[] bytes = longsImage(); + bytes[PreambleUtil.LG_MAX_MAP_SIZE_BYTE] = 30; + bytes[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = 30; + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = org.apache.datasketches.common.AllocationTracker.allocatedBytes( + () -> FrequentLongsSketch.getInstance(seg)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + final FrequentLongsSketch fls = FrequentLongsSketch.getInstance(seg); + assertEquals(fls.getNumActiveItems(), 5); + assertEquals(fls.getEstimate(4), 5); + } + + @Test + public void checkBinaryDeserInconsistentSizes() { + final byte[] curTooBig = longsImage(); + curTooBig[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = (byte) (curTooBig[PreambleUtil.LG_MAX_MAP_SIZE_BYTE] + 1); + final byte[] curTooSmall = longsImage(); + curTooSmall[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = 2; + final byte[] tooManyActive = longsImage(); + org.apache.datasketches.common.ByteArrayUtil.putIntLE(tooManyActive, 8, 1 << 20); + final byte[] negActive = longsImage(); + org.apache.datasketches.common.ByteArrayUtil.putIntLE(negActive, 8, -1); + for (final byte[] b : new byte[][] {curTooBig, curTooSmall, tooManyActive, negActive}) { + try { FrequentLongsSketch.getInstance(MemorySegment.ofArray(b)); fail(); } + catch (final SketchesArgumentException e) { } + } + } + + @Test + public void checkBinaryDeserActiveItemsBeyondImage() { + final byte[] bytes = longsImage(); + bytes[PreambleUtil.LG_MAX_MAP_SIZE_BYTE] = 30; + bytes[PreambleUtil.LG_CUR_MAP_SIZE_BYTE] = 30; + org.apache.datasketches.common.ByteArrayUtil.putIntLE(bytes, 8, 1 << 28); //more items than the image holds + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = org.apache.datasketches.common.AllocationTracker.allocatedBytes( + () -> FrequentLongsSketch.getInstance(seg)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + try { FrequentLongsSketch.getInstance(seg); fail(); } + catch (final SketchesArgumentException e) { } + } + } From 8c2e28864f8a581ee381e007fdc808f5330bcae7 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:40:31 -0700 Subject: [PATCH 03/15] Tuple: validate updatable and compact images Updatable (QuickSelectSketch) image: - Check 0 <= lgNomEntries <= 26 (lgK has the same maximum as Theta), lgCurrentCapacity <= max(lgNomEntries + 1, 5) (the table grows to at most twice the nominal entries, but never starts below 32 slots), and 0 <= lgResizeFactor <= 3. - Check that the count is non-negative, smaller than the table, and that the image holds at least 8 bytes per entry. - Check that the table is plausible for the count: either the starting size or at most 32 times the count. This keeps the exact table size of valid images. Compact image (current and legacy forms): - Check that the image holds the preamble, and that the count is non-negative and fits in the image at 8 bytes per entry. Constructors: the QuickSelectSketch constructors now reject nominal entries that are not positive or are above 2^26, as the builder and Theta already do. Before, UpdatableTupleSketch, TupleUnion and the Double, Integer and ArrayOfStrings sketches accepted them. Co-Authored-By: Claude Opus 5.5 --- .../tuple/CompactTupleSketch.java | 18 +++ .../datasketches/tuple/QuickSelectSketch.java | 47 +++++- .../tuple/TupleCorruptImageTest.java | 146 ++++++++++++++++++ 3 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 src/test/java/org/apache/datasketches/tuple/TupleCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/tuple/CompactTupleSketch.java b/src/main/java/org/apache/datasketches/tuple/CompactTupleSketch.java index ab5f06de5..f02da0a21 100644 --- a/src/main/java/org/apache/datasketches/tuple/CompactTupleSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/CompactTupleSketch.java @@ -23,6 +23,7 @@ import static java.lang.foreign.ValueLayout.JAVA_INT_UNALIGNED; import static java.lang.foreign.ValueLayout.JAVA_LONG_UNALIGNED; import static java.lang.foreign.ValueLayout.JAVA_SHORT_UNALIGNED; +import static org.apache.datasketches.common.Util.checkBounds; import static org.apache.datasketches.common.Util.checkSeedHashes; import static org.apache.datasketches.common.Util.computeSeedHash; import static org.apache.datasketches.tuple.PreambleUtil.COMPACT_FLAGS_BYTE; @@ -118,11 +119,13 @@ private enum Flags { IS_RESERVED, IS_READ_ONLY, IS_EMPTY, IS_COMPACT, IS_ORDERED if (version == COMPACT_SER_VER_WITH_CLASS_NAME) { classNameLength = seg.get(JAVA_BYTE, offset++); } + checkBounds(offset, Integer.BYTES, seg.byteSize()); final int count = seg.get(JAVA_INT_UNALIGNED, offset); offset += Integer.BYTES; if (version == COMPACT_SER_VER_WITH_CLASS_NAME) { offset += classNameLength; } + checkCount(count, offset, seg); hashArr_ = new long[count]; for (int i = 0; i < count; i++) { @@ -137,6 +140,7 @@ private enum Flags { IS_RESERVED, IS_READ_ONLY, IS_EMPTY, IS_COMPACT, IS_ORDERED summaryArr_ = null; } } else { // current serial format + checkBounds(0, (long) preambleLongs * Long.BYTES, seg.byteSize()); final byte flags = seg.get(JAVA_BYTE, COMPACT_FLAGS_BYTE); final short seedHash = seg.get(JAVA_SHORT_UNALIGNED, SEED_HASH_SHORT); empty_ = (flags & (1 << Flags.IS_EMPTY.ordinal())) > 0; @@ -154,6 +158,7 @@ private enum Flags { IS_RESERVED, IS_READ_ONLY, IS_EMPTY, IS_COMPACT, IS_ORDERED } } int offset = preambleLongs * Long.BYTES; // the entries start after the preamble + checkCount(count, offset, seg); hashArr_ = new long[count]; for (int i = 0; i < count; i++) { @@ -164,6 +169,19 @@ private enum Flags { IS_RESERVED, IS_READ_ONLY, IS_EMPTY, IS_COMPACT, IS_ORDERED } } + /** + * Checks the retained entries count read from an image. Each entry is at least an 8-byte hash, + * so the remaining bytes of the image bound the count. + * @param count the retained entries count + * @param offset the offset of the first entry + * @param seg the image + */ + private static void checkCount(final int count, final int offset, final MemorySegment seg) { + if ((count < 0) || (((long) count * Long.BYTES) > (seg.byteSize() - offset))) { + throw new SketchesArgumentException("Possible corruption: invalid retained entries count: " + count); + } + } + @SuppressWarnings({"unchecked"}) private int readSummary(final MemorySegment seg, final int offset, final int i, final int count, final SummaryDeserializer deserializer) { diff --git a/src/main/java/org/apache/datasketches/tuple/QuickSelectSketch.java b/src/main/java/org/apache/datasketches/tuple/QuickSelectSketch.java index 5f3777f33..ce62e817e 100644 --- a/src/main/java/org/apache/datasketches/tuple/QuickSelectSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/QuickSelectSketch.java @@ -140,6 +140,10 @@ private enum Flags { IS_RESERVED, IS_IN_SAMPLING_MODE, IS_EMPTY, HAS_ENTRIES, IS final float samplingProbability, final SummaryFactory summaryFactory, final short seedHash) { + if ((nomEntries <= 0) || (nomEntries > (1 << ThetaUtil.MAX_LG_NOM_LONGS))) { + throw new SketchesArgumentException("Nominal Entries must be > 0 and <= " + + (1 << ThetaUtil.MAX_LG_NOM_LONGS) + ": " + nomEntries); + } this( nomEntries, lgResizeFactor, @@ -264,9 +268,17 @@ long validate( SerializerDeserializer.validateType(seg.get(JAVA_BYTE, SKETCH_TYPE_BYTE), SerializerDeserializer.SketchType.QuickSelectSketch); final byte flags = seg.get(JAVA_BYTE, LEGACY_FLAGS_BYTE); - myNomEntries = 1 << seg.get(JAVA_BYTE, LG_NOM_ENTRIES_BYTE); - myLgCurrentCapacity = seg.get(JAVA_BYTE, LG_CUR_CAPACITY_BYTE); + final int lgNomEntries = seg.get(JAVA_BYTE, LG_NOM_ENTRIES_BYTE); + final int lgCurCapacity = seg.get(JAVA_BYTE, LG_CUR_CAPACITY_BYTE); myLgResizeFactor = seg.get(JAVA_BYTE, LG_RESIZE_FACTOR_BYTE); + // lgK is at most 26, as for Theta. The hash table grows to at most twice the nominal entries, + // but is never smaller than the minimum table size. + if ((lgNomEntries < 0) || (lgNomEntries > ThetaUtil.MAX_LG_NOM_LONGS) || (lgCurCapacity < 0) + || (lgCurCapacity > Math.max(lgNomEntries + 1, ThetaUtil.MIN_LG_ARR_LONGS))) { + throw new SketchesArgumentException("Possible corruption: lgNomEntries: " + lgNomEntries + + ", lgCurrentCapacity: " + lgCurCapacity); + } + myNomEntries = 1 << lgNomEntries; int offset = QS_VARIABLE_START; // the remaining fields are packed and optional checkBounds(0, preambleLongs * 8L, seg.byteSize()); @@ -287,9 +299,17 @@ long validate( int count = 0; final boolean hasEntries = (flags & (1 << Flags.HAS_ENTRIES.ordinal())) > 0; if (hasEntries) { + checkBounds(offset, Integer.BYTES, seg.byteSize()); count = seg.get(JAVA_INT_UNALIGNED, offset); offset += Integer.BYTES; } + // each entry is at least an 8-byte hash, so the image bounds the count + if ((count < 0) || (count >= (1L << lgCurCapacity)) + || (((long) count * Long.BYTES) > (seg.byteSize() - offset))) { + throw new SketchesArgumentException("Possible corruption: invalid retained entries count: " + count); + } + checkLgCurrentCapacity(lgCurCapacity, count, lgNomEntries, myLgResizeFactor); + myLgCurrentCapacity = lgCurCapacity; final int currentCapacity = 1 << myLgCurrentCapacity; myHashTable = new long[currentCapacity]; for (int i = 0; i < count; i++) { @@ -315,6 +335,29 @@ long validate( } //end class Validate + /** + * Checks that the hash table size declared by an image is plausible for the number of entries it holds. + * A table only grows when its entries reach a threshold, so its size is either the starting size or + * at most about 16 times the number of entries, just after a resize with the largest resize factor. + * This bounds the allocation by the size of the image. + * @param lgCurCapacity log2 of the current table size declared by the image + * @param count the number of entries in the image + * @param lgNomEntries log2 of the nominal entries + * @param lgResizeFactor log2 of the resize factor + */ + private static void checkLgCurrentCapacity(final int lgCurCapacity, final int count, final int lgNomEntries, + final int lgResizeFactor) { + if ((lgResizeFactor < 0) || (lgResizeFactor > ResizeFactor.X8.lg())) { + throw new SketchesArgumentException("Possible corruption: lgResizeFactor: " + lgResizeFactor); + } + final int lgStart = ThetaUtil.startingSubMultiple(lgNomEntries + 1, lgResizeFactor, ThetaUtil.MIN_LG_ARR_LONGS); + final int lgCount = (count <= 1) ? 0 : 32 - Integer.numberOfLeadingZeros(count - 1); //ceiling log2 + if (lgCurCapacity > Math.max(lgStart, lgCount + 5)) { + throw new SketchesArgumentException("Possible corruption: lgCurrentCapacity " + lgCurCapacity + + " is too large for " + count + " retained entries"); + } + } + /** * @return a deep copy of this sketch */ diff --git a/src/test/java/org/apache/datasketches/tuple/TupleCorruptImageTest.java b/src/test/java/org/apache/datasketches/tuple/TupleCorruptImageTest.java new file mode 100644 index 000000000..853a001f5 --- /dev/null +++ b/src/test/java/org/apache/datasketches/tuple/TupleCorruptImageTest.java @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.tuple; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.apache.datasketches.tuple.adouble.DoubleSummary; +import org.apache.datasketches.tuple.adouble.DoubleSummaryDeserializer; +import org.apache.datasketches.tuple.adouble.DoubleSummaryFactory; +import org.apache.datasketches.tuple.adouble.DoubleTupleSketch; +import org.testng.annotations.Test; + +/** + * Checks that corrupt tuple sketch images are rejected without large allocations. + */ +@SuppressWarnings("deprecation") +public class TupleCorruptImageTest { + private static final int LG_NOM_BYTE = 5; + private static final int LG_CUR_BYTE = 6; + private static final int COMPACT_COUNT_INT = 8; + private final DoubleSummaryDeserializer deser = new DoubleSummaryDeserializer(); + + private static UpdatableTupleSketch sketch(final int n) { + final UpdatableTupleSketch sk = + new UpdatableTupleSketchBuilder<>(new DoubleSummaryFactory(DoubleSummary.Mode.Sum)).build(); + for (int i = 0; i < n; i++) { sk.update(i, 1.0); } + return sk; + } + + private void checkRejected(final byte[] bytes) { + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = AllocationTracker.allocatedBytes(() -> TupleSketch.heapifySketch(seg, deser)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> TupleSketch.heapifySketch(seg, deser)); + } + + @Test + public void checkValidImagesStillWork() { + for (final int n : new int[] {0, 1, 10, 100, 10_000}) { + final UpdatableTupleSketch sk = sketch(n); + assertEquals(TupleSketch.heapifySketch(MemorySegment.ofArray(sk.toByteArray()), deser).getRetainedEntries(), + sk.getRetainedEntries()); + assertEquals(TupleSketch.heapifySketch(MemorySegment.ofArray(sk.compact().toByteArray()), deser) + .getRetainedEntries(), sk.getRetainedEntries()); + } + } + + @Test + public void checkUpdatableLargeTable() { + final byte[] bytes = sketch(10).toByteArray(); + bytes[LG_NOM_BYTE] = 25; + bytes[LG_CUR_BYTE] = 26; //far too large for 10 entries + checkRejected(bytes); + } + + @Test + public void checkSmallNominalEntriesStillWork() { + //for nominal entries below 16 the table starts at the minimum size, which is more than twice nominal + for (int lgK = 0; lgK <= 5; lgK++) { + for (int lgRF = 0; lgRF <= 3; lgRF++) { + for (final int n : new int[] {0, 1, 3, 100}) { + final DoubleTupleSketch sk = new DoubleTupleSketch(lgK, lgRF, 1.0F, DoubleSummary.Mode.Sum); + for (int i = 0; i < n; i++) { sk.update(i, 1.0); } + final byte[] bytes = sk.toByteArray(); + assertEquals(TupleSketch.heapifySketch(MemorySegment.ofArray(bytes), deser).getRetainedEntries(), + sk.getRetainedEntries()); + assertEquals(new DoubleTupleSketch(MemorySegment.ofArray(bytes), DoubleSummary.Mode.Sum) + .getRetainedEntries(), sk.getRetainedEntries()); + } + } + } + } + + @Test + public void checkUpdatableTableLargerThanNominal() { + final byte[] bytes = sketch(10).toByteArray(); + bytes[LG_CUR_BYTE] = (byte) (bytes[LG_NOM_BYTE] + 2); + checkRejected(bytes); + final byte[] bytes2 = sketch(10).toByteArray(); + bytes2[LG_NOM_BYTE] = 30; + checkRejected(bytes2); + final byte[] bytes3 = sketch(10).toByteArray(); + bytes3[LG_NOM_BYTE] = 27; //lgK above the Theta maximum of 26 + checkRejected(bytes3); + } + + @Test + public void checkConstructorNominalEntriesLimit() { + final DoubleSummaryFactory factory = new DoubleSummaryFactory(DoubleSummary.Mode.Sum); + assertEquals(new UpdatableTupleSketch(1 << 26, 3, 1f, factory).getNominalEntries(), 1 << 26); + expectThrows(SketchesArgumentException.class, + () -> new UpdatableTupleSketch((1 << 26) + 1, 3, 1f, factory)); + expectThrows(SketchesArgumentException.class, + () -> new UpdatableTupleSketch(0, 3, 1f, factory)); + expectThrows(SketchesArgumentException.class, () -> new TupleUnion(1 << 27, null)); + expectThrows(SketchesArgumentException.class, + () -> new DoubleTupleSketch(27, 3, 1f, DoubleSummary.Mode.Sum)); + } + + @Test + public void checkUpdatableCountBeyondImage() { + final byte[] bytes = sketch(10).toByteArray(); + ByteArrayUtil.putIntLE(bytes, 8, 1 << 28); //count: no sampling and no theta, so it is at byte 8 + checkRejected(bytes); + } + + @Test + public void checkCompactCountBeyondImage() { + final byte[] bytes = sketch(10).compact().toByteArray(); + ByteArrayUtil.putIntLE(bytes, COMPACT_COUNT_INT, 1 << 28); + checkRejected(bytes); + final byte[] negative = sketch(10).compact().toByteArray(); + ByteArrayUtil.putIntLE(negative, COMPACT_COUNT_INT, -1); + checkRejected(negative); + } + + @Test + public void checkCompactTruncatedPreamble() { + final byte[] bytes = java.util.Arrays.copyOf(sketch(10).compact().toByteArray(), 12); + checkRejected(bytes); + } +} From 55d9b770e1c0914ef218ac83b072961c29284bf6 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:43:52 -0700 Subject: [PATCH 04/15] AoD: validate sketch and union images - Compact images (heapify and wrap): check that numValues is positive, the count is non-negative, and the image holds count keys and count * numValues values. - QuickSelect images (heapify and wrap, and the union, which contains one): check that numValues is positive, 0 <= lgNomEntries <= 26, lgCurrentCapacity <= max(lgNomEntries + 1, 5) (the table never starts below 32 slots; unions accept nominal entries below 16), 0 <= count <= capacity, and that the image holds the whole table, which the serialized form always includes. Co-Authored-By: Claude Opus 5.5 --- .../ArrayOfDoublesCompactSketch.java | 31 ++++ .../ArrayOfDoublesQuickSelectSketch.java | 29 ++++ .../DirectArrayOfDoublesCompactSketch.java | 1 + ...DirectArrayOfDoublesQuickSelectSketch.java | 1 + .../HeapArrayOfDoublesCompactSketch.java | 1 + .../HeapArrayOfDoublesQuickSelectSketch.java | 1 + .../arrayofdoubles/AodCorruptImageTest.java | 154 ++++++++++++++++++ 7 files changed, 218 insertions(+) create mode 100644 src/test/java/org/apache/datasketches/tuple/arrayofdoubles/AodCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesCompactSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesCompactSketch.java index 1cf3ac040..fb217a7c2 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesCompactSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesCompactSketch.java @@ -19,6 +19,14 @@ package org.apache.datasketches.tuple.arrayofdoubles; +import static java.lang.foreign.ValueLayout.JAVA_BYTE; +import static java.lang.foreign.ValueLayout.JAVA_INT_UNALIGNED; +import static org.apache.datasketches.common.Util.checkBounds; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.SketchesArgumentException; + /** * Top level compact tuple sketch of type ArrayOfDoubles. Compact sketches are never created * directly. They are created as a result of the compact() method on a QuickSelectSketch @@ -46,6 +54,29 @@ public abstract class ArrayOfDoublesCompactSketch extends ArrayOfDoublesSketch { super(numValues); } + /** + * Checks that the given image is large enough for the number of entries and values it declares, + * before anything is allocated or read from it. + * @param seg the serialized compact sketch + * @throws SketchesArgumentException if the image is too small or its counts are invalid + */ + static void checkCompactImage(final MemorySegment seg) { + checkBounds(0, EMPTY_SIZE, seg.byteSize()); + final int numValues = seg.get(JAVA_BYTE, NUM_VALUES_BYTE); + if (numValues < 1) { + throw new SketchesArgumentException("Possible corruption: numValues must be positive: " + numValues); + } + final boolean hasEntries = (seg.get(JAVA_BYTE, FLAGS_BYTE) & (1 << Flags.HAS_ENTRIES.ordinal())) != 0; + if (!hasEntries) { return; } + checkBounds(0, ENTRIES_START, seg.byteSize()); + final int count = seg.get(JAVA_INT_UNALIGNED, RETAINED_ENTRIES_INT); + if (count < 0) { + throw new SketchesArgumentException("Possible corruption: negative retained entries: " + count); + } + checkBounds(ENTRIES_START, (long) count * (SIZE_OF_KEY_BYTES + ((long) SIZE_OF_VALUE_BYTES * numValues)), + seg.byteSize()); + } + @Override public int getCurrentBytes() { final int count = getRetainedEntries(); diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesQuickSelectSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesQuickSelectSketch.java index ddf102451..61febdbcb 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesQuickSelectSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/ArrayOfDoublesQuickSelectSketch.java @@ -19,7 +19,10 @@ package org.apache.datasketches.tuple.arrayofdoubles; +import static java.lang.foreign.ValueLayout.JAVA_BYTE; +import static java.lang.foreign.ValueLayout.JAVA_INT_UNALIGNED; import static org.apache.datasketches.common.Util.ceilingPowerOf2; +import static org.apache.datasketches.common.Util.checkBounds; import java.lang.foreign.MemorySegment; @@ -54,6 +57,32 @@ abstract class ArrayOfDoublesQuickSelectSketch extends ArrayOfDoublesUpdatableSk static final int DEFAULT_LG_RESIZE_FACTOR = 3; + /** + * Checks that the given image holds the whole hash table it declares, and that its counts are valid, + * before anything is allocated or read from it. The serialized form always includes the whole table. + * @param seg the serialized QuickSelect sketch + * @throws SketchesArgumentException if the image is too small or its counts are invalid + */ + static void checkQuickSelectImage(final MemorySegment seg) { + checkBounds(0, ENTRIES_START, seg.byteSize()); + final int numValues = seg.get(JAVA_BYTE, NUM_VALUES_BYTE); + final int lgCurCapacity = seg.get(JAVA_BYTE, LG_CUR_CAPACITY_BYTE); + final int lgNomEntries = seg.get(JAVA_BYTE, LG_NOM_ENTRIES_BYTE); + //the table grows to at most twice the nominal entries, but is never smaller than the minimum table size + if ((numValues < 1) || (lgNomEntries < 0) || (lgNomEntries > ThetaUtil.MAX_LG_NOM_LONGS) || (lgCurCapacity < 0) + || (lgCurCapacity > Math.max(lgNomEntries + 1, ThetaUtil.MIN_LG_ARR_LONGS))) { + throw new SketchesArgumentException("Possible corruption: numValues: " + numValues + + ", lgNomEntries: " + lgNomEntries + ", lgCurrentCapacity: " + lgCurCapacity); + } + final long capacity = 1L << lgCurCapacity; + final int count = seg.get(JAVA_INT_UNALIGNED, RETAINED_ENTRIES_INT); + if ((count < 0) || (count > capacity)) { + throw new SketchesArgumentException("Possible corruption: invalid retained entries: " + count); + } + checkBounds(ENTRIES_START, capacity * (SIZE_OF_KEY_BYTES + ((long) SIZE_OF_VALUE_BYTES * numValues)), + seg.byteSize()); + } + // these can be derived from other things, but are kept here for performance int rebuildThreshold_; //absolute value relative to current capacity int lgCurrentCapacity_; diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesCompactSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesCompactSketch.java index e110a24fa..6342ba717 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesCompactSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesCompactSketch.java @@ -157,6 +157,7 @@ final class DirectArrayOfDoublesCompactSketch extends ArrayOfDoublesCompactSketc * @param seed See seed */ DirectArrayOfDoublesCompactSketch(final MemorySegment seg, final long seed) { + checkCompactImage(seg); super(seg.get(JAVA_BYTE, NUM_VALUES_BYTE)); seg_ = seg; SerializerDeserializer.validateFamily(seg.get(JAVA_BYTE, FAMILY_ID_BYTE), diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesQuickSelectSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesQuickSelectSketch.java index f32571d0b..1a7e6585f 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesQuickSelectSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/DirectArrayOfDoublesQuickSelectSketch.java @@ -126,6 +126,7 @@ private static final void checkMemorySegment( final MemorySegment seg, final long seed) { checkSerVer(seg); + checkQuickSelectImage(seg); super(seg.get(JAVA_BYTE, NUM_VALUES_BYTE), seed); seg_ = seg; SerializerDeserializer.validateFamily(seg.get(JAVA_BYTE, FAMILY_ID_BYTE), diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesCompactSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesCompactSketch.java index 762324543..c52b34cb7 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesCompactSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesCompactSketch.java @@ -115,6 +115,7 @@ final class HeapArrayOfDoublesCompactSketch extends ArrayOfDoublesCompactSketch * @param seed See seed */ HeapArrayOfDoublesCompactSketch(final MemorySegment seg, final long seed) { + checkCompactImage(seg); super(seg.get(JAVA_BYTE, NUM_VALUES_BYTE)); seedHash_ = seg.get(JAVA_SHORT_UNALIGNED, SEED_HASH_SHORT); SerializerDeserializer.validateFamily(seg.get(JAVA_BYTE, FAMILY_ID_BYTE), diff --git a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesQuickSelectSketch.java b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesQuickSelectSketch.java index 370d6cb44..4108201b2 100644 --- a/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesQuickSelectSketch.java +++ b/src/main/java/org/apache/datasketches/tuple/arrayofdoubles/HeapArrayOfDoublesQuickSelectSketch.java @@ -89,6 +89,7 @@ final class HeapArrayOfDoublesQuickSelectSketch extends ArrayOfDoublesQuickSelec * @param seed See seed */ HeapArrayOfDoublesQuickSelectSketch(final MemorySegment seg, final long seed) { + checkQuickSelectImage(seg); super(seg.get(JAVA_BYTE, NUM_VALUES_BYTE), seed); SerializerDeserializer.validateFamily(seg.get(JAVA_BYTE, FAMILY_ID_BYTE), seg.get(JAVA_BYTE, PREAMBLE_LONGS_BYTE)); diff --git a/src/test/java/org/apache/datasketches/tuple/arrayofdoubles/AodCorruptImageTest.java b/src/test/java/org/apache/datasketches/tuple/arrayofdoubles/AodCorruptImageTest.java new file mode 100644 index 000000000..9dfa9b621 --- /dev/null +++ b/src/test/java/org/apache/datasketches/tuple/arrayofdoubles/AodCorruptImageTest.java @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.tuple.arrayofdoubles; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; +import java.util.function.Function; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt ArrayOfDoubles sketch and union images are rejected without large allocations, + * including when the result of a wrap is used later. + */ +public class AodCorruptImageTest { + private static final int NUM_VALUES_BYTE = 5; + private static final int COMPACT_COUNT_INT = 16; + private static final int QS_LG_NOM_BYTE = 16; + private static final int QS_LG_CUR_BYTE = 17; + private static final int QS_COUNT_INT = 24; + private static final int UNION_PREAMBLE = 16; + + private static ArrayOfDoublesUpdatableSketch sketch(final int n) { + final ArrayOfDoublesUpdatableSketch sk = new ArrayOfDoublesUpdatableSketchBuilder().setNumberOfValues(2).build(); + for (int i = 0; i < n; i++) { sk.update(i, new double[] {1.0, 2.0}); } + return sk; + } + + private static byte[] unionImage() { + final ArrayOfDoublesUnion union = new ArrayOfDoublesSetOperationBuilder().setNumberOfValues(2).buildUnion(); + union.union(sketch(10)); + return union.toByteArray(); + } + + private static void checkRejected(final byte[] bytes, final Function deser) { + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = AllocationTracker.allocatedBytes(() -> deser.apply(seg)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> deser.apply(seg)); + } + + @Test + public void checkValidImagesStillWork() { + final ArrayOfDoublesUpdatableSketch sk = sketch(100); + final byte[] c = sk.compact().toByteArray(); + assertEquals(ArrayOfDoublesSketch.heapifySketch(MemorySegment.ofArray(c)).getRetainedEntries(), 100); + assertEquals(ArrayOfDoublesSketch.wrapSketch(MemorySegment.ofArray(c)).getRetainedEntries(), 100); + final byte[] u = sk.toByteArray(); + assertEquals(ArrayOfDoublesSketch.heapifyUpdatableSketch(MemorySegment.ofArray(u)).getRetainedEntries(), 100); + assertEquals(ArrayOfDoublesSketch.wrapUpdatableSketch(MemorySegment.ofArray(u)).getRetainedEntries(), 100); + assertEquals(ArrayOfDoublesUnion.heapify(MemorySegment.ofArray(unionImage())).getResult().getRetainedEntries(), 10); + assertEquals(ArrayOfDoublesUnion.wrap(MemorySegment.ofArray(unionImage())).getResult().getRetainedEntries(), 10); + final byte[] e = sketch(0).compact().toByteArray(); + assertTrue(ArrayOfDoublesSketch.heapifySketch(MemorySegment.ofArray(e)).isEmpty()); + } + + @Test + public void checkSmallNominalEntriesStillWork() { + //for nominal entries below 16 the table starts at the minimum size, which is more than twice nominal. + //The set operation builder, unlike the sketch builder, accepts nominal entries below 16. + for (final int k : new int[] {1, 2, 4, 8, 16, 32}) { + for (final int n : new int[] {0, 1, 3, 100}) { + final ArrayOfDoublesSetOperationBuilder bldr = + new ArrayOfDoublesSetOperationBuilder().setNominalEntries(k).setNumberOfValues(2); + final ArrayOfDoublesUnion heapUnion = bldr.buildUnion(); + final ArrayOfDoublesUnion directUnion = bldr.buildUnion(MemorySegment.ofArray(new byte[1 << 16])); + heapUnion.union(sketch(n)); + directUnion.union(sketch(n)); + final int expected = heapUnion.getResult().getRetainedEntries(); + for (final byte[] u : new byte[][] {heapUnion.toByteArray(), directUnion.toByteArray()}) { + assertEquals(ArrayOfDoublesUnion.heapify(MemorySegment.ofArray(u)).getResult().getRetainedEntries(), expected); + assertEquals(ArrayOfDoublesUnion.wrap(MemorySegment.ofArray(u)).getResult().getRetainedEntries(), expected); + } + } + } + } + + @Test + public void checkCompactCountBeyondImage() { + final byte[] bytes = sketch(10).compact().toByteArray(); + ByteArrayUtil.putIntLE(bytes, COMPACT_COUNT_INT, 1 << 26); + checkRejected(bytes, ArrayOfDoublesSketch::heapifySketch); + checkRejected(bytes, seg -> ArrayOfDoublesSketch.wrapSketch(seg).getValues()); + final byte[] negative = sketch(10).compact().toByteArray(); + ByteArrayUtil.putIntLE(negative, COMPACT_COUNT_INT, -1); + checkRejected(negative, ArrayOfDoublesSketch::heapifySketch); + } + + @Test + public void checkCompactNumValues() { + final byte[] bytes = sketch(10).compact().toByteArray(); + bytes[NUM_VALUES_BYTE] = (byte) 127; //more values than the image holds + checkRejected(bytes, ArrayOfDoublesSketch::heapifySketch); + bytes[NUM_VALUES_BYTE] = (byte) -1; + checkRejected(bytes, ArrayOfDoublesSketch::heapifySketch); + checkRejected(bytes, ArrayOfDoublesSketch::wrapSketch); + } + + @Test + public void checkUpdatableLargeTable() { + final byte[] bytes = sketch(10).toByteArray(); + bytes[QS_LG_NOM_BYTE] = 26; + bytes[QS_LG_CUR_BYTE] = 27; //the image does not hold this table + checkRejected(bytes, ArrayOfDoublesSketch::heapifyUpdatableSketch); + checkRejected(bytes, seg -> ArrayOfDoublesSketch.wrapUpdatableSketch(seg).getValues()); + final byte[] bytes2 = sketch(10).toByteArray(); + bytes2[QS_LG_CUR_BYTE] = (byte) (bytes2[QS_LG_NOM_BYTE] + 2); + checkRejected(bytes2, ArrayOfDoublesSketch::heapifyUpdatableSketch); + } + + @Test + public void checkUpdatableCount() { + final byte[] bytes = sketch(10).toByteArray(); + ByteArrayUtil.putIntLE(bytes, QS_COUNT_INT, 1 << 26); //more than the table holds + checkRejected(bytes, ArrayOfDoublesSketch::heapifyUpdatableSketch); + checkRejected(bytes, seg -> ArrayOfDoublesSketch.wrapUpdatableSketch(seg).getValues()); + } + + @Test + public void checkUnion() { + final byte[] bytes = unionImage(); + bytes[UNION_PREAMBLE + QS_LG_NOM_BYTE] = 26; + bytes[UNION_PREAMBLE + QS_LG_CUR_BYTE] = 27; + checkRejected(bytes, ArrayOfDoublesUnion::heapify); + checkRejected(bytes, seg -> ArrayOfDoublesUnion.wrap(seg).getResult()); + final byte[] bytes2 = unionImage(); + ByteArrayUtil.putIntLE(bytes2, UNION_PREAMBLE + QS_COUNT_INT, 1 << 26); + checkRejected(bytes2, ArrayOfDoublesUnion::heapify); + checkRejected(bytes2, seg -> ArrayOfDoublesUnion.wrap(seg).getResult()); + } +} From 4dc80febf09d093982f38c90f4fd18cfd6b1707a Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:48:30 -0700 Subject: [PATCH 05/15] Theta: validate sketch, union and intersection images - checkSegIntegrity(), used for updatable, alpha and union images (heapify and wrap): check lgNomLongs <= 26, 0 <= lgArrLongs <= lgNomLongs + 1 and 0 <= curCount <= table size, and compute the required size in long arithmetic. - Compact images (heapify, and wrap of a MemorySegment or byte[]): check that curCount is non-negative and that the image holds the hashes. Images in hash-table form must hold the whole table. - Compressed images (serial version 4): check preLongs, entryBits and numEntriesBytes, and that the image holds numEntries * entryBits bits. - Intersection images (heapify and wrap): check 0 <= lgArrLongs <= 27, -1 <= curCount <= table size (-1 means no result yet), and that the image holds the table when curCount > 0. Co-Authored-By: Claude Opus 5.5 --- .../datasketches/theta/CompactOperations.java | 7 + .../theta/CompactThetaSketch.java | 1 + .../theta/DirectCompactCompressedSketch.java | 1 + .../theta/DirectCompactSketch.java | 1 + .../datasketches/theta/PreambleUtil.java | 44 ++++++ .../theta/ThetaIntersectionImpl.java | 15 ++ .../theta/UpdatableThetaSketch.java | 16 +- .../theta/WrappedCompactCompressedSketch.java | 3 + .../theta/WrappedCompactSketch.java | 1 + .../theta/ThetaCorruptImageTest.java | 137 ++++++++++++++++++ 10 files changed, 223 insertions(+), 3 deletions(-) create mode 100644 src/test/java/org/apache/datasketches/theta/ThetaCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/theta/CompactOperations.java b/src/main/java/org/apache/datasketches/theta/CompactOperations.java index 99d406600..a09d80cac 100644 --- a/src/main/java/org/apache/datasketches/theta/CompactOperations.java +++ b/src/main/java/org/apache/datasketches/theta/CompactOperations.java @@ -52,6 +52,7 @@ import org.apache.datasketches.common.Family; import org.apache.datasketches.common.SketchesArgumentException; import org.apache.datasketches.common.SketchesStateException; +import org.apache.datasketches.thetacommon.ThetaUtil; /** * @author Lee Rhodes @@ -174,9 +175,15 @@ static CompactThetaSketch segmentToCompact( //extract hashArr > 1 final long[] hashArr; if (srcCompactFlag) { + PreambleUtil.checkCompactImage(srcSeg); hashArr = new long[curCount]; MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, srcPreLongs << 3, hashArr, 0, curCount); } else { //update sketch, thus hashTable form + if ((srcLgArrLongs > (ThetaUtil.MAX_LG_NOM_LONGS + 1)) + || (srcSeg.byteSize() < (((long) srcPreLongs << 3) + (8L << srcLgArrLongs)))) { + throw new SketchesArgumentException("Possible corruption: lgArrLongs " + srcLgArrLongs + + " is too large for the MemorySegment size " + srcSeg.byteSize()); + } final int srcCacheLen = 1 << srcLgArrLongs; final long[] tempHashArr = new long[srcCacheLen]; MemorySegment.copy(srcSeg, JAVA_LONG_UNALIGNED, srcPreLongs << 3, tempHashArr, 0, srcCacheLen); diff --git a/src/main/java/org/apache/datasketches/theta/CompactThetaSketch.java b/src/main/java/org/apache/datasketches/theta/CompactThetaSketch.java index 10f2f4986..94461dfda 100644 --- a/src/main/java/org/apache/datasketches/theta/CompactThetaSketch.java +++ b/src/main/java/org/apache/datasketches/theta/CompactThetaSketch.java @@ -397,6 +397,7 @@ private static CompactThetaSketch heapifyV4(final MemorySegment srcSeg, final lo final int numEntriesBytes = extractNumEntriesBytesV4(srcSeg); final short seedHash = (short) extractSeedHash(srcSeg); PreambleUtil.checkSegmentSeedHash(srcSeg, seed); + PreambleUtil.checkCompressedImage(srcSeg); int offsetBytes = 8; long theta = Long.MAX_VALUE; if (preLongs > 1) { diff --git a/src/main/java/org/apache/datasketches/theta/DirectCompactCompressedSketch.java b/src/main/java/org/apache/datasketches/theta/DirectCompactCompressedSketch.java index 9f37241d2..b3c757709 100644 --- a/src/main/java/org/apache/datasketches/theta/DirectCompactCompressedSketch.java +++ b/src/main/java/org/apache/datasketches/theta/DirectCompactCompressedSketch.java @@ -58,6 +58,7 @@ final class DirectCompactCompressedSketch extends DirectCompactSketch { */ static DirectCompactCompressedSketch wrapInstance(final MemorySegment srcSeg, final short seedHash) { Util.checkSeedHashes((short) extractSeedHash(srcSeg), seedHash); + PreambleUtil.checkCompressedImage(srcSeg); return new DirectCompactCompressedSketch(srcSeg); } diff --git a/src/main/java/org/apache/datasketches/theta/DirectCompactSketch.java b/src/main/java/org/apache/datasketches/theta/DirectCompactSketch.java index 6c01d78a9..9a28383a6 100644 --- a/src/main/java/org/apache/datasketches/theta/DirectCompactSketch.java +++ b/src/main/java/org/apache/datasketches/theta/DirectCompactSketch.java @@ -68,6 +68,7 @@ class DirectCompactSketch extends CompactThetaSketch { */ static DirectCompactSketch wrapInstance(final MemorySegment srcSeg, final short seedHash) { Util.checkSeedHashes((short) extractSeedHash(srcSeg), seedHash); + PreambleUtil.checkCompactImage(srcSeg); return new DirectCompactSketch(srcSeg); } diff --git a/src/main/java/org/apache/datasketches/theta/PreambleUtil.java b/src/main/java/org/apache/datasketches/theta/PreambleUtil.java index ec6ad22aa..c135cd952 100644 --- a/src/main/java/org/apache/datasketches/theta/PreambleUtil.java +++ b/src/main/java/org/apache/datasketches/theta/PreambleUtil.java @@ -527,6 +527,50 @@ static int checkSegPreambleCap(final MemorySegment seg) { } } + /** + * Checks that a serial version 3 compact image, which is not empty and not a single item, holds the + * number of hashes it declares, before anything is allocated or read from it. + * @param seg the compact image + * @throws SketchesArgumentException if the image is too small or the count is invalid + */ + static void checkCompactImage(final MemorySegment seg) { + final int preLongs = checkSegPreambleCap(seg); + final int curCount = (preLongs > 1) ? extractCurCount(seg) : 1; + final long reqBytes = ((long) preLongs << 3) + ((long) curCount << 3); + if ((curCount < 0) || (seg.byteSize() < reqBytes)) { + throw new SketchesArgumentException("Possible corruption: curCount " + curCount + + " needs " + reqBytes + " bytes, MemorySegment has " + seg.byteSize()); + } + } + + /** + * Checks that a serial version 4 (compressed) compact image holds the number of entries it + * declares, before anything is allocated or read from it. + * @param seg the compressed compact image + * @throws SketchesArgumentException if the image is too small or its fields are invalid + */ + static void checkCompressedImage(final MemorySegment seg) { + final int preLongs = checkSegPreambleCap(seg); + final int entryBits = extractEntryBitsV4(seg); + final int numEntriesBytes = extractNumEntriesBytesV4(seg); + if ((preLongs < 1) || (preLongs > 2) || (entryBits < 1) || (entryBits > 63) + || (numEntriesBytes < 1) || (numEntriesBytes > 4)) { + throw new SketchesArgumentException("Possible corruption: preLongs: " + preLongs + + ", entryBits: " + entryBits + ", numEntriesBytes: " + numEntriesBytes); + } + final int offsetBytes = preLongs << 3; + Util.checkBounds(0, offsetBytes + numEntriesBytes, seg.byteSize()); + long numEntries = 0; + for (int i = 0; i < numEntriesBytes; i++) { + numEntries |= (long) Byte.toUnsignedInt(seg.get(JAVA_BYTE, offsetBytes + i)) << (i << 3); + } + if (numEntries > Integer.MAX_VALUE) { + throw new SketchesArgumentException("Possible corruption: numEntries: " + numEntries); + } + final long packedBytes = ((numEntries * entryBits) + 7) >>> 3; + Util.checkBounds(offsetBytes + numEntriesBytes, packedBytes, seg.byteSize()); + } + static short checkSegmentSeedHash(final MemorySegment seg, final long seed) { final short seedHashSeg = (short) extractSeedHash(seg); Util.checkSeedHashes(seedHashSeg, Util.computeSeedHash(seed)); //throws if bad seedHash diff --git a/src/main/java/org/apache/datasketches/theta/ThetaIntersectionImpl.java b/src/main/java/org/apache/datasketches/theta/ThetaIntersectionImpl.java index 3dfdf4fa1..df3a77f71 100644 --- a/src/main/java/org/apache/datasketches/theta/ThetaIntersectionImpl.java +++ b/src/main/java/org/apache/datasketches/theta/ThetaIntersectionImpl.java @@ -641,6 +641,21 @@ private static void segChecks(final MemorySegment srcSeg) { } //empty = true AND curCount_ = 0: OK } //else empty = false, curCount could be anything + //The hash table holds curCount entries and is no larger than for the largest nominal entries. + //A curCount of -1 means the intersection has no result yet. + final int lgArrLongs = extractLgArrLongs(srcSeg); + if ((lgArrLongs < 0) || (lgArrLongs > (ThetaUtil.MAX_LG_NOM_LONGS + 1)) + || (curCount < -1) || (curCount > (1 << lgArrLongs))) { + throw new SketchesArgumentException("Possible corruption: lgArrLongs: " + lgArrLongs + + ", curCount: " + curCount); + } + if (curCount > 0) { //the image must hold the hash table + final long reqBytes = ((long) CONST_PREAMBLE_LONGS << 3) + (8L << lgArrLongs); + if (srcSeg.byteSize() < reqBytes) { + throw new SketchesArgumentException("Possible corruption: MemorySegment size " + srcSeg.byteSize() + + " < required " + reqBytes); + } + } } } diff --git a/src/main/java/org/apache/datasketches/theta/UpdatableThetaSketch.java b/src/main/java/org/apache/datasketches/theta/UpdatableThetaSketch.java index f3b5898c5..cc6bfd322 100644 --- a/src/main/java/org/apache/datasketches/theta/UpdatableThetaSketch.java +++ b/src/main/java/org/apache/datasketches/theta/UpdatableThetaSketch.java @@ -30,13 +30,13 @@ import static org.apache.datasketches.theta.PreambleUtil.SER_VER; import static org.apache.datasketches.theta.PreambleUtil.checkSegPreambleCap; import static org.apache.datasketches.theta.PreambleUtil.checkSegmentSeedHash; +import static org.apache.datasketches.theta.PreambleUtil.extractCurCount; import static org.apache.datasketches.theta.PreambleUtil.extractFamilyID; import static org.apache.datasketches.theta.PreambleUtil.extractFlags; import static org.apache.datasketches.theta.PreambleUtil.extractLgResizeFactor; import static org.apache.datasketches.theta.PreambleUtil.extractP; import static org.apache.datasketches.theta.PreambleUtil.extractSerVer; import static org.apache.datasketches.theta.PreambleUtil.extractThetaLong; -import static org.apache.datasketches.theta.PreambleUtil.getUpdatableSegBytes; import static org.apache.datasketches.theta.UpdateReturnState.RejectedNullOrEmpty; import java.lang.foreign.MemorySegment; @@ -562,9 +562,19 @@ static void checkSegIntegrity(final MemorySegment srcSeg, final long expectedSee final short seedHash = checkSegmentSeedHash(srcSeg, expectedSeed); //byte 6,7 Util.checkSeedHashes(seedHash, Util.computeSeedHash(expectedSeed)); - //Check seg capacity, lgArrLongs + //Check lgNomLongs, lgArrLongs and curCount. The hash table grows to at most twice the nominal entries. + if ((lgNomLongs > ThetaUtil.MAX_LG_NOM_LONGS) || (lgArrLongs < 0) || (lgArrLongs > (lgNomLongs + 1))) { + throw new SketchesArgumentException("Possible corruption: lgNomLongs: " + lgNomLongs + + ", lgArrLongs: " + lgArrLongs); + } + final int curCount = extractCurCount(srcSeg); + if ((curCount < 0) || (curCount > (1 << lgArrLongs))) { + throw new SketchesArgumentException("Possible corruption: invalid curCount: " + curCount); + } + + //Check seg capacity, computed in long arithmetic so that it cannot overflow final long curCapBytes = srcSeg.byteSize(); - final int minReqBytes = getUpdatableSegBytes(lgArrLongs, preambleLongs); + final long minReqBytes = (8L << lgArrLongs) + ((long) preambleLongs << 3); if (curCapBytes < minReqBytes) { throw new SketchesArgumentException( "Possible corruption: Current MemorySegment size < min required size: " diff --git a/src/main/java/org/apache/datasketches/theta/WrappedCompactCompressedSketch.java b/src/main/java/org/apache/datasketches/theta/WrappedCompactCompressedSketch.java index d709188ed..bcf752690 100644 --- a/src/main/java/org/apache/datasketches/theta/WrappedCompactCompressedSketch.java +++ b/src/main/java/org/apache/datasketches/theta/WrappedCompactCompressedSketch.java @@ -24,6 +24,8 @@ import static org.apache.datasketches.theta.PreambleUtil.PREAMBLE_LONGS_BYTE; import static org.apache.datasketches.theta.PreambleUtil.wholeBytesToHoldBits; +import java.lang.foreign.MemorySegment; + import org.apache.datasketches.common.ByteArrayUtil; import org.apache.datasketches.common.Util; @@ -52,6 +54,7 @@ final class WrappedCompactCompressedSketch extends WrappedCompactSketch { */ static WrappedCompactCompressedSketch wrapInstance(final byte[] bytes, final short seedHash) { Util.checkSeedHashes(ByteArrayUtil.getShortLE(bytes, PreambleUtil.SEED_HASH_SHORT), seedHash); + PreambleUtil.checkCompressedImage(MemorySegment.ofArray(bytes)); return new WrappedCompactCompressedSketch(bytes); } diff --git a/src/main/java/org/apache/datasketches/theta/WrappedCompactSketch.java b/src/main/java/org/apache/datasketches/theta/WrappedCompactSketch.java index 53ef5c176..4697f1ce2 100644 --- a/src/main/java/org/apache/datasketches/theta/WrappedCompactSketch.java +++ b/src/main/java/org/apache/datasketches/theta/WrappedCompactSketch.java @@ -62,6 +62,7 @@ class WrappedCompactSketch extends CompactThetaSketch { */ static WrappedCompactSketch wrapInstance(final byte[] bytes, final short seedHash) { Util.checkSeedHashes(getShortLE(bytes, SEED_HASH_SHORT), seedHash); + PreambleUtil.checkCompactImage(MemorySegment.ofArray(bytes)); return new WrappedCompactSketch(bytes); } diff --git a/src/test/java/org/apache/datasketches/theta/ThetaCorruptImageTest.java b/src/test/java/org/apache/datasketches/theta/ThetaCorruptImageTest.java new file mode 100644 index 000000000..e9374c182 --- /dev/null +++ b/src/test/java/org/apache/datasketches/theta/ThetaCorruptImageTest.java @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.theta; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; +import java.util.function.Function; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt theta sketch, union and intersection images are rejected without large + * allocations, including when the result of a wrap is used later. + */ +public class ThetaCorruptImageTest { + private static final int LG_ARR_BYTE = 4; + private static final int CUR_COUNT_INT = 8; + + private static UpdatableThetaSketch sketch(final int n) { + final UpdatableThetaSketch sk = UpdatableThetaSketch.builder().build(); + for (int i = 0; i < n; i++) { sk.update(i); } + return sk; + } + + private static void checkRejected(final byte[] bytes, final Function deser) { + final long allocated = AllocationTracker.allocatedBytes(() -> deser.apply(MemorySegment.ofArray(bytes))); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> deser.apply(MemorySegment.ofArray(bytes))); + } + + @Test + public void checkValidImagesStillWork() { + final UpdatableThetaSketch sk = sketch(1000); + final byte[] c = sk.compact().toByteArray(); + assertEquals(CompactThetaSketch.heapify(MemorySegment.ofArray(c)).getRetainedEntries(), 1000); + assertEquals(CompactThetaSketch.wrap(MemorySegment.ofArray(c)).getRetainedEntries(), 1000); + assertEquals(CompactThetaSketch.wrap(c).getRetainedEntries(), 1000); + final byte[] z = sk.compact().toByteArrayCompressed(); + assertEquals(CompactThetaSketch.heapify(MemorySegment.ofArray(z)).getRetainedEntries(), 1000); + assertEquals(CompactThetaSketch.wrap(MemorySegment.ofArray(z)).getRetainedEntries(), 1000); + assertEquals(CompactThetaSketch.wrap(z).getRetainedEntries(), 1000); + assertEquals(UpdatableThetaSketch.heapify(MemorySegment.ofArray(sk.toByteArray())).getRetainedEntries(), 1000); + final ThetaUnion union = ThetaSetOperation.builder().buildUnion(); + union.union(sk); + assertEquals(((ThetaUnion) ThetaSetOperation.heapify(MemorySegment.ofArray(union.toByteArray()))) + .getResult().getRetainedEntries(), 1000); + final ThetaIntersection inter = ThetaSetOperation.builder().buildIntersection(); + inter.intersect(sk); + assertEquals(((ThetaIntersection) ThetaSetOperation.heapify(MemorySegment.ofArray(inter.toByteArray()))) + .getResult().getRetainedEntries(), 1000); + } + + @Test + public void checkCompactCountBeyondImage() { + final byte[] bytes = sketch(100).compact().toByteArray(); + ByteArrayUtil.putIntLE(bytes, CUR_COUNT_INT, 1 << 28); + checkRejected(bytes, CompactThetaSketch::heapify); + checkRejected(bytes, ThetaSketch::heapify); + checkRejected(bytes, seg -> CompactThetaSketch.wrap(seg).getCache()); + checkRejected(bytes, _ -> CompactThetaSketch.wrap(bytes).getCache()); + final byte[] negative = sketch(100).compact().toByteArray(); + ByteArrayUtil.putIntLE(negative, CUR_COUNT_INT, -1); + checkRejected(negative, CompactThetaSketch::heapify); + } + + @Test + public void checkCompressedCountBeyondImage() { + final byte[] bytes = sketch(100).compact().toByteArrayCompressed(); + final int numEntriesBytes = bytes[PreambleUtil.NUM_ENTRIES_BYTES_BYTE_V4]; + final int offset = (bytes[0] & 0x3F) << 3; + for (int i = 0; i < numEntriesBytes; i++) { bytes[offset + i] = (byte) 0xFF; } + checkRejected(bytes, CompactThetaSketch::heapify); + checkRejected(bytes, seg -> CompactThetaSketch.wrap(seg).getCache()); + checkRejected(bytes, _ -> CompactThetaSketch.wrap(bytes).getCache()); + final byte[] badBits = sketch(100).compact().toByteArrayCompressed(); + badBits[PreambleUtil.ENTRY_BITS_BYTE_V4] = 0; + checkRejected(badBits, CompactThetaSketch::heapify); + } + + @Test + public void checkUpdatableLargeTable() { + final byte[] bytes = sketch(100).toByteArray(); + bytes[LG_ARR_BYTE] = 30; //(8 << 30) overflows int + checkRejected(bytes, UpdatableThetaSketch::heapify); + checkRejected(bytes, ThetaSketch::heapify); + checkRejected(bytes, UpdatableThetaSketch::wrap); + final byte[] count = sketch(100).toByteArray(); + ByteArrayUtil.putIntLE(count, CUR_COUNT_INT, 1 << 28); + checkRejected(count, UpdatableThetaSketch::heapify); + } + + @Test + public void checkUnionLargeTable() { + final ThetaUnion union = ThetaSetOperation.builder().buildUnion(); + union.union(sketch(100)); + final byte[] bytes = union.toByteArray(); + bytes[LG_ARR_BYTE] = 30; + checkRejected(bytes, ThetaSetOperation::heapify); + checkRejected(bytes, seg -> ThetaUnion.wrap(seg).getResult()); + } + + @Test + public void checkIntersectionLargeTable() { + final ThetaIntersection inter = ThetaSetOperation.builder().buildIntersection(); + inter.intersect(sketch(100)); + final byte[] bytes = inter.toByteArray(); + bytes[LG_ARR_BYTE] = 30; + checkRejected(bytes, ThetaSetOperation::heapify); + checkRejected(bytes, seg -> ThetaIntersection.wrap(seg).getResult()); + final byte[] count = inter.toByteArray(); + ByteArrayUtil.putIntLE(count, CUR_COUNT_INT, 1 << 20); //more than the table holds + checkRejected(count, ThetaSetOperation::heapify); + } +} From f8f41cd283ffba8e6aba2d10ddb2cdcb51fc8da0 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:51:24 -0700 Subject: [PATCH 06/15] HLL: check lgK and declared sizes checkPreamble(), which heapify, wrap, writableWrap and HllUnion.heapify all call, now also checks: - that lgK is in its documented range of [4, 21]; - that a list count is no more than the list capacity; - that a set count is between 0 and K and set lgArr is at most lgK; - that an aux count is between 0 and K and aux lgArr is at most lgK. writableWrap() now validates the preamble before using lgK. Co-Authored-By: Claude Opus 5.5 --- .../apache/datasketches/hll/HllSketch.java | 2 +- .../org/apache/datasketches/hll/HllUtil.java | 36 ++++++ .../datasketches/hll/HllCorruptImageTest.java | 114 ++++++++++++++++++ 3 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 src/test/java/org/apache/datasketches/hll/HllCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/hll/HllSketch.java b/src/main/java/org/apache/datasketches/hll/HllSketch.java index b1e958651..b49474464 100644 --- a/src/main/java/org/apache/datasketches/hll/HllSketch.java +++ b/src/main/java/org/apache/datasketches/hll/HllSketch.java @@ -254,12 +254,12 @@ static final HllSketch writableWrap( final MemorySegment srcWseg, final boolean "Cannot perform a writableWrap of a sketch image that is in compact form. " + "Compact sketches are by definition immutable."); } + final CurMode curMode = checkPreamble(srcWseg); //checks lgK before it is used final int lgConfigK = extractLgK(srcWseg); final TgtHllType tgtHllType = extractTgtHllType(srcWseg); final long minBytes = getMaxUpdatableSerializationBytes(lgConfigK, tgtHllType); final long capBytes = srcWseg.byteSize(); HllUtil.checkSegSize(minBytes, capBytes); - final CurMode curMode = checkPreamble(srcWseg); final HllSketch directSketch; if (curMode == CurMode.HLL) { if (tgtHllType == TgtHllType.HLL_4) { diff --git a/src/main/java/org/apache/datasketches/hll/HllUtil.java b/src/main/java/org/apache/datasketches/hll/HllUtil.java index b87b4e2d2..721382200 100644 --- a/src/main/java/org/apache/datasketches/hll/HllUtil.java +++ b/src/main/java/org/apache/datasketches/hll/HllUtil.java @@ -23,8 +23,13 @@ import static org.apache.datasketches.hll.PreambleUtil.HASH_SET_PREINTS; import static org.apache.datasketches.hll.PreambleUtil.HLL_PREINTS; import static org.apache.datasketches.hll.PreambleUtil.LIST_PREINTS; +import static org.apache.datasketches.hll.PreambleUtil.extractAuxCount; import static org.apache.datasketches.hll.PreambleUtil.extractCurMode; import static org.apache.datasketches.hll.PreambleUtil.extractFamilyId; +import static org.apache.datasketches.hll.PreambleUtil.extractHashSetCount; +import static org.apache.datasketches.hll.PreambleUtil.extractLgArr; +import static org.apache.datasketches.hll.PreambleUtil.extractLgK; +import static org.apache.datasketches.hll.PreambleUtil.extractListCount; import static org.apache.datasketches.hll.PreambleUtil.extractPreInts; import static org.apache.datasketches.hll.PreambleUtil.extractSerVer; @@ -111,9 +116,40 @@ static CurMode checkPreamble(final MemorySegment seg) { throw new SketchesArgumentException("Possible Corruption, Invalid Preamble:" + PreambleUtil.toString(seg)); } + checkSizes(seg, curMode); return curMode; } + /** + * Checks lgK and the sizes declared by the image. + * LgK must be in the documented range. The coupon hash set and the auxiliary (exception) hash map + * are never larger than K, and a coupon list never holds more than its initial capacity. + * @param seg the image + * @param curMode the current mode of the image + */ + private static void checkSizes(final MemorySegment seg, final CurMode curMode) { + final int lgK = extractLgK(seg); + checkLgK(lgK); + if (curMode == CurMode.LIST) { + final int count = extractListCount(seg); + if (count > (1 << LG_INIT_LIST_SIZE)) { + throw new SketchesArgumentException("Possible Corruption: list count " + count); + } + } else if (curMode == CurMode.SET) { + final int count = extractHashSetCount(seg); + final int lgArr = extractLgArr(seg); + if ((count < 0) || (count > (1 << lgK)) || (lgArr > lgK)) { + throw new SketchesArgumentException("Possible Corruption: set count " + count + ", lgArr " + lgArr); + } + } else { //HLL + final int auxCount = extractAuxCount(seg); + final int lgArr = extractLgArr(seg); + if ((auxCount < 0) || (auxCount > (1 << lgK)) || ((auxCount > 0) && (lgArr > lgK))) { + throw new SketchesArgumentException("Possible Corruption: aux count " + auxCount + ", lgArr " + lgArr); + } + } + } + //Exceptions static void noWriteAccess() { throw new SketchesReadOnlyException( diff --git a/src/test/java/org/apache/datasketches/hll/HllCorruptImageTest.java b/src/test/java/org/apache/datasketches/hll/HllCorruptImageTest.java new file mode 100644 index 000000000..6f32b3d8f --- /dev/null +++ b/src/test/java/org/apache/datasketches/hll/HllCorruptImageTest.java @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.hll; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; +import java.util.function.Function; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt HLL images are rejected without large allocations. + */ +public class HllCorruptImageTest { + private static final int LG_K_BYTE = 3; + private static final int LG_ARR_BYTE = 4; + private static final int LIST_COUNT_BYTE = 6; + private static final int HASH_SET_COUNT_INT = 8; + private static final int AUX_COUNT_INT = 36; + + private static HllSketch sketch(final TgtHllType type, final int n) { + final HllSketch sk = new HllSketch(12, type); + for (int i = 0; i < n; i++) { sk.update(i); } + return sk; + } + + private static void checkRejected(final byte[] bytes, final Function deser) { + final long allocated = AllocationTracker.allocatedBytes(() -> deser.apply(MemorySegment.ofArray(bytes))); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> deser.apply(MemorySegment.ofArray(bytes))); + } + + private static void checkAllPaths(final byte[] bytes) { + checkRejected(bytes, HllSketch::heapify); + checkRejected(bytes, HllSketch::wrap); + checkRejected(bytes, HllSketch::writableWrap); + checkRejected(bytes, HllUnion::heapify); + } + + @Test + public void checkValidImagesStillWork() { + for (final TgtHllType type : TgtHllType.values()) { + for (final int n : new int[] {0, 5, 200, 10_000}) { + final HllSketch sk = sketch(type, n); + for (final byte[] b : new byte[][] {sk.toCompactByteArray(), sk.toUpdatableByteArray()}) { + assertEquals(HllSketch.heapify(MemorySegment.ofArray(b)).getEstimate(), sk.getEstimate()); + assertEquals(HllSketch.wrap(MemorySegment.ofArray(b)).getEstimate(), sk.getEstimate()); + } + } + } + } + + @Test + public void checkLgKTooLarge() { + for (final TgtHllType type : TgtHllType.values()) { + for (final int n : new int[] {5, 200, 10_000}) { //list, set and HLL modes + final byte[] bytes = sketch(type, n).toUpdatableByteArray(); + bytes[LG_K_BYTE] = 30; + checkAllPaths(bytes); + } + } + } + + @Test + public void checkListCount() { + final byte[] bytes = sketch(TgtHllType.HLL_8, 5).toUpdatableByteArray(); + bytes[LIST_COUNT_BYTE] = (byte) 200; + checkAllPaths(bytes); + } + + @Test + public void checkSetSizes() { + final byte[] lgArr = sketch(TgtHllType.HLL_8, 200).toUpdatableByteArray(); + lgArr[LG_ARR_BYTE] = 28; + checkAllPaths(lgArr); + final byte[] count = sketch(TgtHllType.HLL_8, 200).toCompactByteArray(); + ByteArrayUtil.putIntLE(count, HASH_SET_COUNT_INT, 1 << 28); + checkAllPaths(count); + } + + @Test + public void checkAuxSizes() { + final byte[] bytes = sketch(TgtHllType.HLL_4, 10_000).toUpdatableByteArray(); + ByteArrayUtil.putIntLE(bytes, AUX_COUNT_INT, 1 << 28); + checkAllPaths(bytes); + final byte[] lgArr = sketch(TgtHllType.HLL_4, 10_000).toUpdatableByteArray(); + if (ByteArrayUtil.getIntLE(lgArr, AUX_COUNT_INT) == 0) { ByteArrayUtil.putIntLE(lgArr, AUX_COUNT_INT, 1); } + lgArr[LG_ARR_BYTE] = 28; + checkAllPaths(lgArr); + } +} From e004666395cebe0bf1eb2f71916dac9fb045c526 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:54:06 -0700 Subject: [PATCH 07/15] Bloom: check the image size in heapify() For a non-empty image, check that the remaining bytes hold the bit array. An empty image still allocates the configured bit array, as designed. Co-Authored-By: Claude Opus 5.5 --- .../filters/bloomfilter/HeapBitArray.java | 4 ++ .../BloomFilterCorruptImageTest.java | 67 +++++++++++++++++++ 2 files changed, 71 insertions(+) create mode 100644 src/test/java/org/apache/datasketches/filters/bloomfilter/BloomFilterCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/filters/bloomfilter/HeapBitArray.java b/src/main/java/org/apache/datasketches/filters/bloomfilter/HeapBitArray.java index cb0e22727..0abd43f43 100644 --- a/src/main/java/org/apache/datasketches/filters/bloomfilter/HeapBitArray.java +++ b/src/main/java/org/apache/datasketches/filters/bloomfilter/HeapBitArray.java @@ -78,6 +78,10 @@ static HeapBitArray heapify(final PositionalSegment posSeg, final boolean isEmpt // will be -1 if dirty final long numBitsSet = posSeg.getLong(); + if (posSeg.getRemaining() < ((long) numLongs * Long.BYTES)) { + throw new SketchesArgumentException("Possible corruption: bit array of " + numLongs + + " longs does not fit in the remaining " + posSeg.getRemaining() + " bytes"); + } final long[] data = new long[numLongs]; posSeg.getLongArray(data, 0, numLongs); return new HeapBitArray(numBitsSet, data); diff --git a/src/test/java/org/apache/datasketches/filters/bloomfilter/BloomFilterCorruptImageTest.java b/src/test/java/org/apache/datasketches/filters/bloomfilter/BloomFilterCorruptImageTest.java new file mode 100644 index 000000000..14250be1f --- /dev/null +++ b/src/test/java/org/apache/datasketches/filters/bloomfilter/BloomFilterCorruptImageTest.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.filters.bloomfilter; + +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt BloomFilter images are rejected, or handled, without large allocations. + */ +public class BloomFilterCorruptImageTest { + private static final int NUM_LONGS_INT = 16; + private static final int LARGE = 1 << 28; //2 GB of longs + + @Test + public void checkNonEmptyArrayBeyondImage() { + final BloomFilter bf = BloomFilterBuilder.createBySize(1000, 3, 123L); + bf.update(1L); + final byte[] bytes = bf.toByteArray(); + ByteArrayUtil.putIntLE(bytes, NUM_LONGS_INT, LARGE); + final long allocated = AllocationTracker.allocatedBytes(() -> BloomFilter.heapify(MemorySegment.ofArray(bytes))); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> BloomFilter.heapify(MemorySegment.ofArray(bytes))); + } + + @SuppressWarnings("deprecation") + @Test + public void checkEmptyImageIsUsable() { + final BloomFilter bf = BloomFilter.heapify(MemorySegment.ofArray(BloomFilterBuilder.createBySize(1000, 3, 123L).toByteArray())); + assertFalse(bf.query(1L)); + bf.update(1L); + assertTrue(bf.query(1L)); + final BloomFilter bf2 = BloomFilter.heapify(MemorySegment.ofArray(bf.toByteArray())); + assertTrue(bf2.query(1L)); + final BloomFilter empty = BloomFilter.heapify(MemorySegment.ofArray(BloomFilterBuilder.createBySize(1000, 3, 123L).toByteArray())); + empty.union(bf); + assertTrue(empty.query(1L)); + final BloomFilter empty2 = BloomFilter.heapify(MemorySegment.ofArray(BloomFilterBuilder.createBySize(1000, 3, 123L).toByteArray())); + empty2.invert(); + assertTrue(empty2.query(2L)); + } +} From 3930b5ada9ac1a6222aa037255ec101b7a333f52 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:54:42 -0700 Subject: [PATCH 08/15] ArrayOfNumbersSerDe: check the item count Each item takes at least one byte (its type code), so check that the input can hold numItems items. PR 783 makes the same check in the UTF-8 and UTF-16 string SerDes. Co-Authored-By: Claude Opus 5.5 --- .../common/ArrayOfNumbersSerDe.java | 1 + .../common/ArrayOfNumbersSerDeBoundsTest.java | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 src/test/java/org/apache/datasketches/common/ArrayOfNumbersSerDeBoundsTest.java diff --git a/src/main/java/org/apache/datasketches/common/ArrayOfNumbersSerDe.java b/src/main/java/org/apache/datasketches/common/ArrayOfNumbersSerDe.java index 797d8d3b0..d04fb6167 100644 --- a/src/main/java/org/apache/datasketches/common/ArrayOfNumbersSerDe.java +++ b/src/main/java/org/apache/datasketches/common/ArrayOfNumbersSerDe.java @@ -127,6 +127,7 @@ public Number[] deserializeFromMemorySegment(final MemorySegment seg, final int public Number[] deserializeFromMemorySegment(final MemorySegment seg, final long offsetBytes, final int numItems) { Objects.requireNonNull(seg, "MemorySegment must not be null"); if (numItems <= 0) { return new Number[0]; } + Util.checkBounds(offsetBytes, numItems, seg.byteSize()); //each item has at least a 1-byte type code final Number[] array = new Number[numItems]; long offset = offsetBytes; for (int i = 0; i < numItems; i++) { diff --git a/src/test/java/org/apache/datasketches/common/ArrayOfNumbersSerDeBoundsTest.java b/src/test/java/org/apache/datasketches/common/ArrayOfNumbersSerDeBoundsTest.java new file mode 100644 index 000000000..90ba33686 --- /dev/null +++ b/src/test/java/org/apache/datasketches/common/ArrayOfNumbersSerDeBoundsTest.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.common; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.testng.annotations.Test; + +/** + * Checks that ArrayOfNumbersSerDe does not allocate from an item count that the input cannot hold. + */ +public class ArrayOfNumbersSerDeBoundsTest { + + @Test + public void checkImpossibleItemCount() { + final ArrayOfNumbersSerDe serDe = new ArrayOfNumbersSerDe(); + final MemorySegment seg = MemorySegment.ofArray(serDe.serializeToByteArray(new Number[] {1L, 2.0, 3})); + final int numItems = 1 << 28; //each item needs at least one byte + final long allocated = AllocationTracker.allocatedBytes(() -> serDe.deserializeFromMemorySegment(seg, 0, numItems)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> serDe.deserializeFromMemorySegment(seg, 0, numItems)); + assertEquals(serDe.deserializeFromMemorySegment(seg, 0, 3).length, 3); + } +} From 4e024585d37c0eebb3653c55529f64ac42688364 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 14:56:25 -0700 Subject: [PATCH 09/15] Quantiles items: check n against the image size No serialized item takes less than one bit (booleans are packed 8 per byte), so QuantilesItemsSketch.heapify() now requires that n is non-negative and that the number of retained items is no more than 8 times the remaining image bytes. Co-Authored-By: Claude Opus 5.5 --- .../quantiles/QuantilesItemsSketch.java | 9 ++- .../QuantilesItemsCorruptImageTest.java | 72 +++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 src/test/java/org/apache/datasketches/quantiles/QuantilesItemsCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/quantiles/QuantilesItemsSketch.java b/src/main/java/org/apache/datasketches/quantiles/QuantilesItemsSketch.java index 17c94940d..f007ac392 100644 --- a/src/main/java/org/apache/datasketches/quantiles/QuantilesItemsSketch.java +++ b/src/main/java/org/apache/datasketches/quantiles/QuantilesItemsSketch.java @@ -214,9 +214,16 @@ public static QuantilesItemsSketch heapify( //Not empty, must have valid preamble + min, max final long n = extractN(srcSeg); - //can't check MemorySegment capacity here, not enough information + //can't check the exact MemorySegment capacity here, not enough information final int extra = 2; //for min, max final int numSegItems = computeRetainedItems(k, n) + extra; + //No serialized item takes less than one bit (booleans are packed 8 per byte), + //so the image bounds the number of items before the buffer is allocated. + final long itemsBytes = segCapBytes - ((long) preambleLongs * Long.BYTES); + if ((n < 0) || ((long) numSegItems > (8L * itemsBytes))) { + throw new SketchesArgumentException("Possible corruption: n = " + n + " requires " + numSegItems + + " items, but only " + itemsBytes + " bytes remain"); + } //set class members sk.n_ = n; diff --git a/src/test/java/org/apache/datasketches/quantiles/QuantilesItemsCorruptImageTest.java b/src/test/java/org/apache/datasketches/quantiles/QuantilesItemsCorruptImageTest.java new file mode 100644 index 000000000..9335a5115 --- /dev/null +++ b/src/test/java/org/apache/datasketches/quantiles/QuantilesItemsCorruptImageTest.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.quantiles; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; +import java.util.Comparator; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ArrayOfLongsSerDe; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that a corrupt QuantilesItemsSketch image is rejected. + */ +public class QuantilesItemsCorruptImageTest { + private static final int N_LONG = 8; + private final ArrayOfLongsSerDe serDe = new ArrayOfLongsSerDe(); + + private byte[] image(final int n) { + final QuantilesItemsSketch sk = QuantilesItemsSketch.getInstance(Long.class, 32768, Comparator.naturalOrder()); + for (long i = 0; i < n; i++) { sk.update(i); } + return sk.toByteArray(serDe); + } + + @Test + public void checkLargeN() { + final byte[] bytes = image(1); + ByteArrayUtil.putLongLE(bytes, N_LONG, 1L << 60); + final MemorySegment seg = MemorySegment.ofArray(bytes); + final long allocated = AllocationTracker.allocatedBytes( + () -> QuantilesItemsSketch.heapify(Long.class, seg, Comparator.naturalOrder(), serDe)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, + () -> QuantilesItemsSketch.heapify(Long.class, seg, Comparator.naturalOrder(), serDe)); + final byte[] negative = image(1); + ByteArrayUtil.putLongLE(negative, N_LONG, -1L); + expectThrows(SketchesArgumentException.class, + () -> QuantilesItemsSketch.heapify(Long.class, MemorySegment.ofArray(negative), Comparator.naturalOrder(), serDe)); + } + + @Test + public void checkValidImagesStillWork() { + for (final int n : new int[] {1, 100, 100_000}) { + final QuantilesItemsSketch sk = + QuantilesItemsSketch.heapify(Long.class, MemorySegment.ofArray(image(n)), Comparator.naturalOrder(), serDe); + assertEquals(sk.getN(), n); + } + } +} From 024dc74e051a6ad194687598d78bdeb6f347d7a7 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 15:03:00 -0700 Subject: [PATCH 10/15] REQ: validate compactor fields - Check that k is even and in [4, 1024] before reading the compactors. - Check that 1 <= sectionSize <= k, numSections >= 1, and sectionSize * numSections <= k * k. Simulating the section growth rule for every legal k gives at most 0.96 * k * k. - Check that the count is non-negative and that the image holds the items. - Size each reconstructed buffer to the items it holds. The buffer grows by delta as needed, as it does after compaction. ReqCompactorTest compared the rebuilt buffer's capacity and backing array. It now compares the count, the sorted flag and the items. Co-Authored-By: Claude Opus 5.5 --- .../org/apache/datasketches/req/ReqSerDe.java | 25 ++++- .../datasketches/req/ReqCompactorTest.java | 18 ++-- .../datasketches/req/ReqCorruptImageTest.java | 96 +++++++++++++++++++ 3 files changed, 127 insertions(+), 12 deletions(-) create mode 100644 src/test/java/org/apache/datasketches/req/ReqCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/req/ReqSerDe.java b/src/main/java/org/apache/datasketches/req/ReqSerDe.java index c1f038bdd..89a65e345 100644 --- a/src/main/java/org/apache/datasketches/req/ReqSerDe.java +++ b/src/main/java/org/apache/datasketches/req/ReqSerDe.java @@ -147,6 +147,10 @@ static ReqSketch heapify(final MemorySegment seg) { final int k = posSeg.getShort() & 0xFFFF; final int numCompactors = posSeg.getByte() & 0xFF; final int numRawItems = posSeg.getByte() & 0xFF; + if (((k & 1) > 0) || (k < ReqSketch.MIN_K) || (k > 1024)) { + throw new SketchesArgumentException( + "Possible corruption: K must be even and in the range [4, 1024]: " + k); + } // extract different serialization formats final SerDeFormat deserFormat = getDeserFormat(empty, rawItems, numCompactors); switch (deserFormat) { @@ -162,7 +166,7 @@ static ReqSketch heapify(final MemorySegment seg) { } case EXACT: { assert preInts == 2; - final Compactor compactor = extractCompactor(posSeg, lvl0Sorted, hra); + final Compactor compactor = extractCompactor(posSeg, k, lvl0Sorted, hra); checkLgWeight(compactor, 0); //Construct sketch final long totalN = compactor.count; @@ -184,7 +188,7 @@ static ReqSketch heapify(final MemorySegment seg) { final List compactors = new ArrayList<>(); for (int i = 0; i < numCompactors; i++) { final boolean level0sorted = i == 0 ? lvl0Sorted : true; - final Compactor compactor = extractCompactor(posSeg, level0sorted, hra); + final Compactor compactor = extractCompactor(posSeg, k, level0sorted, hra); checkLgWeight(compactor, i); compactors.add(compactor.reqCompactor); } @@ -206,7 +210,7 @@ private static void checkLgWeight(final Compactor compactor, final int level) { } } - static final Compactor extractCompactor(final PositionalSegment posSeg, final boolean lvl0Sorted, + static final Compactor extractCompactor(final PositionalSegment posSeg, final int k, final boolean lvl0Sorted, final boolean hra) { final long state = posSeg.getLong(); final float sectionSizeFlt = posSeg.getFloat(); @@ -215,6 +219,17 @@ static final Compactor extractCompactor(final PositionalSegment posSeg, final bo final byte numSections = posSeg.getByte(); posSeg.incrementPosition(2); final int count = posSeg.getInt(); + // The section size starts at k and only shrinks, and as the number of sections doubles the section size + // shrinks by sqrt(2) to no less than MIN_K, so sectionSize * numSections stays below k * k. + if ((sectionSize < 1) || (sectionSize > k) || (numSections < 1) + || (((long) sectionSize * numSections) > ((long) k * k))) { + throw new SketchesArgumentException("Possible corruption: sectionSize: " + sectionSize + + ", numSections: " + numSections + ", k: " + k); + } + if ((count < 0) || (((long) count * Float.BYTES) > posSeg.getRemaining())) { + throw new SketchesArgumentException("Possible corruption: compactor count " + count + + " does not fit in the remaining " + posSeg.getRemaining() + " bytes"); + } final float[] arr = new float[count]; posSeg.getFloatArray(arr, 0, count); float minItem = Float.POSITIVE_INFINITY; @@ -224,8 +239,8 @@ static final Compactor extractCompactor(final PositionalSegment posSeg, final bo maxItem = max(maxItem, arr[i]); } final int delta = 2 * sectionSize * numSections; - final int nomCap = 2 * delta; - final int cap = max(count, nomCap); + // The buffer grows as needed, so it only has to hold the items in the image. + final int cap = count; final FloatBuffer fltBuf = FloatBuffer.reconstruct(arr, count, cap, delta, lvl0Sorted, hra); final ReqCompactor reqCompactor = new ReqCompactor(lgWt, hra, state, sectionSizeFlt, numSections, fltBuf); diff --git a/src/test/java/org/apache/datasketches/req/ReqCompactorTest.java b/src/test/java/org/apache/datasketches/req/ReqCompactorTest.java index 95f35831a..59e874272 100644 --- a/src/test/java/org/apache/datasketches/req/ReqCompactorTest.java +++ b/src/test/java/org/apache/datasketches/req/ReqCompactorTest.java @@ -62,7 +62,6 @@ public void checkSerDe() { private static void checkSerDeImpl(final int k, final boolean hra) { final ReqCompactor c1 = new ReqCompactor((byte)0, hra, k, null); final int nomCap = 2 * 3 * k; - final int cap = 2 * nomCap; final int delta = nomCap; final FloatBuffer fbuf = c1.getBuffer(); @@ -81,7 +80,7 @@ private static void checkSerDeImpl(final int k, final boolean hra) { final byte[] c1ser = c1.toByteArray(); //now deserialize final PositionalSegment posSeg = PositionalSegment.wrap(MemorySegment.ofArray(c1ser)); - final Compactor compactor = ReqSerDe.extractCompactor(posSeg, sorted, c1hra); + final Compactor compactor = ReqSerDe.extractCompactor(posSeg, k, sorted, c1hra); final ReqCompactor c2 = compactor.reqCompactor; assertEquals(compactor.minItem, minV); assertEquals(compactor.maxItem, maxV); @@ -93,9 +92,14 @@ private static void checkSerDeImpl(final int k, final boolean hra) { assertEquals(c2.getLgWeight(), lgWeight); assertEquals(c2.isHighRankAccuracy(), c1hra); final FloatBuffer fbuf2 = c2.getBuffer(); - assertEquals(fbuf2.getCapacity(), cap); + //deserialization sizes the buffer to the items it holds; the buffer grows as needed + assertTrue(fbuf2.getCapacity() >= compactor.count); assertEquals(fbuf2.getDelta(), delta); - assertTrue(fbuf.isEqualTo(fbuf2)); + assertEquals(fbuf2.getCount(), fbuf.getCount()); + assertEquals(fbuf2.isSorted(), fbuf.isSorted()); + for (int i = 0; i < fbuf.getCount(); i++) { + assertEquals(fbuf2.getItem(i), fbuf.getItem(i)); + } } @Test @@ -114,7 +118,7 @@ private static void checkSerDeNegativeImpl(final int k, final boolean hra) { } final byte[] c1ser = c1.toByteArray(); final PositionalSegment posSeg = PositionalSegment.wrap(MemorySegment.ofArray(c1ser)); - final Compactor compactor = ReqSerDe.extractCompactor(posSeg, fbuf.isSorted(), hra); + final Compactor compactor = ReqSerDe.extractCompactor(posSeg, k, fbuf.isSorted(), hra); assertEquals(compactor.minItem, -nomCap); assertEquals(compactor.maxItem, -1f); } @@ -136,7 +140,7 @@ private static void checkSerDeMixedImpl(final int k, final boolean hra) { } final byte[] c1ser = c1.toByteArray(); final PositionalSegment posSeg = PositionalSegment.wrap(MemorySegment.ofArray(c1ser)); - final Compactor compactor = ReqSerDe.extractCompactor(posSeg, fbuf.isSorted(), hra); + final Compactor compactor = ReqSerDe.extractCompactor(posSeg, k, fbuf.isSorted(), hra); assertEquals(compactor.minItem, (float) -half); assertEquals(compactor.maxItem, (float) (half - 1)); } @@ -158,7 +162,7 @@ private static void checkSerDeInfiniteImpl(final int k, final boolean hra) { final byte[] c1ser = c1.toByteArray(); final PositionalSegment posSeg = PositionalSegment.wrap(MemorySegment.ofArray(c1ser)); - final Compactor compactor = ReqSerDe.extractCompactor(posSeg, fbuf.isSorted(), hra); + final Compactor compactor = ReqSerDe.extractCompactor(posSeg, k, fbuf.isSorted(), hra); assertEquals(compactor.minItem, Float.NEGATIVE_INFINITY); assertEquals(compactor.maxItem, Float.POSITIVE_INFINITY); } diff --git a/src/test/java/org/apache/datasketches/req/ReqCorruptImageTest.java b/src/test/java/org/apache/datasketches/req/ReqCorruptImageTest.java new file mode 100644 index 000000000..c8f7ca939 --- /dev/null +++ b/src/test/java/org/apache/datasketches/req/ReqCorruptImageTest.java @@ -0,0 +1,96 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.req; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt ReqSketch images are rejected without large allocations. + */ +public class ReqCorruptImageTest { + private static final int EST_SECTION_SIZE_FLOAT = 32; //first compactor, estimation mode + private static final int EST_NUM_SECTIONS_BYTE = 37; + private static final int EST_COUNT_INT = 40; + private static final int EXACT_COUNT_INT = 24; //the only compactor, exact mode + + private static byte[] image(final int n) { + final ReqSketch sk = ReqSketch.builder().setK(12).build(); + for (int i = 0; i < n; i++) { sk.update(i); } + return sk.toByteArray(); + } + + private static void checkRejected(final byte[] bytes) { + final long allocated = AllocationTracker.allocatedBytes(() -> ReqSketch.heapify(MemorySegment.ofArray(bytes))); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> ReqSketch.heapify(MemorySegment.ofArray(bytes))); + } + + @Test + public void checkValidImagesStillWork() { + for (final int n : new int[] {0, 5, 50, 10_000, 1_000_000}) { + final ReqSketch sk = ReqSketch.heapify(MemorySegment.ofArray(image(n))); + assertEquals(sk.getN(), n); + if (n > 0) { + sk.update(-1f); //the reconstructed buffers grow as needed + assertEquals(sk.getN(), n + 1); + } + } + } + + @Test + public void checkCountBeyondImage() { + final byte[] est = image(10_000); + ByteArrayUtil.putIntLE(est, EST_COUNT_INT, 1 << 28); + checkRejected(est); + final byte[] exact = image(50); + ByteArrayUtil.putIntLE(exact, EXACT_COUNT_INT, 1 << 28); + checkRejected(exact); + } + + @Test + public void checkSectionsTooLarge() { + final byte[] size = image(10_000); + ByteArrayUtil.putFloatLE(size, EST_SECTION_SIZE_FLOAT, 1e8f); + checkRejected(size); + final byte[] sections = image(10_000); + sections[EST_NUM_SECTIONS_BYTE] = 127; + ByteArrayUtil.putFloatLE(sections, EST_SECTION_SIZE_FLOAT, 12f); + checkRejected(sections); + final byte[] negative = image(10_000); + negative[EST_NUM_SECTIONS_BYTE] = -1; + checkRejected(negative); + } + + @Test + public void checkInvalidK() { + final byte[] bytes = image(10_000); + ByteArrayUtil.putShortLE(bytes, 4, (short) 60_000); + checkRejected(bytes); + } +} From 5061899e0965fb6e9e5d1479a1b51c75f73f0fe3 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 15:05:34 -0700 Subject: [PATCH 11/15] VarOpt: check region counts and the image size - Check the region counts against k: a sketch in sampling mode holds exactly k items, and one in warm-up holds at most k. - Check that the image holds the H-region weights. - Size the weight array to hCount, and the item, weight and mark lists to the items in the image. The ArrayLists grow as needed, so this does not change behavior. Co-Authored-By: Claude Opus 5.5 --- .../sampling/VarOptItemsSketch.java | 18 +++- .../sampling/VarOptCorruptImageTest.java | 98 +++++++++++++++++++ 2 files changed, 111 insertions(+), 5 deletions(-) create mode 100644 src/test/java/org/apache/datasketches/sampling/VarOptCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/sampling/VarOptItemsSketch.java b/src/main/java/org/apache/datasketches/sampling/VarOptItemsSketch.java index 121b6678a..1f3c05a9b 100644 --- a/src/main/java/org/apache/datasketches/sampling/VarOptItemsSketch.java +++ b/src/main/java/org/apache/datasketches/sampling/VarOptItemsSketch.java @@ -345,6 +345,11 @@ public static VarOptItemsSketch heapify(final MemorySegment srcSeg, final int preLongBytes = numPreLongs << 3; + // In sampling mode the sketch holds exactly k items; while warming up it holds at most k. + if (((rCount > 0) && (((long) hCount + rCount) != k)) || ((rCount == 0) && (hCount > k))) { + throw new SketchesArgumentException("Possible Corruption: H and R region counts " + hCount + " and " + + rCount + " are not valid for k = " + k); + } final int totalItems = hCount + rCount; int allocatedItems = k + 1; // default to full @@ -361,10 +366,13 @@ public static VarOptItemsSketch heapify(final MemorySegment srcSeg, } } - // allocate full-sized ArrayLists, but we store only hCount weights at any moment + // The ArrayLists grow as needed, so size them to the items in the image rather than to k. + // We store only hCount weights at any moment. final long weightOffsetBytes = TOTAL_WEIGHT_R_DOUBLE + (rCount > 0 ? Double.BYTES : 0); - final ArrayList weightList = new ArrayList<>(allocatedItems); - final double[] wts = new double[allocatedItems]; + Util.checkBounds(weightOffsetBytes, (long) hCount * Double.BYTES, srcSeg.byteSize()); + final int listCapacity = (int) Math.min(allocatedItems, (long) totalItems + 1); + final ArrayList weightList = new ArrayList<>(listCapacity); + final double[] wts = new double[hCount]; MemorySegment.copy(srcSeg, JAVA_DOUBLE_UNALIGNED, weightOffsetBytes, wts, 0, hCount); // can't use Arrays.asList(wts) since double[] rather than Double[] @@ -383,7 +391,7 @@ public static VarOptItemsSketch heapify(final MemorySegment srcSeg, if (isGadget) { final long markOffsetBytes = preLongBytes + ((long) hCount * Double.BYTES); markBytes = ArrayOfBooleansSerDe.computeBytesNeeded(hCount); - markList = new ArrayList<>(allocatedItems); + markList = new ArrayList<>(listCapacity); final ArrayOfBooleansSerDe booleansSerDe = new ArrayOfBooleansSerDe(); final Boolean[] markArray = booleansSerDe.deserializeFromMemorySegment( @@ -398,7 +406,7 @@ public static VarOptItemsSketch heapify(final MemorySegment srcSeg, final long offsetBytes = preLongBytes + ((long) hCount * Double.BYTES) + markBytes; final T[] data = serDe.deserializeFromMemorySegment(srcSeg.asSlice(offsetBytes), 0, totalItems); final List wrappedData = Arrays.asList(data); - final ArrayList dataList = new ArrayList<>(allocatedItems); + final ArrayList dataList = new ArrayList<>(listCapacity); dataList.addAll(wrappedData.subList(0, hCount)); // Load items in R as needed diff --git a/src/test/java/org/apache/datasketches/sampling/VarOptCorruptImageTest.java b/src/test/java/org/apache/datasketches/sampling/VarOptCorruptImageTest.java new file mode 100644 index 000000000..706b58d85 --- /dev/null +++ b/src/test/java/org/apache/datasketches/sampling/VarOptCorruptImageTest.java @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.sampling; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ArrayOfLongsSerDe; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.ResizeFactor; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt VarOptItemsSketch images are rejected without large allocations. + */ +public class VarOptCorruptImageTest { + private static final int K_INT = 4; + private static final int H_COUNT_INT = 16; + private final ArrayOfLongsSerDe serDe = new ArrayOfLongsSerDe(); + + private byte[] image(final int k, final int n, final ResizeFactor rf) { + final VarOptItemsSketch sk = VarOptItemsSketch.newInstance(k, rf); + for (long i = 0; i < n; i++) { sk.update(i, 1.0 + (i % 7)); } + return sk.toByteArray(serDe); + } + + private void checkRejected(final byte[] bytes) { + final long allocated = AllocationTracker.allocatedBytes( + () -> VarOptItemsSketch.heapify(MemorySegment.ofArray(bytes), serDe)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> VarOptItemsSketch.heapify(MemorySegment.ofArray(bytes), serDe)); + } + + @Test + public void checkValidImagesStillWork() { + for (final ResizeFactor rf : ResizeFactor.values()) { + for (final int n : new int[] {1, 10, 32, 50, 1000}) { + final VarOptItemsSketch sk = VarOptItemsSketch.heapify(MemorySegment.ofArray(image(32, n, rf)), serDe); + assertEquals(sk.getN(), n); + for (long i = 0; i < 100; i++) { sk.update(i, 2.0); } //grows as needed + assertEquals(sk.getN(), n + 100); + } + } + } + + @Test + public void checkLargeKInSamplingMode() { + final byte[] bytes = image(32, 1000, ResizeFactor.X8); + ByteArrayUtil.putIntLE(bytes, K_INT, 1 << 28); + checkRejected(bytes); + } + + @Test + public void checkLargeKInWarmupMode() { + final byte[] bytes = image(32, 10, ResizeFactor.X1); //X1 would allocate the full k at once + ByteArrayUtil.putIntLE(bytes, K_INT, 1 << 28); + final long allocated = AllocationTracker.allocatedBytes( + () -> VarOptItemsSketch.heapify(MemorySegment.ofArray(bytes), serDe)); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + final VarOptItemsSketch sk = VarOptItemsSketch.heapify(MemorySegment.ofArray(bytes), serDe); + assertEquals(sk.getK(), 1 << 28); + assertEquals(sk.getNumSamples(), 10); + sk.update(100L, 1.0); + assertEquals(sk.getNumSamples(), 11); + } + + @Test + public void checkCountsBeyondImage() { + final byte[] h = image(32, 10, ResizeFactor.X8); + ByteArrayUtil.putIntLE(h, H_COUNT_INT, 1 << 28); //more than k + checkRejected(h); + final byte[] h2 = image(1 << 30, 10, ResizeFactor.X8); //large legal k, so only the image size limits hCount + ByteArrayUtil.putIntLE(h2, H_COUNT_INT, 1 << 28); + checkRejected(h2); + } +} From 9cf8254cd577a779ee924e8f4c4dfe989ddd77a5 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Thu, 29 Oct 2026 15:09:39 -0700 Subject: [PATCH 12/15] CPC: check stream lengths and coupon counts - Compute the required image size in long arithmetic, and bounds-check each stream against the image. - Check that 0 <= numCoupons <= 64k (k rows of 64 columns), that 0 <= numSv <= numCoupons, and that numSv <= 32 * svLengthInts, since each surprising value takes at least one bit of the compressed stream. Co-Authored-By: Claude Opus 5.5 --- .../datasketches/cpc/CompressedState.java | 23 ++++- .../apache/datasketches/cpc/PreambleUtil.java | 2 + .../datasketches/cpc/CpcCorruptImageTest.java | 86 +++++++++++++++++++ 3 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 src/test/java/org/apache/datasketches/cpc/CpcCorruptImageTest.java diff --git a/src/main/java/org/apache/datasketches/cpc/CompressedState.java b/src/main/java/org/apache/datasketches/cpc/CompressedState.java index a0e88df50..cde401d7c 100644 --- a/src/main/java/org/apache/datasketches/cpc/CompressedState.java +++ b/src/main/java/org/apache/datasketches/cpc/CompressedState.java @@ -49,6 +49,7 @@ import java.lang.foreign.MemorySegment; +import org.apache.datasketches.common.SketchesArgumentException; import org.apache.datasketches.common.Util; /** @@ -115,7 +116,7 @@ int getWindowOffset() { long getRequiredSerializedBytes() { final Format format = getFormat(); final int preInts = getDefinedPreInts(format); - return 4L * (preInts + csvLengthInts + cwLengthInts); + return 4L * ((long) preInts + csvLengthInts + cwLengthInts); //long arithmetic, so it cannot overflow } static CompressedState importFromSegment(final MemorySegment seg) { @@ -216,10 +217,28 @@ static CompressedState importFromSegment(final MemorySegment seg) { } } checkCapacity(seg.byteSize(), - 4L * (getPreInts(seg) + state.csvLengthInts + state.cwLengthInts)); + 4L * ((long) getPreInts(seg) + state.csvLengthInts + state.cwLengthInts)); + checkCounts(state); return state; } + /** + * Checks the coupon and surprising-value counts read from an image, before they are used to size + * the arrays that hold the uncompressed values. The matrix has k rows of 64 columns, so there are + * at most 64k coupons. The surprising values are a subset of the coupons, and each one takes at + * least one bit of the compressed stream. + * @param state the imported state + */ + private static void checkCounts(final CompressedState state) { + final long k = 1L << state.lgK; + if ((state.numCoupons < 0) || (state.numCoupons > (64L * k)) + || (state.numCsv < 0) || (state.numCsv > state.numCoupons) + || (state.numCsv > (32L * state.csvLengthInts))) { + throw new SketchesArgumentException("Possible corruption: numCoupons: " + state.numCoupons + + ", numSv: " + state.numCsv + ", svLengthInts: " + state.csvLengthInts); + } + } + void exportToSegment(final MemorySegment wseg) { final Format format = getFormat(); switch (format) { diff --git a/src/main/java/org/apache/datasketches/cpc/PreambleUtil.java b/src/main/java/org/apache/datasketches/cpc/PreambleUtil.java index f24c86dd1..89cb71aca 100644 --- a/src/main/java/org/apache/datasketches/cpc/PreambleUtil.java +++ b/src/main/java/org/apache/datasketches/cpc/PreambleUtil.java @@ -375,6 +375,7 @@ static long getWStreamOffset(final MemorySegment seg) { static int[] getSvStream(final MemorySegment seg) { final long offset = getSvStreamOffset(seg); final int svLengthInts = getSvLengthInts(seg); + checkBounds(offset, (long) svLengthInts * Integer.BYTES, seg.byteSize()); //before allocating final int[] svStream = new int[svLengthInts]; MemorySegment.copy(seg, JAVA_INT_UNALIGNED, offset, svStream, 0, svLengthInts); return svStream; @@ -383,6 +384,7 @@ static int[] getSvStream(final MemorySegment seg) { static int[] getWStream(final MemorySegment seg) { final long offset = getWStreamOffset(seg); final int wLength = getWLengthInts(seg); + checkBounds(offset, (long) wLength * Integer.BYTES, seg.byteSize()); //before allocating final int[] wStream = new int[wLength]; MemorySegment.copy(seg, JAVA_INT_UNALIGNED, offset, wStream, 0, wLength); return wStream; diff --git a/src/test/java/org/apache/datasketches/cpc/CpcCorruptImageTest.java b/src/test/java/org/apache/datasketches/cpc/CpcCorruptImageTest.java new file mode 100644 index 000000000..9dca351a8 --- /dev/null +++ b/src/test/java/org/apache/datasketches/cpc/CpcCorruptImageTest.java @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.datasketches.cpc; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import java.lang.foreign.MemorySegment; + +import org.apache.datasketches.common.AllocationTracker; +import org.apache.datasketches.common.ByteArrayUtil; +import org.apache.datasketches.common.SketchesArgumentException; +import org.testng.annotations.Test; + +/** + * Checks that corrupt CpcSketch images are rejected without large allocations. + */ +public class CpcCorruptImageTest { + + private static byte[] image(final int n) { + final CpcSketch sk = new CpcSketch(11); + for (int i = 0; i < n; i++) { sk.update(i); } + return sk.toByteArray(); + } + + @Test + public void checkValidImagesStillWork() { + for (final int n : new int[] {0, 5, 100, 3000, 100_000}) { + final CpcSketch sk = new CpcSketch(11); + for (int i = 0; i < n; i++) { sk.update(i); } + assertEquals(CpcSketch.heapify(MemorySegment.ofArray(sk.toByteArray())).getEstimate(), sk.getEstimate()); + } + } + + private static void checkRejected(final byte[] bytes) { + final long allocated = AllocationTracker.allocatedBytes(() -> CpcSketch.heapify(MemorySegment.ofArray(bytes))); + assertTrue(allocated < (1 << 20), "allocated " + allocated); + expectThrows(SketchesArgumentException.class, () -> CpcSketch.heapify(MemorySegment.ofArray(bytes))); + } + + private static int offset(final byte[] bytes, final PreambleUtil.HiField field) { + return (int) PreambleUtil.getHiFieldOffset(PreambleUtil.getFormat(MemorySegment.ofArray(bytes)), field); + } + + @Test + public void checkNumCouponsTooLarge() { + for (final int n : new int[] {100, 3000, 100_000}) { //sparse, hybrid and pinned or sliding formats + for (final int numCoupons : new int[] {Integer.MAX_VALUE, 1 << 25, -1}) { + final byte[] bytes = image(n); + ByteArrayUtil.putIntLE(bytes, offset(bytes, PreambleUtil.HiField.NUM_COUPONS), numCoupons); + checkRejected(bytes); + } + } + } + + @Test + public void checkStreamLengthTooLarge() { + for (final int n : new int[] {100, 3000, 100_000}) { + for (final int len : new int[] {Integer.MAX_VALUE, 1 << 28, -1}) { + final byte[] bytes = image(n); + final PreambleUtil.HiField field = (n > 3000) ? PreambleUtil.HiField.W_LENGTH_INTS + : PreambleUtil.HiField.SV_LENGTH_INTS; + ByteArrayUtil.putIntLE(bytes, offset(bytes, field), len); + checkRejected(bytes); + } + } + } +} From e9752c7352cf4c8711ed2c4ec0a5fae8e3a5dffb Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Wed, 4 Nov 2026 12:50:42 -0800 Subject: [PATCH 13/15] Start 9.2.0-SNAPSHOT, update NOTICE and the users list address Co-Authored-By: Claude Opus 5.5 --- NOTICE | 2 +- pom.xml | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/NOTICE b/NOTICE index 17bbe0d20..e367c88ea 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ Apache DataSketches Java -Copyright 2025 The Apache Software Foundation +Copyright 2026 The Apache Software Foundation Copyright 2015-2018 Yahoo Inc. Copyright 2019-2020 Verizon Media diff --git a/pom.xml b/pom.xml index 40fa32fe6..e3b169664 100644 --- a/pom.xml +++ b/pom.xml @@ -27,7 +27,7 @@ under the License. org.apache.datasketches datasketches-java - 9.1.0-SNAPSHOT + 9.2.0-SNAPSHOT jar ${project.artifactId} @@ -53,9 +53,9 @@ under the License. DataSketches Users - user-subscribe@datasketches.apache.org - user-unsubscribe@datasketches.apache.org - user@datasketches.apache.org + users-subscribe@datasketches.apache.org + users-unsubscribe@datasketches.apache.org + users@datasketches.apache.org https://mail-archives.apache.org/mod_mbox/datasketches-user From 745504b7dba5529824822843ebe995361d56ca62 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Wed, 4 Nov 2026 14:10:48 -0800 Subject: [PATCH 14/15] Allow Maven builds from the source release, which has no .git directory Co-Authored-By: Claude Opus 5.5 --- pom.xml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index e3b169664..fce945fcc 100644 --- a/pom.xml +++ b/pom.xml @@ -574,6 +574,8 @@ under the License. + + false true ${project.basedir}/git.properties @@ -606,7 +608,7 @@ under the License. **/*.yml **/.* **/test-output/**/* - + **/git.properties **/doc/** **/docs/** **/*.sk From 689a9461d205e181ab90920eabe6c5d6dbcf9984 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Fri, 6 Nov 2026 13:37:56 -0800 Subject: [PATCH 15/15] NOTICE: ASF copyright is 2019-2026 The project's first ASF release was in 2019. This matches datasketches-cpp. Co-Authored-By: Claude Opus 5.5 --- NOTICE | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NOTICE b/NOTICE index e367c88ea..fe8efd6f4 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ Apache DataSketches Java -Copyright 2026 The Apache Software Foundation +Copyright 2019-2026 The Apache Software Foundation Copyright 2015-2018 Yahoo Inc. Copyright 2019-2020 Verizon Media