Skip to content

Commit

Permalink
Move to xorshiro for rand generation (#190)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jackkleeman authored Nov 21, 2023
1 parent 0dea824 commit fdaae83
Show file tree
Hide file tree
Showing 6 changed files with 143 additions and 132 deletions.
4 changes: 2 additions & 2 deletions src/restate_context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

// ----------------------------------------------------------------------------
Expand Down
2 changes: 1 addition & 1 deletion src/restate_context_impl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
7 changes: 5 additions & 2 deletions src/server/restate_lambda_handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
93 changes: 58 additions & 35 deletions src/utils/rand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand All @@ -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]] +
Expand Down
12 changes: 8 additions & 4 deletions test/protoutils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 || []),
Expand Down Expand Up @@ -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[] {
Expand Down
Loading

0 comments on commit fdaae83

Please sign in to comment.