Move to xorshiro for rand generation (#190)

By using the full 256 bits from the sha256 hash, we get extremely good state collision resistance; it is now pretty much impossible for two different invocation ids to produce the same state. As a result, we can be a lot more confident in the quality of our pseudorandom numbers.

Author: jackkleeman

src/restate_context.ts

@@ -191,13 +191,13 @@ export interface Rand {
    * Equivalent of JS `Math.random()` but deterministic; seeded by the invocation ID of the current invocation,
    * each call will return a new pseudorandom float within the range [0,1)
    */
-  random(): number
+  random(): number;
 
   /**
    * Using the same random source and seed as random(), produce a UUID version 4 string. This is inherently predictable
    * based on the invocation ID and should not be used in cryptographic contexts
    */
-  uuidv4(): string
+  uuidv4(): string;
 }
 
 // ----------------------------------------------------------------------------

src/restate_context_impl.ts

@@ -59,7 +59,7 @@ import { rlog } from "./utils/logger";
 import { Client, SendClient } from "./types/router";
 import { RpcRequest, RpcResponse } from "./generated/proto/dynrpc";
 import { requestFromArgs } from "./utils/assumpsions";
-import {RandImpl} from "./utils/rand";
+import { RandImpl } from "./utils/rand";
 
 export enum CallContexType {
   None,

src/server/restate_lambda_handler.ts

@@ -28,7 +28,7 @@ import { Message } from "../types/types";
 import { StateMachine } from "../state_machine";
 import { ensureError } from "../types/errors";
 import { KeyedRouter, UnKeyedRouter } from "../public_api";
-import {OUTPUT_STREAM_ENTRY_MESSAGE_TYPE} from "../types/protocol";
+import { OUTPUT_STREAM_ENTRY_MESSAGE_TYPE } from "../types/protocol";
 
 /**
  * Creates an Restate entrypoint for services deployed on AWS Lambda and invoked
@@ -224,7 +224,10 @@ export class LambdaRestateServer extends BaseRestateServer {
       let decodedEntries: Message[] | null = decodeLambdaBody(event.body);
       const journalBuilder = new InvocationBuilder(method);
       decodedEntries.forEach((e: Message) => journalBuilder.handleMessage(e));
-      const alreadyCompleted = decodedEntries.find((e: Message) => e.messageType === OUTPUT_STREAM_ENTRY_MESSAGE_TYPE) !== undefined
+      const alreadyCompleted =
+        decodedEntries.find(
+          (e: Message) => e.messageType === OUTPUT_STREAM_ENTRY_MESSAGE_TYPE
+        ) !== undefined;
       decodedEntries = null;
 
       // set up and invoke the state machine

src/utils/rand.ts

@@ -12,72 +12,95 @@
 //! Some parts copied from https://github.com/uuidjs/uuid/blob/main/src/stringify.js
 //! License MIT
 
-import {Rand} from "../restate_context";
-import {ErrorCodes, TerminalError} from "../types/errors";
-import {CallContexType, RestateGrpcContextImpl} from "../restate_context_impl";
-import {createHash} from "crypto";
+import { Rand } from "../restate_context";
+import { ErrorCodes, TerminalError } from "../types/errors";
+import {
+  CallContexType,
+  RestateGrpcContextImpl,
+} from "../restate_context_impl";
+import { createHash } from "crypto";
 
 export class RandImpl implements Rand {
-  private randstate64: bigint;
+  private randstate256: [bigint, bigint, bigint, bigint];
 
-  constructor(id: Buffer | bigint) {
-    if (typeof id == "bigint") {
-      this.randstate64 = id
-    } else {
+  constructor(id: Buffer | [bigint, bigint, bigint, bigint]) {
+    if (id instanceof Buffer) {
       // hash the invocation ID, which is known to contain 74 bits of entropy
-      const hash = createHash('sha256')
-        .update(id)
-        .digest();
-
-      // seed using first 64 bits of the hash
-      this.randstate64 = hash.readBigUInt64LE(0);
+      const hash = createHash("sha256").update(id).digest();
+
+      this.randstate256 = [
+        hash.readBigUInt64LE(0),
+        hash.readBigUInt64LE(8),
+        hash.readBigUInt64LE(16),
+        hash.readBigUInt64LE(24),
+      ];
+    } else {
+      this.randstate256 = id;
     }
   }
 
-  static U64_MASK = ((1n << 64n) - 1n)
+  static U64_MASK = (1n << 64n) - 1n;
 
-  // splitmix64
-  // https://prng.di.unimi.it/splitmix64.c - public domain
+  // xoshiro256++
+  // https://prng.di.unimi.it/xoshiro256plusplus.c - public domain
   u64(): bigint {
-    this.randstate64 = (this.randstate64 + 0x9e3779b97f4a7c15n) & RandImpl.U64_MASK;
-    let next: bigint = this.randstate64;
-    next = ((next ^ (next >> 30n)) * 0xbf58476d1ce4e5b9n) & RandImpl.U64_MASK;
-    next = ((next ^ (next >> 27n)) * 0x94d049bb133111ebn) & RandImpl.U64_MASK;
-    next = next ^ (next >> 31n);
-    return next
+    const result: bigint =
+      (RandImpl.rotl(
+        (this.randstate256[0] + this.randstate256[3]) & RandImpl.U64_MASK,
+        23n
+      ) +
+        this.randstate256[0]) &
+      RandImpl.U64_MASK;
+
+    const t: bigint = (this.randstate256[1] << 17n) & RandImpl.U64_MASK;
+
+    this.randstate256[2] ^= this.randstate256[0];
+    this.randstate256[3] ^= this.randstate256[1];
+    this.randstate256[1] ^= this.randstate256[2];
+    this.randstate256[0] ^= this.randstate256[3];
+
+    this.randstate256[2] ^= t;
+
+    this.randstate256[3] = RandImpl.rotl(this.randstate256[3], 45n);
+
+    return result;
   }
 
-  static U53_MASK = ((1n << 53n) - 1n)
+  static rotl(x: bigint, k: bigint): bigint {
+    return ((x << k) & RandImpl.U64_MASK) | (x >> (64n - k));
+  }
 
   checkContext() {
     const context = RestateGrpcContextImpl.callContext.getStore();
     if (context && context.type === CallContexType.SideEffect) {
       throw new TerminalError(
         `You may not call methods on Rand from within a side effect.`,
-        {errorCode: ErrorCodes.INTERNAL}
+        { errorCode: ErrorCodes.INTERNAL }
       );
     }
   }
 
+  static U53_MASK = (1n << 53n) - 1n;
+
   public random(): number {
-    this.checkContext()
+    this.checkContext();
 
     // first generate a uint in range [0,2^53), which can be mapped 1:1 to a float64 in [0,1)
-    const u53 = this.u64() & RandImpl.U53_MASK
+    const u53 = this.u64() & RandImpl.U53_MASK;
     // then divide by 2^53, which will simply update the exponent
-    return Number(u53) / 2 ** 53
+    return Number(u53) / 2 ** 53;
   }
 
   public uuidv4(): string {
-    this.checkContext()
+    this.checkContext();
 
     const buf = Buffer.alloc(16);
     buf.writeBigUInt64LE(this.u64(), 0);
     buf.writeBigUInt64LE(this.u64(), 8);
     // Per 4.4, set bits for version and `clock_seq_hi_and_reserved`
     buf[6] = (buf[6] & 0x0f) | 0x40;
     buf[8] = (buf[8] & 0x3f) | 0x80;
-    return uuidStringify(buf)
+    return uuidStringify(buf);
   }
 }
 
@@ -102,16 +125,16 @@ function uuidStringify(arr: Buffer, offset = 0) {
     byteToHex[arr[offset + 1]] +
     byteToHex[arr[offset + 2]] +
     byteToHex[arr[offset + 3]] +
-    '-' +
+    "-" +
     byteToHex[arr[offset + 4]] +
     byteToHex[arr[offset + 5]] +
-    '-' +
+    "-" +
     byteToHex[arr[offset + 6]] +
     byteToHex[arr[offset + 7]] +
-    '-' +
+    "-" +
     byteToHex[arr[offset + 8]] +
     byteToHex[arr[offset + 9]] +
-    '-' +
+    "-" +
     byteToHex[arr[offset + 10]] +
     byteToHex[arr[offset + 11]] +
     byteToHex[arr[offset + 12]] +

test/protoutils.ts

@@ -65,7 +65,10 @@ export function startMessage(
   return new Message(
     START_MESSAGE_TYPE,
     StartMessage.create({
-      id: Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"),
+      id: Buffer.from(
+        "f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2",
+        "hex"
+      ),
       debugId: "8xHx_cuYY_AAYvTQA7NfWm1RyBOd2IYsg",
       knownEntries: knownEntries, // only used for the Lambda case. For bidi streaming, this will be imputed by the testdriver
       stateMap: toStateEntries(state || []),
@@ -441,9 +444,10 @@ export function getAwakeableId(entryIndex: number): string {
   const encodedEntryIndex = Buffer.alloc(4 /* Size of u32 */);
   encodedEntryIndex.writeUInt32BE(entryIndex);
 
-  return Buffer.concat([Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"), encodedEntryIndex]).toString(
-    "base64url"
-  );
+  return Buffer.concat([
+    Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"),
+    encodedEntryIndex,
+  ]).toString("base64url");
 }
 
 export function keyVal(key: string, value: any): Buffer[] {