fix: shift with BigInt instead of JS number (#7703)

Co-authored-by: Begoña Álvarez de la Cruz <balvarez@boxfish.studio>
iotaledger · Nov 16, 2023 · a56001b · a56001b
1 parent c9559f0
commit a56001b
Showing 1 changed file with 20 additions and 106 deletions.
diff --git a/packages/shared/lib/core/layer-2/classes/special-stream.class.ts b/packages/shared/lib/core/layer-2/classes/special-stream.class.ts
@@ -46,127 +46,41 @@ export class ReadSpecialStream extends ReadStream {
     }
 }
 
-function shiftRight(s: bigint, n: bigint): bigint {
-    const result = BigInteger(s).divide(BigInteger(2).pow(BigInteger(n)))
-    return BigInt(result.toString())
+function shiftRight(n: bigint, shift: bigint): bigint {
+    return n >> shift
 }
-// from https://github.com/iotaledger/wasp/blob/12845adea4fc097813a30a061853af4a43407d3c/packages/util/rwutil/convert.go#L113
+
 function size64Encode(n: bigint): Buffer {
-    if (n < BigInt(0x80)) {
-        return Buffer.from([Number(n)])
-    } else if (n < BigInt(0x4000)) {
-        return Buffer.from([Number(n | BigInt(0x80)), Number(shiftRight(n, BigInt(7)))])
-    } else if (n < BigInt(0x20_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14))),
-        ])
-    } else if (n < BigInt(0x1000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21))),
-        ])
-    } else if (n < BigInt(0x8_0000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28))),
-        ])
-    } else if (n < BigInt(0x400_0000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(35))),
-        ])
-    } else if (n < BigInt(0x2_0000_0000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(35)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(42))),
-        ])
-    } else if (n < BigInt(0x100_0000_0000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(35)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(42)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(49))),
-        ])
-    } else if (n < BigInt(0x8000_0000_0000_0000)) {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(35)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(42)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(49)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(56))),
-        ])
-    } else {
-        return Buffer.from([
-            Number(n | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(7)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(14)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(21)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(28)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(35)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(42)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(49)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(56)) | BigInt(0x80)),
-            Number(shiftRight(n, BigInt(63))),
-        ])
+    const parts: number[] = []
+    while (n >= 0x80) {
+        parts.push(Number((n & BigInt(0x7f)) | BigInt(0x80)))
+        n = shiftRight(n, BigInt(7))
     }
+    parts.push(Number(n)) // Last byte, without the continuation bit
+
+    return Buffer.from(parts)
 }
 
-// Adapted from WASP golang implementation https://github.com/iotaledger/wasp/blob/7f880a7983d24d0dcd225e994d67b29741b318bc/packages/util/rwutil/convert.go#L76
 function size64Decode(readByte: () => number): [bigint, null | Error] {
     let byte = readByte()
-
-    if (!byte) {
-        return [BigInt(0), new Error('no more bytes')]
-    }
-
     if (byte < 0x80) {
         return [BigInt(byte), null]
     }
 
-    let value = BigInt(byte) & BigInt(0x7f)
-
-    for (let shift = 7; shift < 63; shift += 7) {
+    let value = BigInt(byte & 0x7f)
+    let shift = 7
+    while (shift < 64) {
         byte = readByte()
-        if (!byte) {
-            return [BigInt(0), new Error('no more bytes')]
+        if (byte === -1) {
+            // Assuming -1 signifies end of data or error in readByte
+            return [BigInt(0), new Error('Unexpected end of data')]
         }
+        value |= BigInt(byte & 0x7f) << BigInt(shift)
         if (byte < 0x80) {
-            return [value | (BigInt(byte) << BigInt(shift)), null]
+            return [value, null]
         }
-        value |= (BigInt(byte) & BigInt(0x7f)) << BigInt(shift)
-    }
-
-    byte = readByte()
-    if (!byte) {
-        return [BigInt(0), new Error('no more bytes')]
-    }
-    if (byte > 0x01) {
-        return [BigInt(0), new Error('size64 overflow')]
+        shift += 7
     }
 
-    return [value | (BigInt(byte) << BigInt(63)), null]
+    return [BigInt(0), new Error('size64 overflow')]
 }