From d20aedda3f14d1881ca509c19334125bb3baabca Mon Sep 17 00:00:00 2001
From: moonrailgun <moonrailgun@gmail.com>
Date: Tue, 3 Oct 2023 15:56:09 +0800
Subject: [PATCH] refactor: update trpc entry style

---
 src/server/trpc/index.ts         | 12 +--------
 src/server/trpc/routers/index.ts | 12 +++++++++
 src/server/trpc/trpc.ts          | 46 +++++++++++++++++++++-----------
 3 files changed, 43 insertions(+), 27 deletions(-)
 create mode 100644 src/server/trpc/routers/index.ts

diff --git a/src/server/trpc/index.ts b/src/server/trpc/index.ts
index d3b5287..369f051 100644
--- a/src/server/trpc/index.ts
+++ b/src/server/trpc/index.ts
@@ -1,16 +1,6 @@
 import * as trpcExpress from '@trpc/server/adapters/express';
 import { createContext, router } from './trpc';
-import { notificationRouter } from './routers/notification';
-import { websiteRouter } from './routers/website';
-import { monitorRouter } from './routers/monitor';
-
-const appRouter = router({
-  website: websiteRouter,
-  notification: notificationRouter,
-  monitor: monitorRouter,
-});
-
-export type AppRouter = typeof appRouter;
+import { appRouter } from './routers';
 
 export const trpcExpressMiddleware = trpcExpress.createExpressMiddleware({
   router: appRouter,
diff --git a/src/server/trpc/routers/index.ts b/src/server/trpc/routers/index.ts
new file mode 100644
index 0000000..8d10d5f
--- /dev/null
+++ b/src/server/trpc/routers/index.ts
@@ -0,0 +1,12 @@
+import { router } from '../trpc';
+import { notificationRouter } from './notification';
+import { websiteRouter } from './website';
+import { monitorRouter } from './monitor';
+
+export const appRouter = router({
+  website: websiteRouter,
+  notification: notificationRouter,
+  monitor: monitorRouter,
+});
+
+export type AppRouter = typeof appRouter;
diff --git a/src/server/trpc/trpc.ts b/src/server/trpc/trpc.ts
index 4039bff..0dc0677 100644
--- a/src/server/trpc/trpc.ts
+++ b/src/server/trpc/trpc.ts
@@ -1,25 +1,16 @@
 import { initTRPC, inferAsyncReturnType, TRPCError } from '@trpc/server';
-import * as trpcExpress from '@trpc/server/adapters/express';
 import _ from 'lodash';
 import { z } from 'zod';
 import { jwtVerify } from '../middleware/auth';
 import { getWorkspaceUser } from '../model/workspace';
 import { ROLES, SYSTEM_ROLES } from '../utils/const';
+import type { IncomingMessage } from 'http';
 
-export function createContext({
-  req,
-  res,
-}: trpcExpress.CreateExpressContextOptions) {
+export function createContext({ req }: { req: IncomingMessage }) {
   const authorization = req.headers['authorization'] ?? '';
   const token = authorization.replace('Bearer ', '');
 
-  try {
-    const user = jwtVerify(token);
-
-    return { user };
-  } catch (err) {
-    throw new TRPCError({ code: 'UNAUTHORIZED' });
-  }
+  return { token };
 }
 
 type Context = inferAsyncReturnType<typeof createContext>;
@@ -29,8 +20,31 @@ export const middleware = t.middleware;
 export const router = t.router;
 export const publicProcedure = t.procedure;
 
-const isSystemAdmin = middleware(async (opts) => {
+const isUser = middleware(async (opts) => {
+  const token = opts.ctx.token;
+
+  if (!token) {
+    throw new TRPCError({ code: 'UNAUTHORIZED', message: 'NoToken' });
+  }
+
+  try {
+    const user = jwtVerify(token);
+
+    return opts.next({
+      ctx: {
+        user,
+      },
+    });
+  } catch (err) {
+    throw new TRPCError({ code: 'UNAUTHORIZED', message: 'TokenInvalid' });
+  }
+});
+
+export const protectProedure = t.procedure.use(isUser);
+
+const isSystemAdmin = isUser.unstable_pipe(async (opts) => {
   const { ctx, input } = opts;
+
   if (ctx.user.role !== SYSTEM_ROLES.admin) {
     throw new TRPCError({ code: 'FORBIDDEN' });
   }
@@ -39,14 +53,14 @@ const isSystemAdmin = middleware(async (opts) => {
 });
 
 export const systemAdminProcedure = t.procedure.use(isSystemAdmin);
-export const workspaceProcedure = t.procedure
+export const workspaceProcedure = protectProedure
   .input(
     z.object({
       workspaceId: z.string().uuid(),
     })
   )
   .use(createWorkspacePermissionMiddleware());
-export const workspaceOwnerProcedure = t.procedure
+export const workspaceOwnerProcedure = protectProedure
   .input(
     z.object({
       workspaceId: z.string().uuid(),
@@ -58,7 +72,7 @@ export const workspaceOwnerProcedure = t.procedure
  * Create a trpc middleware which help user check workspace permission
  */
 function createWorkspacePermissionMiddleware(roles: ROLES[] = []) {
-  return middleware(async (opts) => {
+  return isUser.unstable_pipe(async (opts) => {
     const { ctx, input } = opts;
 
     const workspaceId = _.get(input, 'workspaceId', '');