Merge pull request #29 from KennethWussmann/task/28

package.json

@@ -39,7 +39,9 @@
     "next": "^13.5.4",
     "next-auth": "^4.23.0",
     "pdfjs": "^2.5.2",
+    "pg": "^8.11.3",
     "qrcode": "^1.5.3",
+    "rate-limiter-flexible": "^4.0.1",
     "react": "18.2.0",
     "react-datepicker": "^4.21.0",
     "react-dom": "18.2.0",
@@ -58,6 +60,7 @@
     "@types/eslint": "^8.44.2",
     "@types/lodash": "^4.14.201",
     "@types/node": "^18.16.0",
+    "@types/pg": "^8.11.0",
     "@types/qrcode": "^1.5.5",
     "@types/react": "^18.2.20",
     "@types/react-datepicker": "^4.19.1",

prisma/schema.prisma

@@ -208,3 +208,9 @@ model LabelTemplate {
     team        Team?             @relation(fields: [teamId], references: [id], onDelete: Cascade)
     teamId      String?
 }
+
+model RateLimit {
+    key    String @id
+    points Int    @default(0)
+    expire BigInt
+}

src/env.mjs

@@ -25,6 +25,7 @@ export const env = createEnv({
     APP_BASE_URL: z.string().url(),
     MEILI_URL: z.string().url(),
     MEILI_MASTER_KEY: z.string(),
+    DISABLE_RATE_LIMIT: z.boolean().optional().default(false),
 
     PASSWORD_AUTH_ENABLED: z
       .string()
@@ -67,6 +68,7 @@ export const env = createEnv({
     APP_BASE_URL: process.env.APP_BASE_URL,
     MEILI_URL: process.env.MEILI_URL,
     MEILI_MASTER_KEY: process.env.MEILI_MASTER_KEY,
+    DISABLE_RATE_LIMIT: process.env.DISABLE_RATE_LIMIT,
 
     COGNITO_CLIENT_ID: process.env.COGNITO_CLIENT_ID,
     COGNITO_CLIENT_SECRET: process.env.COGNITO_CLIENT_SECRET,

src/server/api/routers/user.ts

@@ -12,6 +12,9 @@ export const userRouter = createTRPCRouter({
       if (ctx.session?.user) {
         throw new Error("User already logged in");
       }
-      await ctx.applicationContext.userService.register(input);
+      await ctx.applicationContext.userService.register(
+        ctx.remoteAddress,
+        input
+      );
     }),
 });

src/server/api/trpc.ts

@@ -16,6 +16,8 @@ import { ApplicationContext } from "~/server/lib/applicationContext";
 
 import { getServerAuthSession } from "~/server/auth/auth";
 import { db } from "~/server/db";
+import { RateLimitType } from "../lib/user/rateLimitService";
+import { env } from "~/env.mjs";
 
 /**
  * 1. CONTEXT
@@ -27,6 +29,7 @@ import { db } from "~/server/db";
 
 interface CreateContextOptions {
   session: Session | null;
+  remoteAddress: string;
   applicationContext: ApplicationContext;
 }
 
@@ -59,8 +62,24 @@ export const createTRPCContext = async (opts: CreateNextContextOptions) => {
   // Get the session from the server using the getServerSession wrapper function
   const session = await getServerAuthSession({ req, res });
 
+  const xForwardedFor = req.headers["x-forwarded-for"];
+  let remoteAddress = req.socket.remoteAddress;
+  if (xForwardedFor) {
+    remoteAddress = Array.isArray(xForwardedFor)
+      ? xForwardedFor[0]
+      : xForwardedFor.split(",")[0];
+  }
+
+  if (!remoteAddress || remoteAddress.length === 0) {
+    throw new TRPCError({
+      code: "INTERNAL_SERVER_ERROR",
+      message: "Could not determine remote address",
+    });
+  }
+
   return createInnerTRPCContext({
     session,
+    remoteAddress,
     applicationContext: new ApplicationContext(),
   });
 };
@@ -101,14 +120,33 @@ const t = initTRPC.context<typeof createTRPCContext>().create({
  */
 export const createTRPCRouter = t.router;
 
+export const rateLimit = (type: RateLimitType) =>
+  t.middleware(async ({ ctx, next }) => {
+    try {
+      await ctx.applicationContext.rateLimitService.consume(
+        type,
+        ctx.remoteAddress,
+        ctx.session?.user?.id
+      );
+    } catch (e) {
+      if (e instanceof Error) {
+        throw new TRPCError({ code: "TOO_MANY_REQUESTS", message: e.message });
+      } else {
+        throw new TRPCError({ code: "TOO_MANY_REQUESTS" });
+      }
+    }
+
+    return next();
+  });
+
 /**
  * Public (unauthenticated) procedure
  *
  * This is the base piece you use to build new queries and mutations on your tRPC API. It does not
  * guarantee that a user querying is authorized, but you can still access user session data if they
  * are logged in.
  */
-export const publicProcedure = t.procedure;
+export const publicProcedure = t.procedure.use(rateLimit("request"));
 
 /** Reusable middleware that enforces users are logged in before running the procedure. */
 const enforceUserIsAuthed = t.middleware(async ({ ctx, next }) => {
@@ -134,4 +172,6 @@ const enforceUserIsAuthed = t.middleware(async ({ ctx, next }) => {
  *
  * @see https://trpc.io/docs/procedures
  */
-export const protectedProcedure = t.procedure.use(enforceUserIsAuthed);
+export const protectedProcedure = t.procedure
+  .use(enforceUserIsAuthed)
+  .use(rateLimit("request"));