fix(tokens): Hash tokens in tokens module to resist ND2DB-style tim…

…ing attack
rivet-gg · May 1, 2024 · b2cefc7 · b2cefc7
1 parent 1317b2b
commit b2cefc7
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 18 deletions.
diff --git a/...rations/20240307005945_init/migration.sql → ...rations/20240501200641_init/migration.sql b/...rations/20240307005945_init/migration.sql → ...rations/20240501200641_init/migration.sql
@@ -1,7 +1,7 @@
 -- CreateTable
 CREATE TABLE "Token" (
     "id" UUID NOT NULL,
-    "token" TEXT NOT NULL,
+    "tokenHash" TEXT NOT NULL,
     "type" TEXT NOT NULL,
     "meta" JSONB NOT NULL,
     "trace" JSONB NOT NULL,
@@ -13,4 +13,4 @@ CREATE TABLE "Token" (
 );
 
 -- CreateIndex
-CREATE UNIQUE INDEX "Token_token_key" ON "Token"("token");
+CREATE UNIQUE INDEX "Token_tokenHash_key" ON "Token"("tokenHash");
diff --git a/modules/tokens/db/schema.prisma b/modules/tokens/db/schema.prisma
@@ -5,7 +5,7 @@ datasource db {
 
 model Token {
   id        String    @id @default(uuid()) @db.Uuid
-  token     String    @unique
+  tokenHash String    @unique
   type      String
   meta      Json
   trace     Json

diff --git a/modules/tokens/scripts/create.ts b/modules/tokens/scripts/create.ts
@@ -1,6 +1,5 @@
 import { ScriptContext } from "../_gen/scripts/create.ts";
-import { TokenWithSecret } from "../utils/types.ts";
-import { tokenFromRow } from "../utils/types.ts";
+import { TokenWithSecret, tokenFromRow, hash } from "../utils/types.ts";
 
 export interface Request {
 	type: string;
@@ -17,10 +16,9 @@ export async function run(
 	req: Request,
 ): Promise<Response> {
 	const tokenStr = generateToken(req.type);
-
 	const token = await ctx.db.token.create({
 		data: {
-			token: tokenStr,
+			tokenHash: await hash(tokenStr),
 			type: req.type,
 			meta: req.meta,
 			trace: ctx.trace,
@@ -29,7 +27,7 @@ export async function run(
 	});
 
 	return {
-		token: tokenFromRow(token),
+		token: tokenFromRow(token, () => tokenStr),
 	};
 }
 

diff --git a/modules/tokens/scripts/extend.ts b/modules/tokens/scripts/extend.ts
@@ -1,6 +1,5 @@
 import { ScriptContext } from "../_gen/scripts/extend.ts";
-import { TokenWithSecret } from "../types/common.ts";
-import { tokenFromRow } from "../types/common.ts";
+import { TokenWithSecret, tokenFromRow } from "../utils/types.ts";
 
 export interface Request {
     token: string;
@@ -32,6 +31,6 @@ export async function run(
 
     // Return the updated token
 	return {
-		token: tokenFromRow(newToken),
+		token: tokenFromRow(newToken, () => req.token),
 	};
 }
diff --git a/modules/tokens/scripts/get.ts b/modules/tokens/scripts/get.ts
@@ -1,6 +1,5 @@
 import { ScriptContext } from "../_gen/scripts/get.ts";
-import { Token } from "../utils/types.ts";
-import { tokenFromRow } from "../utils/types.ts";
+import { Token, tokenFromRow } from "../utils/types.ts";
 
 export interface Request {
 	tokenIds: string[];
@@ -23,7 +22,7 @@ export async function run(
 		},
 	});
 
-	const tokens = rows.map(tokenFromRow);
+	const tokens = rows.map(row => tokenFromRow(row, () => ""));
 
 	return { tokens };
 }
diff --git a/modules/tokens/scripts/get_by_token.ts b/modules/tokens/scripts/get_by_token.ts
@@ -1,5 +1,5 @@
 import { ScriptContext } from "../_gen/scripts/get_by_token.ts";
-import { Token, tokenFromRow } from "../utils/types.ts";
+import { Token, tokenFromRow, hash } from "../utils/types.ts";
 
 export interface Request {
 	tokens: string[];
@@ -13,18 +13,21 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
+	const hashed = await Promise.all(req.tokens.map(hash));
 	const rows = await ctx.db.token.findMany({
 		where: {
-			token: {
-				in: req.tokens,
+			tokenHash: {
+				in: hashed,
 			},
 		},
 		orderBy: {
 			createdAt: "desc",
 		},
 	});
 
-	const tokens = rows.map(tokenFromRow);
+	const hashMap = Object.fromEntries(req.tokens.map((token, i) => [hashed[i], token]));
+
+	const tokens = rows.map(row => tokenFromRow(row, h => hashMap[h]));
 
 	return { tokens };
 }
diff --git a/modules/tokens/utils/types.ts b/modules/tokens/utils/types.ts
@@ -15,11 +15,22 @@ export interface TokenWithSecret extends Token {
 
 export function tokenFromRow(
 	row: prisma.Prisma.TokenGetPayload<any>,
+	hashToOrig: (hash: string) => string,
 ): TokenWithSecret {
 	return {
 		...row,
 		createdAt: row.createdAt.toISOString(),
 		expireAt: row.expireAt?.toISOString() ?? null,
 		revokedAt: row.revokedAt?.toISOString() ?? null,
+		token: hashToOrig(row.tokenHash),
 	};
 }
+
+export async function hash(token: string): Promise<string> {
+	const encoder = new TextEncoder();
+	const data = encoder.encode(token);
+	const hash = await crypto.subtle.digest("SHA-256", data);
+	const digest = Array.from(new Uint8Array(hash));
+	const strDigest = digest.map(b => b.toString(16).padStart(2, "0")).join("");
+	return strDigest;
+}