feat(text-normalization): add an IP based ratelimit

functions/package.json

@@ -24,6 +24,7 @@
     "@google-cloud/storage": "7.13.0",
     "@sign-mt/browsermt": "0.2.3",
     "@unkey/api": "0.26.2",
+    "@unkey/ratelimit": "0.4.5",
     "cors": "2.8.5",
     "express": "4.21.1",
     "express-async-errors": "3.1.1",
@@ -32,16 +33,18 @@
     "http-errors": "2.0.0",
     "http-proxy-middleware": "^3.0.3",
     "node-fetch": "2.6.7",
-    "openai": "4.67.3"
+    "openai": "4.67.3",
+    "request-ip": "3.3.0"
   },
   "devDependencies": {
     "@firebase/firestore-types": "3.0.2",
     "@firebase/rules-unit-testing": "3.0.4",
     "@types/http-errors": "2.0.4",
     "@types/jest": "29.5.13",
     "@types/node-fetch": "2.6.11",
-    "@typescript-eslint/eslint-plugin": "8.8.1",
-    "@typescript-eslint/parser": "8.8.1",
+    "@types/request-ip": "0.0.41",
+    "@typescript-eslint/eslint-plugin": "8.9.0",
+    "@typescript-eslint/parser": "8.9.0",
     "eslint": "8.57.0",
     "firebase-functions-test": "3.3.0",
     "firebase-tools": "13.22.0",

functions/src/gateway/controller.ts

@@ -1,56 +1,22 @@
 import * as express from 'express';
-import {Application, NextFunction, Request, Response} from 'express';
+import {Application} from 'express';
 
 import * as cors from 'cors';
 import {errorMiddleware} from '../middlewares/error.middleware';
 import {onRequest} from 'firebase-functions/v2/https';
-import {unkeyAuth} from './middleware';
+import {unkeyAuth} from '../middlewares/unkey-auth.middleware';
 import {HttpsOptions} from 'firebase-functions/lib/v2/providers/https';
-import {getAppCheck} from 'firebase-admin/app-check';
 import {avatars} from './avatars';
 import {me} from './me';
 import {spokenToSigned} from './spoken-to-signed';
 import {optionsRequest} from '../middlewares/options.request';
-
-// The public APP ID of the sign-mt web app
-const APP_ID = '1:665830225099:web:18e0669d5847a4b047974e';
-
-// Create and cache an App Check token for the sign-mt web app
-let appCheckAPIKey = Promise.resolve({token: '', expires: 0});
-
-export async function getAppCheckKey(req: Request, res: Response, next: NextFunction) {
-  async function safeGetToken() {
-    // If there was a failure to get the token, reset the promise before trying again
-    try {
-      return await appCheckAPIKey;
-    } catch (e) {
-      appCheckAPIKey = Promise.resolve({token: '', expires: 0});
-      throw e;
-    }
-  }
-
-  let {token, expires} = await safeGetToken();
-  if (expires < Date.now()) {
-    appCheckAPIKey = getAppCheck()
-      .createToken(APP_ID)
-      .then(({token, ttlMillis}) => ({token, expires: Date.now() - 1000 + ttlMillis}));
-    ({token, expires} = await safeGetToken());
-  }
-
-  res.locals.appCheckToken = token;
-  res.locals.headers = {
-    'X-Firebase-AppCheck': token,
-    'X-AppCheck-Token': token,
-  };
-
-  return next();
-}
+import {createAppCheckKey} from '../middlewares/create-appcheck.middleware';
 
 // API Documentation: https://app.swaggerhub.com/apis/AmitMoryossef/sign_mt/
 const app: Application = express();
 app.use(cors());
 app.use(unkeyAuth);
-app.use(getAppCheckKey);
+app.use(createAppCheckKey);
 app.options('*', optionsRequest);
 
 spokenToSigned(app);

functions/src/middlewares/create-appcheck.middleware.ts

@@ -0,0 +1,36 @@
+import {NextFunction, Request, Response} from 'express';
+import {getAppCheck} from 'firebase-admin/app-check';
+
+// The public APP ID of the sign-mt web app
+const APP_ID = '1:665830225099:web:18e0669d5847a4b047974e';
+
+// Create and cache an App Check token for the sign-mt web app
+let appCheckAPIKey = Promise.resolve({token: '', expires: 0});
+
+export async function createAppCheckKey(req: Request, res: Response, next: NextFunction) {
+  async function safeGetToken() {
+    // If there was a failure to get the token, reset the promise before trying again
+    try {
+      return await appCheckAPIKey;
+    } catch (e) {
+      appCheckAPIKey = Promise.resolve({token: '', expires: 0});
+      throw e;
+    }
+  }
+
+  let {token, expires} = await safeGetToken();
+  if (expires < Date.now()) {
+    appCheckAPIKey = getAppCheck()
+      .createToken(APP_ID)
+      .then(({token, ttlMillis}) => ({token, expires: Date.now() - 1000 + ttlMillis}));
+    ({token, expires} = await safeGetToken());
+  }
+
+  res.locals.appCheckToken = token;
+  res.locals.headers = {
+    'X-Firebase-AppCheck': token,
+    'X-AppCheck-Token': token,
+  };
+
+  return next();
+}

functions/src/gateway/middleware.ts → functions/src/middlewares/unkey-auth.middleware.ts

@@ -1,8 +1,12 @@
 import {verifyKey} from '@unkey/api';
 import * as httpErrors from 'http-errors';
 import {NextFunction, Request, Response} from 'express';
+import {rateLimitHeaders} from './unkey-ratelimit.middleware';
+import {RatelimitResponse} from '@unkey/ratelimit';
 
 export async function unkeyAuth(req: Request, res: Response, next: NextFunction) {
+  const apiId = 'api_4LtAUnGvWjPZGJV9hQDoJtum53GK'; // Public API ID
+
   const authHeader = req.headers['authorization'];
   let key = authHeader?.toString().replace('Bearer ', '');
 
@@ -14,12 +18,16 @@ export async function unkeyAuth(req: Request, res: Response, next: NextFunction)
     throw new httpErrors.BadRequest('Missing/invalid API key');
   }
 
-  const {result, error} = await verifyKey(key);
+  const {result, error} = await verifyKey({key, apiId});
   if (error) {
     // This may happen on network errors
     throw new httpErrors.InternalServerError('Could not verify API key, please try again later');
   }
 
+  if (result?.ratelimit) {
+    rateLimitHeaders(res, result.ratelimit as RatelimitResponse);
+  }
+
   if (!result.valid) {
     throw new httpErrors.Unauthorized('Invalid API key');
   }

functions/src/middlewares/unkey-ratelimit.middleware.ts

@@ -0,0 +1,39 @@
+import {Duration, Ratelimit, RatelimitResponse} from '@unkey/ratelimit';
+import * as httpErrors from 'http-errors';
+import * as requestIp from 'request-ip';
+import {NextFunction, Request, Response} from 'express';
+import {defineString} from 'firebase-functions/params';
+
+export function rateLimitHeaders(res: Response, ratelimitResponse: RatelimitResponse, duration?: Duration) {
+  res.setHeader('X-RateLimit-Limit', ratelimitResponse.limit.toString());
+  res.setHeader('X-RateLimit-Remaining', ratelimitResponse.remaining.toString());
+  res.setHeader('X-RateLimit-Reset', ratelimitResponse.reset.toString());
+  if (duration) {
+    res.setHeader('X-RateLimit-Policy', `${ratelimitResponse.limit};w=${duration}`);
+  }
+}
+
+export function unkeyRatelimit(namespace: string, limit: number, duration: Duration) {
+  const unkeyRootKey = defineString('UNKEY_ROOT_KEY');
+
+  return async function (req: Request, res: Response, next: NextFunction) {
+    const identifier = requestIp.getClientIp(req) ?? 'unknown';
+
+    const rateLimit = new Ratelimit({
+      rootKey: unkeyRootKey.value(),
+      namespace,
+      limit,
+      duration,
+      async: true,
+    });
+
+    const ratelimitResponse = await rateLimit.limit(identifier);
+    rateLimitHeaders(res, ratelimitResponse, duration);
+
+    if (!ratelimitResponse.success) {
+      throw new httpErrors.TooManyRequests('Too many requests, please try again later');
+    }
+
+    return next();
+  };
+}

functions/src/text-normalization/controller.ts

@@ -10,6 +10,7 @@ import {defineString} from 'firebase-functions/params';
 import {appCheckVerification} from '../middlewares/appcheck.middleware';
 import type {StringParam} from 'firebase-functions/lib/params/types';
 import {optionsRequest} from '../middlewares/options.request';
+import {unkeyRatelimit} from '../middlewares/unkey-ratelimit.middleware';
 
 export class TextNormalizationEndpoint {
   constructor(private database: FirebaseDatabase, private OpenAIApiKey: StringParam) {}
@@ -107,6 +108,7 @@ export const textNormalizationFunctions = (database: FirebaseDatabase) => {
   const app = express();
   app.use(cors());
   app.use(appCheckVerification);
+  app.use(unkeyRatelimit('api.text-normalization', 250, '30m'));
   app.options('*', optionsRequest);
   app.get(['/', '/api/text-normalization'], request);
   app.use(errorMiddleware);

package.json

@@ -133,8 +133,8 @@
     "@types/web-app-manifest": "1.0.8",
     "@types/webgl2": "0.0.11",
     "@types/wicg-file-system-access": "2023.10.5",
-    "@typescript-eslint/eslint-plugin": "8.8.1",
-    "@typescript-eslint/parser": "8.8.1",
+    "@typescript-eslint/eslint-plugin": "8.9.0",
+    "@typescript-eslint/parser": "8.9.0",
     "browser-sync": "3.0.3",
     "deepmerge": "4.3.1",
     "dotenv": "16.4.5",
@@ -155,7 +155,7 @@
     "karma-safari-launcher": "1.0.0",
     "karma-spec-reporter": "0.0.36",
     "lint-staged": "15.2.10",
-    "marked": "14.1.2",
+    "marked": "14.1.3",
     "npm-license-crawler": "0.2.1",
     "open": "10.1.0",
     "pwa-asset-generator": "6.3.2",