feat: impl rate limit

rivet-gg · Mar 13, 2024 · 64eac6b · 64eac6b
1 parent f7560ea
commit 64eac6b
Show file tree

Hide file tree

Showing 32 changed files with 824 additions and 76 deletions.
diff --git a/.github/workflows/test-all.yml b/.github/workflows/test-all.yml
@@ -3,7 +3,7 @@ on:
   - push
 
 env:
-  CURRENT_WORKING_ENGINE_COMMIT: 920ac277dd0f5378cf8c3cc1ae414340dc6d76a6
+  CURRENT_WORKING_ENGINE_COMMIT: cbdd1e1b364b1906347b3ef541a4a5057651650f
 
 jobs:
   build:
@@ -40,4 +40,4 @@ jobs:
 
       # Run tests on all modules in the registry
       - name: Run Tests for all modules
-        run: cd ./opengb-registry/tests/basic && opengb test
+        run: cd ./opengb-registry/tests/basic && opengb test --strict-schemas --force-deploy-migrations
diff --git a/modules/auth/scripts/auth_email_passwordless.ts b/modules/auth/scripts/auth_email_passwordless.ts
@@ -15,9 +15,9 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	if (!ctx.userConfig.email) throw new RuntimeError("PROVIDER_DISABLED");
+	await ctx.modules.rateLimit.throttlePublic({});
 
-	await ctx.modules.rateLimit.throttle({});
+	if (!ctx.userConfig.email) throw new RuntimeError("PROVIDER_DISABLED");
 
 	// Fetch existing user if session token is provided
 	let userId: string | undefined;

diff --git a/modules/auth/scripts/verify_email_passwordless.ts b/modules/auth/scripts/verify_email_passwordless.ts
@@ -18,7 +18,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({});
+	await ctx.modules.rateLimit.throttlePublic({});
 
 	const code = req.code.toUpperCase();
 

