From 58ae6b7fc512f84456dcd4980cfa3ac3517de0b7 Mon Sep 17 00:00:00 2001
From: Richard Moore <github@ricmoo.com>
Date: Fri, 30 Sep 2022 19:56:13 -0400
Subject: [PATCH] tests: fixed locking randomBytes test

---
 src.ts/_tests/test-crypto-algoswap.ts | 26 +++++++++++++++++++-------
 src.ts/crypto/index.ts                |  1 +
 src.ts/crypto/random.ts               |  2 +-
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/src.ts/_tests/test-crypto-algoswap.ts b/src.ts/_tests/test-crypto-algoswap.ts
index de7837a92..c63fefbe9 100644
--- a/src.ts/_tests/test-crypto-algoswap.ts
+++ b/src.ts/_tests/test-crypto-algoswap.ts
@@ -2,16 +2,14 @@ import assert from "assert";
 
 import {
     lock,
-
     computeHmac,
-
     keccak256, ripemd160, sha256, sha512,
-
-    pbkdf2, scrypt, scryptSync
+    pbkdf2, scrypt, scryptSync,
+    randomBytes
 } from "../index.js";
 
 interface Algorithm {
-    (...args: Array<any>): string | Promise<string>;
+    (...args: Array<any>): string | Uint8Array | Promise<string>;
 
     register: (func: any) => void;
     lock: () => void;
@@ -23,6 +21,7 @@ interface TestCase {
     params: Array<any>;
     algorithm: Algorithm;
     hijackTag: string;
+    postCheck?: (value: any) => boolean;
 }
 
 
@@ -85,9 +84,18 @@ describe("test registration", function() {
             hijackTag: 'hijacked computeHmac: ["sha256",{},{}]',
             algorithm: computeHmac
         },
+        {
+            name: "randomBytes",
+            params: [ 32 ],
+            hijackTag: "hijacked randomBytes: [32]",
+            algorithm: randomBytes,
+            postCheck: (value: any) => {
+                return (value instanceof Uint8Array && value.length === 32);
+            }
+        }
     ];
 
-    tests.forEach(({ name, params, hijackTag, algorithm }) => {
+    tests.forEach(({ name, params, hijackTag, algorithm, postCheck }) => {
         it(`swaps in hijacked callback: ${ name }`, async function() {
             const initial = await algorithm(...params);
 
@@ -97,7 +105,11 @@ describe("test registration", function() {
             assert.equal(hijack, hijackTag);
 
             algorithm.register(algorithm._);
-            assert.equal(await algorithm(...params), initial);
+            if (postCheck) {
+                assert.ok(postCheck(await algorithm(...params)));
+            } else {
+                assert.equal(await algorithm(...params), initial);
+            }
         });
     });
 
diff --git a/src.ts/crypto/index.ts b/src.ts/crypto/index.ts
index 2ce7efbf2..cd23d9cd0 100644
--- a/src.ts/crypto/index.ts
+++ b/src.ts/crypto/index.ts
@@ -33,6 +33,7 @@ function lock(): void {
     scryptSync.lock();
     sha256.lock();
     sha512.lock();
+    randomBytes.lock();
 }
 
 export { lock };
diff --git a/src.ts/crypto/random.ts b/src.ts/crypto/random.ts
index 72fb619d9..90daff63d 100644
--- a/src.ts/crypto/random.ts
+++ b/src.ts/crypto/random.ts
@@ -15,7 +15,7 @@ export function randomBytes(length: number): Uint8Array {
 randomBytes._ = _randomBytes;
 randomBytes.lock = function(): void { locked = true; }
 randomBytes.register = function(func: (length: number) => Uint8Array) {
-    if (locked) { throw new Error("random is locked"); }
+    if (locked) { throw new Error("randomBytes is locked"); }
     __randomBytes = func;
 }
 Object.freeze(randomBytes);