Use a random impl with rejection sampling

Author: dbkr

spec/unit/randomstring.spec.ts

@@ -24,6 +24,10 @@ import {
 } from "../../src/randomstring";
 
 describe("Random strings", () => {
+    afterEach(() => {
+        jest.restoreAllMocks();
+    });
+
     it.each([8, 16, 32])("secureRandomBase64 generates %i valid base64 bytes", (n: number) => {
         const randb641 = secureRandomBase64Url(n);
         const randb642 = secureRandomBase64Url(n);
@@ -70,4 +74,41 @@ describe("Random strings", () => {
             expect(rand1.toUpperCase()).toEqual(rand1);
         },
     );
+
+    it("throws if given character set less than 2 characters", () => {
+        expect(() => secureRandomStringFrom(8, "a")).toThrow();
+    });
+
+    it("throws if given character set more than 256 characters", () => {
+        const charSet = Array.from({ length: 257 }, (_, i) => "a").join("");
+
+        expect(() => secureRandomStringFrom(8, charSet)).toThrow();
+    });
+
+    it("throws if given length less than 1", () => {
+        expect(() => secureRandomStringFrom(0, "abc")).toThrow();
+    });
+
+    it("throws if given length more than 32768", () => {
+        expect(() => secureRandomStringFrom(32769, "abc")).toThrow();
+    });
+
+    it("asks for more entropy if given entropy is unusable", () => {
+        // This is testing the internal implementation details of the function rather
+        // than strictly the public API. The intention is to have some assertion that
+        // the rejection sampling to make the distribution even over all possible characters
+        // is doing what it's supposed to do.
+
+        // mock once to fill with 255 the first time: 255 should be unusable because
+        // we give 10 possible characters below and 256 is not evenly divisible by 10, so
+        // this should force it to call for more entropy.
+        jest.spyOn(globalThis.crypto, "getRandomValues").mockImplementationOnce((arr) => {
+            if (arr === null) throw new Error("Buffer is null");
+            new Uint8Array(arr.buffer).fill(255);
+            return arr;
+        });
+
+        secureRandomStringFrom(8, "0123456789");
+        expect(globalThis.crypto.getRandomValues).toHaveBeenCalledTimes(2);
+    });
 });

src/randomstring.ts

@@ -40,17 +40,48 @@ export function secureRandomString(len: number): string {
 
 /**
  * Generate a cryptographically secure random string using characters given
- * @param len The length of the string to generate
- * @param chars The characters to use in the random string.
+ * @param len The length of the string to generate (must be positive and less than 32768)
+ * @param chars The characters to use in the random string (between 2 and 256 characters long).
  * @returns Random string of characters of length `len`
  */
 export function secureRandomStringFrom(len: number, chars: string): string {
-    const positions = new Uint32Array(chars.length);
-    let ret = "";
-    crypto.getRandomValues(positions);
-    for (let i = 0; i < len; i++) {
-        const currentCharPlace = positions[i % chars.length] % chars.length;
-        ret += chars[currentCharPlace];
+    // This is intended for latin strings so 256 possibilities should be more than enough and
+    // means we can use random bytes, minimising the amount of entropy we need to ask for.
+    if (chars.length < 2 || chars.length > 256) {
+        throw new Error("Character set must be between 2 and 256 characters long");
+    }
+
+    if (len < 1 || len > 32768) {
+        throw new Error("Requested random string length must be between 1 and 32768");
+    }
+
+    // We'll generate random unsigned bytes, so get the largest number less than 256 that is a multiple
+    // of the length of the character set: We'll need to discard any random values that are larger than
+    // this as we can't possibly map them onto the character set while keeping each character equally
+    // likely to be chosen (minus 1 to convert to indices in a string). (Essentially, we're using a d8
+    // to choose between 7 possibilities and re-rolling on an 8, keeping all 7 outcomes equally likely.)
+    const maxRandValue = Math.floor(255 / chars.length) * chars.length - 1;
+
+    // Grab 30% more entropy than we need. This should be enough that we can discard the values that are
+    // too high without having to go back and grab more unless we're super unlucky.
+    const entropyBuffer = new Uint8Array(Math.floor(len * 1.3));
+    // Mark all of this buffer as used to start with (we haven't populated it with entropy yet) so it will
+    // be filled on the first iteration.
+    let entropyBufferPos = entropyBuffer.length;
+
+    const result = [];
+    while (result.length < len) {
+        if (entropyBufferPos === entropyBuffer.length) {
+            globalThis.crypto.getRandomValues(entropyBuffer);
+            entropyBufferPos = 0;
+        }
+
+        const randomByte = entropyBuffer[entropyBufferPos++];
+
+        if (randomByte < maxRandValue) {
+            result.push(chars[randomByte % chars.length]);
+        }
     }
-    return ret;
+
+    return result.join("");
 }