diff --git a/modules/currency/scripts/deposit.ts b/modules/currency/scripts/deposit.ts
@@ -16,7 +16,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 25 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 25 });
 
 	if (req.amount < 0 || !Number.isFinite(req.amount)) {
 		throw new RuntimeError("INVALID_AMOUNT");

diff --git a/modules/currency/scripts/get_balance.ts b/modules/currency/scripts/get_balance.ts
@@ -14,7 +14,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 25 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 25 });
 
 	return {
 		balance: await getBalance(ctx.db, req.userId),

diff --git a/modules/currency/scripts/get_balance_by_token.ts b/modules/currency/scripts/get_balance_by_token.ts
@@ -13,7 +13,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 25 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 25 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/currency/scripts/set_balance.ts b/modules/currency/scripts/set_balance.ts
@@ -14,7 +14,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 25 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 25 });
 
 	if (req.balance < 0 || !Number.isFinite(req.balance)) {
 		throw new RuntimeError("INVALID_AMOUNT");

diff --git a/modules/currency/scripts/withdraw.ts b/modules/currency/scripts/withdraw.ts
@@ -16,7 +16,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 25 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 25 });
 
 	if (req.amount < 0 || !Number.isFinite(req.amount)) {
 		throw new RuntimeError("INVALID_AMOUNT");

diff --git a/modules/currency/tests/e2e.ts b/modules/currency/tests/e2e.ts
@@ -5,9 +5,11 @@ import { faker } from "https://deno.land/x/[email protected]/mod.ts";
 test(
 	"e2e transaction",
 	async (ctx: TestContext) => {
-		const { user: user, token: token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
+		});
+		const { token } = await ctx.modules.users.createUserToken({
+			userId: user.id,
 		});
 
 		const { updatedBalance } = await ctx.modules.currency.deposit({

diff --git a/modules/currency/tests/errors.ts b/modules/currency/tests/errors.ts
@@ -19,9 +19,8 @@ test(
 test(
 	"withdraw more than balance",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const { updatedBalance: _ } = await ctx.modules.currency.deposit({
@@ -39,9 +38,8 @@ test(
 test(
 	"withdraw negative amount",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -54,9 +52,8 @@ test(
 test(
 	"withdraw Infinity",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -72,9 +69,8 @@ test(
 test(
 	"withdraw NaN",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -87,9 +83,8 @@ test(
 test(
 	"deposit Infinity",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -102,9 +97,8 @@ test(
 test(
 	"deposit NaN",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -117,9 +111,8 @@ test(
 test(
 	"deposit negative amount",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -132,9 +125,8 @@ test(
 test(
 	"set balance to negative",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -147,9 +139,8 @@ test(
 test(
 	"set balance to NaN",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {
@@ -163,9 +154,8 @@ test(
 test(
 	"set balance to infinity",
 	async (ctx: TestContext) => {
-		const { user: user, token: _token } = await ctx.modules.users.register({
+		const { user: user } = await ctx.modules.users.createUser({
 			username: faker.internet.userName(),
-			identity: { guest: {} },
 		});
 
 		const error = await assertRejects(async () => {

diff --git a/modules/friends/scripts/accept_request.ts b/modules/friends/scripts/accept_request.ts
@@ -11,7 +11,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 50 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 50 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/decline_request.ts b/modules/friends/scripts/decline_request.ts
@@ -14,7 +14,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 50 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 50 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/list_friends.ts b/modules/friends/scripts/list_friends.ts
@@ -13,7 +13,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 50 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 50 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/list_incoming_friend_requests.ts b/modules/friends/scripts/list_incoming_friend_requests.ts
@@ -13,7 +13,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 50 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 50 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/list_outgoing_friend_requests.ts b/modules/friends/scripts/list_outgoing_friend_requests.ts
@@ -13,7 +13,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({});
+	await ctx.modules.rateLimit.throttlePublic({});
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/remove_friend.ts b/modules/friends/scripts/remove_friend.ts
@@ -11,7 +11,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({ requests: 50 });
+	await ctx.modules.rateLimit.throttlePublic({ requests: 50 });
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/scripts/send_request.ts b/modules/friends/scripts/send_request.ts
@@ -14,7 +14,7 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	await ctx.modules.rateLimit.throttle({});
+	await ctx.modules.rateLimit.throttlePublic({});
 
 	const { userId } = await ctx.modules.users.validateUserToken({
 		userToken: req.userToken,

diff --git a/modules/friends/tests/e2e.ts b/modules/friends/tests/e2e.ts
@@ -3,15 +3,15 @@ import { assertEquals } from "https://deno.land/[email protected]/assert/mod.ts";
 import { faker } from "https://deno.land/x/[email protected]/mod.ts";
 
 test("e2e accept", async (ctx: TestContext) => {
-	const { user: _userA, token: tokenA } = await ctx.modules.users.register({
+	const { user: userA } = await ctx.modules.users.createUser({
 		username: faker.internet.userName(),
-		identity: { guest: {} },
 	});
+	const { token: tokenA } = await ctx.modules.users.createUserToken({ userId: userA.id });
 
-	const { user: userB, token: tokenB } = await ctx.modules.users.register({
+	const { user: userB } = await ctx.modules.users.createUser({
 		username: faker.internet.userName(),
-		identity: { guest: {} },
 	});
+	const { token: tokenB } = await ctx.modules.users.createUserToken({ userId: userB.id });
 
 	const { friendRequest } = await ctx.modules.friends.sendRequest({
 		userToken: tokenA.token,
@@ -57,15 +57,15 @@ test("e2e accept", async (ctx: TestContext) => {
 });
 
 test("e2e reject", async (ctx: TestContext) => {
-	const { user: _userA, token: tokenA } = await ctx.modules.users.register({
+	const { user: userA } = await ctx.modules.users.createUser({
 		username: faker.internet.userName(),
-		identity: { guest: {} },
 	});
+	const { token: tokenA } = await ctx.modules.users.createUserToken({ userId: userA.id });
 
-	const { user: userB, token: tokenB } = await ctx.modules.users.register({
+	const { user: userB } = await ctx.modules.users.createUser({
 		username: faker.internet.userName(),
-		identity: { guest: {} },
 	});
+	const { token: tokenB } = await ctx.modules.users.createUserToken({ userId: userB.id });
 
 	const { friendRequest } = await ctx.modules.friends.sendRequest({
 		userToken: tokenA.token,

diff --git a/modules/rate_limit/db/migrations/20240313051342_init/migration.sql b/modules/rate_limit/db/migrations/20240313051342_init/migration.sql
@@ -0,0 +1,9 @@
+-- CreateTable
+CREATE UNLOGGED TABLE "TokenBuckets" (
+    "type" TEXT NOT NULL,
+    "key" TEXT NOT NULL,
+    "tokens" BIGINT NOT NULL,
+    "lastRefill" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "TokenBuckets_pkey" PRIMARY KEY ("type","key")
+);
diff --git a/modules/rate_limit/db/migrations/migration_lock.toml b/modules/rate_limit/db/migrations/migration_lock.toml
@@ -0,0 +1,3 @@
+# Please do not edit this file manually
+# It should be added in your version-control system (i.e. Git)
+provider = "postgresql"
diff --git a/modules/rate_limit/db/schema.prisma b/modules/rate_limit/db/schema.prisma
@@ -0,0 +1,13 @@
+datasource db {
+    provider = "postgresql"
+    url      = env("DATABASE_URL")
+}
+
+model TokenBuckets {
+    type       String
+    key        String
+    tokens     BigInt
+    lastRefill DateTime
+
+    @@id([type, key])
+}
diff --git a/modules/rate_limit/module.yaml b/modules/rate_limit/module.yaml
@@ -3,4 +3,6 @@ authors:
   - NathanFlurry
 scripts:
   throttle: {}
-errors: {}
+  throttle_public: {}
+errors:
+  RATE_LIMIT_EXCEEDED: {}