feat(core): add response guard (#1542)

* feat(core): add response guard

* refactor(core): print error if not prod

packages/core/src/errors/ServerError/index.ts

@@ -0,0 +1,3 @@
+// Needs to standardize
+
+export default class ServerError extends Error {}

packages/core/src/middleware/koa-error-handler.test.ts

@@ -38,8 +38,9 @@ describe('koaErrorHandler middleware', () => {
     expect(ctx.body).toEqual(mockBody);
   });
 
-  it('expect to throw if error type is not RequestError', async () => {
+  it('expect status 500 if error type is not RequestError', async () => {
     next.mockRejectedValueOnce(new Error('err'));
-    await expect(koaErrorHandler()(ctx, next)).rejects.toThrow();
+    await koaErrorHandler()(ctx, next);
+    expect(ctx.status).toEqual(500);
   });
 });

packages/core/src/middleware/koa-error-handler.ts

@@ -1,17 +1,22 @@
 import { RequestErrorBody } from '@logto/schemas';
 import { Middleware } from 'koa';
 
+import envSet from '@/env-set';
 import RequestError from '@/errors/RequestError';
 
 export default function koaErrorHandler<StateT, ContextT, BodyT>(): Middleware<
   StateT,
   ContextT,
-  BodyT | RequestErrorBody
+  BodyT | RequestErrorBody | { message: string }
 > {
   return async (ctx, next) => {
     try {
       await next();
     } catch (error: unknown) {
+      if (!envSet.values.isProduction) {
+        console.error(error);
+      }
+
       if (error instanceof RequestError) {
         ctx.status = error.status;
         ctx.body = error.body;
@@ -19,7 +24,8 @@ export default function koaErrorHandler<StateT, ContextT, BodyT>(): Middleware<
         return;
       }
 
-      throw error;
+      ctx.status = 500;
+      ctx.body = { message: 'Internal server error.' };
     }
   };
 }

packages/core/src/middleware/koa-guard.ts

@@ -5,35 +5,40 @@ import { IMiddleware, IRouterParamContext } from 'koa-router';
 import { ZodType } from 'zod';
 
 import RequestError from '@/errors/RequestError';
+import ServerError from '@/errors/ServerError';
+import assertThat from '@/utils/assert-that';
 
-export type GuardConfig<QueryT, BodyT, ParametersT> = {
+export type GuardConfig<QueryT, BodyT, ParametersT, ResponseT> = {
   query?: ZodType<QueryT>;
   body?: ZodType<BodyT>;
   params?: ZodType<ParametersT>;
+  response?: ZodType<ResponseT>;
+  status?: number | number[];
 };
 
-export type Guarded<QueryT, BodyT, ParametersT> = {
+export type GuardedRequest<QueryT, BodyT, ParametersT> = {
   query: QueryT;
   body: BodyT;
   params: ParametersT;
 };
 
-export type WithGuardedContext<
+export type WithGuardedRequestContext<
   ContextT extends IRouterParamContext,
   GuardQueryT,
   GuardBodyT,
   GuardParametersT
 > = ContextT & {
-  guard: Guarded<GuardQueryT, GuardBodyT, GuardParametersT>;
+  guard: GuardedRequest<GuardQueryT, GuardBodyT, GuardParametersT>;
 };
 
 export type WithGuardConfig<
   Type,
   GuardQueryT = unknown,
   GuardBodyT = unknown,
-  GuardParametersT = unknown
+  GuardParametersT = unknown,
+  GuardResponseT = unknown
 > = Type & {
-  config: GuardConfig<GuardQueryT, GuardBodyT, GuardParametersT>;
+  config: GuardConfig<GuardQueryT, GuardBodyT, GuardParametersT, GuardResponseT>;
 };
 
 export const isGuardMiddleware = <Type extends IMiddleware>(
@@ -56,30 +61,32 @@ const tryParse = <Output, Definition, Input>(
 export default function koaGuard<
   StateT,
   ContextT extends IRouterParamContext,
-  ResponseBodyT,
   GuardQueryT = undefined,
   GuardBodyT = undefined,
-  GuardParametersT = undefined
+  GuardParametersT = undefined,
+  GuardResponseT = unknown
 >({
   query,
   body,
   params,
-}: GuardConfig<GuardQueryT, GuardBodyT, GuardParametersT>): MiddlewareType<
+  response,
+  status,
+}: GuardConfig<GuardQueryT, GuardBodyT, GuardParametersT, GuardResponseT>): MiddlewareType<
   StateT,
-  WithGuardedContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
-  ResponseBodyT
+  WithGuardedRequestContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
+  GuardResponseT
 > {
   const guard: MiddlewareType<
     StateT,
-    WithGuardedContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
-    ResponseBodyT
+    WithGuardedRequestContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
+    GuardResponseT
   > = async (ctx, next) => {
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
     ctx.guard = {
       query: tryParse('query', query, ctx.request.query),
       body: tryParse('body', body, ctx.request.body),
       params: tryParse('params', params, ctx.params),
-    } as Guarded<GuardQueryT, GuardBodyT, GuardParametersT>; // Have to do this since it's too complicated for TS
+    } as GuardedRequest<GuardQueryT, GuardBodyT, GuardParametersT>; // Have to do this since it's too complicated for TS
 
     return next();
   };
@@ -87,15 +94,28 @@ export default function koaGuard<
   const guardMiddleware: WithGuardConfig<
     MiddlewareType<
       StateT,
-      WithGuardedContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
-      ResponseBodyT
+      WithGuardedRequestContext<ContextT, GuardQueryT, GuardBodyT, GuardParametersT>,
+      GuardResponseT
     >
   > = async function (ctx, next) {
     if (body) {
       return koaBody<StateT, ContextT>()(ctx, async () => guard(ctx, next));
     }
 
-    return guard(ctx, next);
+    await guard(ctx, next);
+
+    if (status !== undefined) {
+      assertThat(
+        Array.isArray(status)
+          ? status.includes(ctx.response.status)
+          : status === ctx.response.status,
+        new ServerError()
+      );
+    }
+
+    if (response !== undefined) {
+      assertThat(response.safeParse(ctx.body).success, new ServerError());
+    }
   };
 
   // Intended

packages/core/src/routes/log.test.ts

@@ -2,8 +2,9 @@ import { LogCondition } from '@/queries/log';
 import logRoutes from '@/routes/log';
 import { createRequester } from '@/utils/test-utils';
 
-const mockLog = { id: 1 };
-const mockLogs = [mockLog, { id: 2 }];
+const mockBody = { type: 'a', payload: {}, createdAt: 123 };
+const mockLog = { id: '1', ...mockBody };
+const mockLogs = [mockLog, { id: '2', ...mockBody }];
 
 /* eslint-disable @typescript-eslint/no-unused-vars */
 const countLogs = jest.fn(async (condition: LogCondition) => ({

packages/core/src/routes/log.ts

@@ -1,3 +1,4 @@
+import { Logs } from '@logto/schemas';
 import { object, string } from 'zod';
 
 import koaGuard from '@/middleware/koa-guard';
@@ -38,7 +39,7 @@ export default function logRoutes<T extends AuthedRouter>(router: T) {
 
   router.get(
     '/logs/:id',
-    koaGuard({ params: object({ id: string().min(1) }) }),
+    koaGuard({ params: object({ id: string().min(1) }), response: Logs.guard }),
     async (ctx, next) => {
       const {
         params: { id },

packages/core/src/routes/status.ts

@@ -1,7 +1,9 @@
+import koaGuard from '@/middleware/koa-guard';
+
 import { AnonymousRouter } from './types';
 
 export default function statusRoutes<T extends AnonymousRouter>(router: T) {
-  router.get('/status', async (ctx, next) => {
+  router.get('/status', koaGuard({ status: 204 }), async (ctx, next) => {
     ctx.status = 204;
 
     return next();

packages/schemas/src/foundations/schemas.ts

@@ -4,10 +4,14 @@ type ParseOptional<K> = undefined extends K
   ? ZodOptional<ZodType<Exclude<K, undefined>>>
   : ZodType<K>;
 
-export type Guard<T extends Record<string, unknown>> = ZodObject<{
+export type CreateGuard<T extends Record<string, unknown>> = ZodObject<{
   [key in keyof T]-?: ParseOptional<T[key]>;
 }>;
 
+export type Guard<T extends Record<string, unknown>> = ZodObject<{
+  [key in keyof T]: ZodType<T[key]>;
+}>;
+
 export type SchemaValuePrimitive = string | number | boolean | undefined;
 export type SchemaValue = SchemaValuePrimitive | Record<string, unknown> | unknown[] | null;
 export type SchemaLike<Key extends string = string> = {
@@ -22,6 +26,7 @@ export type GeneratedSchema<Schema extends SchemaLike> = keyof Schema extends st
         [key in keyof Schema]: string;
       };
       fieldKeys: ReadonlyArray<keyof Schema>;
-      createGuard: Guard<Schema>;
+      createGuard: CreateGuard<Schema>;
+      guard: Guard<Schema>;
     }>
   : never;

packages/schemas/src/gen/index.ts

@@ -174,7 +174,7 @@ const generate = async () => {
       }));
 
       if (tableWithTypes.length > 0) {
-        tsTypes.push('GeneratedSchema', 'Guard');
+        tsTypes.push('GeneratedSchema', 'Guard', 'CreateGuard');
       }
       /* eslint-enable @silverhand/fp/no-mutating-methods */
 

packages/schemas/src/gen/schema.ts

@@ -32,23 +32,38 @@ export const generateSchema = ({ name, fields }: TableWithType) => {
     ),
     '};',
     '',
-    `const createGuard: Guard<${databaseEntryType}> = z.object({`,
+    `const createGuard: CreateGuard<${databaseEntryType}> = z.object({`,
     // eslint-disable-next-line complexity
     ...fields.map(({ name, type, isArray, isEnum, nullable, hasDefaultValue, tsType }) => {
       if (tsType) {
         return `  ${camelcase(name)}: ${camelcase(tsType)}Guard${conditionalString(
-          nullable && !hasDefaultValue && '.nullable()'
+          nullable && '.nullable()'
         )}${conditionalString((nullable || hasDefaultValue) && '.optional()')},`;
       }
 
       return `  ${camelcase(name)}: z.${
         isEnum ? `nativeEnum(${type})` : `${type}()`
       }${conditionalString(isArray && '.array()')}${conditionalString(
-        nullable && !hasDefaultValue && '.nullable()'
+        nullable && '.nullable()'
       )}${conditionalString((nullable || hasDefaultValue) && '.optional()')},`;
     }),
     '  });',
     '',
+    `const guard: Guard<${modelName}> = z.object({`,
+    // eslint-disable-next-line complexity
+    ...fields.map(({ name, type, isArray, isEnum, nullable, tsType }) => {
+      if (tsType) {
+        return `  ${camelcase(name)}: ${camelcase(tsType)}Guard${conditionalString(
+          nullable && '.nullable()'
+        )},`;
+      }
+
+      return `  ${camelcase(name)}: z.${
+        isEnum ? `nativeEnum(${type})` : `${type}()`
+      }${conditionalString(isArray && '.array()')}${conditionalString(nullable && '.nullable()')},`;
+    }),
+    '  });',
+    '',
     `export const ${camelcase(name, {
       pascalCase: true,
     })}: GeneratedSchema<${databaseEntryType}> = Object.freeze({`,
@@ -59,8 +74,9 @@ export const generateSchema = ({ name, fields }: TableWithType) => {
     '  },',
     '  fieldKeys: [',
     ...fields.map(({ name }) => `    '${camelcase(name)}',`),
-    '  ],',
+    '  ] as const,',
     '  createGuard,',
+    '  guard,',
     '});',
   ].join('\n');
 };