Fix overflow when decoding WASM histogram

HdrHistogram · Nov 19, 2020 · 74ebc8b · 74ebc8b
1 parent cc23b3a
commit 74ebc8b
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 34 deletions.
diff --git a/assembly/Histogram.ts b/assembly/Histogram.ts
@@ -605,6 +605,13 @@ export default class Histogram<T, U> extends AbstractHistogramBase<T, U> {
   }
 
   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));
   }

diff --git a/assembly/__tests__/Histogram.spec.ts b/assembly/__tests__/Histogram.spec.ts
@@ -326,9 +326,9 @@ describe("Histogram add & subtract", () => {
     const outputBefore = histogram.outputPercentileDistribution();
     // when
     histogram.add<Storage<Uint8Array, u8>, u8>(histogram2);
-    //histogram.subtract<Storage<Uint8Array, u8>, u8>(histogram2);
+    histogram.subtract<Storage<Uint8Array, u8>, u8>(histogram2);
     // then
-   //expect(histogram.outputPercentileDistribution()).toBe(outputBefore);
+    expect(histogram.outputPercentileDistribution()).toBe(outputBefore);
   });
 
   it("should be equal when another histogram of lower precision is added then subtracted", () => {

diff --git a/assembly/encoding.ts b/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);

diff --git a/src/Histogram.spec.ts b/src/Histogram.spec.ts
@@ -559,16 +559,13 @@ describe("Histogram clearing support", () => {
 });
 
 describe("WASM Histogram bucket size overflow", () => {
-  let wasmHistogram: Histogram;
   beforeAll(initWebAssembly);
-  afterEach(() => wasmHistogram.destroy());
-
   [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
-        wasmHistogram = build({ useWebAssembly: true, bitBucketSize });
+        const wasmHistogram = build({ useWebAssembly: true, bitBucketSize });
 
         //when //then
         try {
@@ -577,42 +574,37 @@ describe("WASM Histogram bucket size overflow", () => {
             wasmHistogram.recordValue(1); 
           }
           fail(`should have failed due to ${bitBucketSize}bits integer overflow (bucket size : ${i})`);
-        } catch (e) {
+        } catch(e) {
           //ok
+        } finally {
+          wasmHistogram.destroy();
         }
       });
 
-       //Cannot use a destroyed histogram error
-    it.skip(`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
-      try {
-        histogram2.add(histogram1);
-        fail(`should have failed due to ${bitBucketSize}bits integer overflow`);
-      } catch (e) {
-        //ok
-      }
-    });
+      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();
+      });
     });
-    //Cannot use a destroyed histogram error
-    it.skip("should fail when decoding an Int32 histogram with one bucket count greater than 16bits", () => {
+
+    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
-      try {
-        const encodedInt32Histogram = int32Histogram.encodeIntoCompressedBase64();
-        decodeFromCompressedBase64(encodedInt32Histogram, 16, true);
-        fail(`should have failed due to 16bits integer overflow`);
-      } catch (e) {
-        //ok
-      }
+      expect(() => decodeFromCompressedBase64(encodedInt32Histogram, 16, true)).toThrow();
+      int32Histogram.destroy();
     });
 });