Validate deserialized fields in MergingDigest.fromBytes()

core/src/main/java/com/tdunning/math/stats/MergingDigest.java

@@ -895,6 +895,42 @@ public void asSmallBytes(ByteBuffer buf) {
         }
     }
 
+    /**
+     * Upper bound on compression parameter to prevent excessive memory allocation.
+     * With compression = 1_000_000, the internal arrays are roughly 2M entries (~16MB),
+     * which is generous for any real use case. Values beyond this are almost certainly
+     * from corrupted or malicious input.
+     */
+    static final double MAX_COMPRESSION = 1_000_000.0;
+
+    /**
+     * Checks that a weight value is finite and positive.
+     * Uses negated form {@code !(w > 0)} so that NaN is correctly rejected
+     * (since NaN > 0 is false, the negation catches it).
+     */
+    private static void checkWeight(double w, int index) {
+        if (!(w > 0) || Double.isInfinite(w)) {
+            throw new IllegalArgumentException(
+                    "Invalid weight " + w + " at centroid " + index + " (must be finite and positive)");
+        }
+    }
+
+    /**
+     * Checks that a mean value is finite and in non-decreasing order.
+     * Uses negated form {@code !(m >= prevMean)} so that NaN is correctly rejected.
+     */
+    private static void checkMean(double m, double prevMean, int index) {
+        if (Double.isNaN(m) || Double.isInfinite(m)) {
+            throw new IllegalArgumentException(
+                    "Invalid mean " + m + " at centroid " + index + " (must be finite)");
+        }
+        if (!(m >= prevMean)) {
+            throw new IllegalArgumentException(
+                    "Centroids not in non-decreasing order at index " + index +
+                            ": " + prevMean + " > " + m);
+        }
+    }
+
     @SuppressWarnings("WeakerAccess")
     public static MergingDigest fromBytes(ByteBuffer buf) {
         int encoding = buf.getInt();
@@ -903,31 +939,79 @@ public static MergingDigest fromBytes(ByteBuffer buf) {
             double max = buf.getDouble();
             double compression = buf.getDouble();
             int n = buf.getInt();
+            if (!(compression > 0) || compression > MAX_COMPRESSION) {
+                throw new IllegalArgumentException(
+                        "Invalid compression: " + compression + " (must be finite and in (0, " + MAX_COMPRESSION + "])");
+            }
+            if (n < 0) {
+                throw new IllegalArgumentException("Invalid centroid count: " + n);
+            }
+            if (n > buf.remaining() / 16) {
+                throw new IllegalArgumentException(
+                        "Centroid count " + n + " requires more data than available in buffer");
+            }
             MergingDigest r = new MergingDigest(compression);
+            if (n > r.weight.length) {
+                throw new IllegalArgumentException(
+                        "Centroid count " + n + " exceeds digest capacity " + r.weight.length);
+            }
             r.setMinMax(min, max);
             r.lastUsedCell = n;
+            double prevMean = Double.NEGATIVE_INFINITY;
             for (int i = 0; i < n; i++) {
                 r.weight[i] = buf.getDouble();
                 r.mean[i] = buf.getDouble();
-
+                checkWeight(r.weight[i], i);
+                checkMean(r.mean[i], prevMean, i);
+                prevMean = r.mean[i];
                 r.totalWeight += r.weight[i];
             }
+            if (r.totalWeight < r.lastUsedCell) {
+                throw new IllegalArgumentException(
+                        "Total weight " + r.totalWeight + " is less than centroid count " + r.lastUsedCell);
+            }
             return r;
         } else if (encoding == Encoding.SMALL_ENCODING.code) {
             double min = buf.getDouble();
             double max = buf.getDouble();
             double compression = buf.getFloat();
             int n = buf.getShort();
             int bufferSize = buf.getShort();
+            if (!(compression > 0) || compression > MAX_COMPRESSION) {
+                throw new IllegalArgumentException(
+                        "Invalid compression: " + compression + " (must be finite and in (0, " + MAX_COMPRESSION + "])");
+            }
+            if (n <= 0) {
+                throw new IllegalArgumentException("Invalid main buffer size: " + n);
+            }
+            if (bufferSize <= 0) {
+                throw new IllegalArgumentException("Invalid buffer size: " + bufferSize);
+            }
             MergingDigest r = new MergingDigest(compression, bufferSize, n);
             r.setMinMax(min, max);
-            r.lastUsedCell = buf.getShort();
+            int lastUsedCell = buf.getShort();
+            if (lastUsedCell < 0 || lastUsedCell > n) {
+                throw new IllegalArgumentException(
+                        "Invalid lastUsedCell " + lastUsedCell + " (must be in [0, " + n + "])");
+            }
+            if (lastUsedCell > buf.remaining() / 8) {
+                throw new IllegalArgumentException(
+                        "lastUsedCell " + lastUsedCell + " requires more data than available in buffer");
+            }
+            r.lastUsedCell = lastUsedCell;
+            double prevMean = Double.NEGATIVE_INFINITY;
             for (int i = 0; i < r.lastUsedCell; i++) {
                 r.weight[i] = buf.getFloat();
                 r.mean[i] = buf.getFloat();
-
+                checkWeight(r.weight[i], i);
+                checkMean(r.mean[i], prevMean, i);
+                prevMean = r.mean[i];
                 r.totalWeight += r.weight[i];
             }
+            if (r.lastUsedCell > 0 && r.totalWeight < r.lastUsedCell) {
+                throw new IllegalArgumentException(
+                        "Total weight " + r.totalWeight + " is less than centroid count " + r.lastUsedCell);
+            }
             return r;
         } else {
             throw new IllegalStateException("Invalid format for serialized histogram");