Fix integer overflows in AssemblyScript & Typescript implementations

assembly/Histogram.ts

@@ -49,6 +49,8 @@ export default class Histogram<T, U> extends AbstractHistogramBase<T, U> {
   counts: T;
   totalCount: u64 = 0;
 
+  maxBucketSize: number;
+
   constructor(
     lowestDiscernibleValue: u64,
     highestTrackableValue: u64,
@@ -120,6 +122,8 @@ export default class Histogram<T, U> extends AbstractHistogramBase<T, U> {
     // @ts-ignore
     this.counts = instantiate<T>(this.countsArrayLength);
 
+    this.maxBucketSize = 2 ** (sizeof<U>() * 8) - 1;
+
     this.leadingZeroCountBase =
       64 - this.unitMagnitude - this.subBucketHalfCountMagnitude - 1;
     this.percentileIterator = new PercentileIterator(this, 1);
@@ -589,16 +593,25 @@ export default class Histogram<T, U> extends AbstractHistogramBase<T, U> {
     // @ts-ignore
     const currentCount = unchecked(this.counts[index]);
     const newCount = currentCount + 1;
-    if (newCount < 0) {
+    if (newCount < currentCount) {
+      const bitSize = <u8>(sizeof<U>() * 8);
+      const overflowAt = (<u64>currentCount + 1); 
       throw new Error(
-        newCount.toString() + " would overflow short integer count"
+        overflowAt.toString() + " would overflow " + bitSize.toString() + "bits integer count"
       );
     }
     // @ts-ignore
     unchecked((this.counts[index] = newCount));
   }
 
   setCountAtIndex(index: i32, value: u64): void {
+    // @ts-ignore
+    if ((<u64>value) as number > this.maxBucketSize) {
+      const bitSize = <u8>(sizeof<U>() * 8);
+      throw new Error(
+        value.toString() + " would overflow " + bitSize.toString() + "bits integer count"
+      );
+    }
     // @ts-ignore
     unchecked((this.counts[index] = <U>value));
   }
@@ -607,8 +620,12 @@ export default class Histogram<T, U> extends AbstractHistogramBase<T, U> {
     // @ts-ignore
     const currentCount = unchecked(this.counts[index]);
     const newCount = currentCount + value;
-    if (newCount < 0) {
-      throw newCount + " would overflow short integer count";
+    const u64NewCount = (<u64>currentCount + value) as number;
+    if (u64NewCount > this.maxBucketSize) {
+      const bitSize = <u8>(sizeof<U>() * 8);
+      throw new Error(
+        u64NewCount.toString() + " would overflow " + bitSize.toString() + "bits integer count"
+      );
     }
     // @ts-ignore
     unchecked((this.counts[index] = <U>newCount));

assembly/__tests__/Histogram.spec.ts

@@ -317,12 +317,12 @@ describe("Histogram add & subtract", () => {
     // given
     const histogram = buildHistogram();
     const histogram2 = buildHistogram();
-    histogram.recordCountAtValue(2, 100);
-    histogram2.recordCountAtValue(1, 100);
-    histogram.recordCountAtValue(2, 200);
-    histogram2.recordCountAtValue(1, 200);
-    histogram.recordCountAtValue(2, 300);
-    histogram2.recordCountAtValue(1, 300);
+    histogram.recordCountAtValue(2, 10);
+    histogram2.recordCountAtValue(1, 10);
+    histogram.recordCountAtValue(2, 20);
+    histogram2.recordCountAtValue(1, 20);
+    histogram.recordCountAtValue(2, 30);
+    histogram2.recordCountAtValue(1, 30);
     const outputBefore = histogram.outputPercentileDistribution();
     // when
     histogram.add<Storage<Uint8Array, u8>, u8>(histogram2);

assembly/encoding.ts

@@ -111,9 +111,9 @@ function fillCountsArrayFromSourceBuffer<T, U>(
   const endPosition = sourceBuffer.position + lengthInBytes;
   while (sourceBuffer.position < endPosition) {
     let zerosCount: i32 = 0;
-    let count = <i32>ZigZagEncoding.decode(sourceBuffer);
+    let count = <i64>ZigZagEncoding.decode(sourceBuffer);
     if (count < 0) {
-      zerosCount = -count;
+      zerosCount = <i32>-count;
       dstIndex += zerosCount; // No need to set zeros in array. Just skip them.
     } else {
       self.setCountAtIndex(dstIndex++, count);

src/Histogram.spec.ts

@@ -3,8 +3,8 @@ import JsHistogram from "./JsHistogram";
 import { NO_TAG } from "./Histogram";
 import Int32Histogram from "./Int32Histogram";
 import { initWebAssembly, WasmHistogram, initWebAssemblySync } from "./wasm";
-import Int8Histogram from "./Int8Histogram";
 import Histogram from "./Histogram";
+import { decodeFromCompressedBase64 } from './encoding';
 
 class HistogramForTests extends JsHistogram {
   //constructor() {}
@@ -557,3 +557,54 @@ describe("Histogram clearing support", () => {
     expect(histogram.mean).toBe(histogram2.mean);
   });
 });
+
+describe("WASM Histogram bucket size overflow", () => {
+  beforeAll(initWebAssembly);
+  [8 as const, 16 as const].forEach(
+    (bitBucketSize) => {
+      const maxBucketSize = (2 ** bitBucketSize) - 1;
+      it(`should fail when recording more than ${maxBucketSize} times the same value`, () => {
+        //given
+        const wasmHistogram = build({ useWebAssembly: true, bitBucketSize });
+
+        //when //then
+        try {
+          let i = 0;
+          for (i; i <= maxBucketSize; i++) {
+            wasmHistogram.recordValue(1); 
+          }
+          fail(`should have failed due to ${bitBucketSize}bits integer overflow (bucket size : ${i})`);
+        } catch(e) {
+          //ok
+        } finally {
+          wasmHistogram.destroy();
+        }
+      });
+
+      it(`should fail when adding two histograms when the same bucket count addition is greater than ${bitBucketSize}bits max integer value`, () => {
+        //given
+        const histogram1 = build({ useWebAssembly: true, bitBucketSize });
+        histogram1.recordValueWithCount(1, maxBucketSize);
+        const histogram2 = build({ useWebAssembly: true, bitBucketSize });
+        histogram2.recordValueWithCount(1, maxBucketSize);
+
+        //when //then
+        expect(() => histogram2.add(histogram1)).toThrow();
+
+        histogram1.destroy();
+        histogram2.destroy();
+      });
+    });
+
+    it("should fail when decoding an Int32 histogram with one bucket count greater than 16bits in a smaller histogram", () => {
+      //given
+      const int32Histogram = build({ useWebAssembly: true, bitBucketSize: 32 }) as WasmHistogram;
+      int32Histogram.recordValueWithCount(1, 2**32 - 1);
+      const encodedInt32Histogram = int32Histogram.encodeIntoCompressedBase64();
+
+      //when //then
+      expect(() => decodeFromCompressedBase64(encodedInt32Histogram, 16, true)).toThrow();
+      int32Histogram.destroy();
+    });
+});
+

src/TypedArrayHistogram.spec.ts

@@ -2,6 +2,7 @@ import Int8Histogram from "./Int8Histogram";
 import Int16Histogram from "./Int16Histogram";
 import Int32Histogram from "./Int32Histogram";
 import Float64Histogram from "./Float64Histogram";
+import { decodeFromCompressedBase64 } from "./encoding";
 
 [Int8Histogram, Int16Histogram, Int32Histogram, Float64Histogram].forEach(
   (Histogram) => {
@@ -67,3 +68,47 @@ import Float64Histogram from "./Float64Histogram";
     });
   }
 );
+
+describe("Histogram bucket size overflow", () => {
+  [Int8Histogram, Int16Histogram].forEach(
+    (Histogram) => {
+      const maxBucketSize = (new Histogram(1, Number.MAX_SAFE_INTEGER, 3)).maxBucketSize;
+      const bitBucketSize = (new Histogram(1, Number.MAX_SAFE_INTEGER, 3))._counts.BYTES_PER_ELEMENT * 8;
+      it(`should fail when recording more than ${maxBucketSize} times the same value for a ${bitBucketSize}bits histogram`, () => {
+        //given;
+        const histogram = new Histogram(1, Number.MAX_SAFE_INTEGER, 3);
+
+        //when //then
+        try {
+          let i = 0;
+          for (i; i <= histogram.maxBucketSize; i++) {
+            histogram.recordValue(1); 
+          }
+          fail(`should have failed due to ${bitBucketSize}bits integer overflow (bucket size: ${i})`);
+        } catch (e) {
+          //ok
+          expect(histogram.getCountAtIndex(1)).toBe(maxBucketSize);
+        }
+      });
+      it(`should fail when adding two histograms when the same bucket count addition is greater than ${bitBucketSize}bits max integer value`, () => {
+        //given
+        const histogram1 = new Histogram(1, Number.MAX_SAFE_INTEGER, 3);
+        histogram1.recordValueWithCount(1, maxBucketSize);
+        const histogram2 = new Histogram(1, Number.MAX_SAFE_INTEGER, 3);
+        histogram2.recordValueWithCount(1, maxBucketSize);
+
+        //when //then
+        expect(() => histogram1.add(histogram2)).toThrow();
+      });
+    });
+    it("should fail when decoding an Int32 histogram with one bucket couunt greater than 16bits", () => {
+      //given
+      const int32Histogram = new Int32Histogram(1, Number.MAX_SAFE_INTEGER, 3);
+      int32Histogram.recordValueWithCount(1, 2**32 - 1);
+      const encodedInt32Histogram = int32Histogram.encodeIntoCompressedBase64();
+
+      //when //then
+      expect(() => decodeFromCompressedBase64(encodedInt32Histogram, 16)).toThrow();
+
+    });
+});

src/TypedArrayHistogram.ts

@@ -18,6 +18,8 @@ class TypedArrayHistogram extends JsHistogram {
   _counts: TypedArray;
   _totalCount: number;
 
+  maxBucketSize: number;
+
   constructor(
     private arrayCtr: new (size: number) => TypedArray,
     lowestDiscernibleValue: number,
@@ -31,6 +33,7 @@ class TypedArrayHistogram extends JsHistogram {
     );
     this._totalCount = 0;
     this._counts = new arrayCtr(this.countsArrayLength);
+    this.maxBucketSize = 2**(this._counts.BYTES_PER_ELEMENT * 8) - 1;
   }
 
   clearCounts() {
@@ -40,27 +43,24 @@ class TypedArrayHistogram extends JsHistogram {
   incrementCountAtIndex(index: number) {
     const currentCount = this._counts[index];
     const newCount = currentCount + 1;
-    if (newCount < 0) {
-      throw newCount + " would overflow short integer count";
+    if (newCount > this.maxBucketSize) {
+      throw newCount + " would overflow " + this._counts.BYTES_PER_ELEMENT * 8 + "bits integer count";
     }
     this._counts[index] = newCount;
   }
 
   addToCountAtIndex(index: number, value: number) {
     const currentCount = this._counts[index];
     const newCount = currentCount + value;
-    if (
-      newCount < Number.MIN_SAFE_INTEGER ||
-      newCount > Number.MAX_SAFE_INTEGER
-    ) {
-      throw newCount + " would overflow integer count";
+    if (newCount > this.maxBucketSize) {
+      throw newCount + " would overflow " + this._counts.BYTES_PER_ELEMENT * 8 + "bits integer count";
     }
     this._counts[index] = newCount;
   }
 
   setCountAtIndex(index: number, value: number) {
-    if (value < Number.MIN_SAFE_INTEGER || value > Number.MAX_SAFE_INTEGER) {
-      throw value + " would overflow integer count";
+    if (value > this.maxBucketSize) {
+      throw value + " would overflow " + this._counts.BYTES_PER_ELEMENT * 8 + "bits integer count";
     }
     this._counts[index] = value;
   }