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/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/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/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/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/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/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/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/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/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);
+ }
+}
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));
+ }
+
}
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);
+ }
+ }
+ }
+}
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));
+ }
+}
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) { }
+ }
+
}
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);
+ }
+}
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);
+ }
+ }
+}
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);
+ }
+}
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);
+ }
+}
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);
+ }
+}
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);
+ }
+}
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());
+ }
+}