feat: support variadic middleware for routes

src/koa-utils.ts

@@ -94,6 +94,11 @@ export type EndpointImplementation<
   context: ContextOfEndpoint<RouteName, Request, RouterType>,
 ) => Response | Promise<Response>;
 
+export type OneSchemaRouterMiddleware<Context> = (
+  context: Context,
+  next: () => Promise<any>,
+) => any | Promise<any>;
+
 export const implementRoute = <
   Route extends string,
   Request,
@@ -103,6 +108,9 @@ export const implementRoute = <
   route: Route,
   router: R,
   parse: (ctx: ContextOfEndpoint<Route, Request, R>, data: unknown) => Request,
+  middlewares: OneSchemaRouterMiddleware<
+    ContextOfEndpoint<Route, Request, R>
+  >[],
   implementation: EndpointImplementation<Route, Request, Response, R>,
 ) => {
   // Separate method and path. e.g. 'POST my/route' => ['POST', 'my/route']
@@ -163,19 +171,19 @@ export const implementRoute = <
   // Register the route + handler on the router.
   switch (method) {
     case 'POST':
-      router.post(path, handler);
+      router.post(path, ...middlewares, handler);
       break;
     case 'GET':
-      router.get(path, handler);
+      router.get(path, ...middlewares, handler);
       break;
     case 'PUT':
-      router.put(path, handler);
+      router.put(path, ...middlewares, handler);
       break;
     case 'PATCH':
-      router.patch(path, handler);
+      router.patch(path, ...middlewares, handler);
       break;
     case 'DELETE':
-      router.delete(path, handler);
+      router.delete(path, ...middlewares, handler);
       break;
     default:
       throw new Error(`Unsupported method detected: ${route}`);

src/koa.ts

@@ -135,6 +135,7 @@ export const implementSchema = <
           data,
         });
       },
+      [],
       routeHandler,
     );
   }

src/router.test.ts

@@ -410,10 +410,11 @@ describe('output validation', () => {
             request: z.object({}),
             response: z.object({ message: z.string() }),
           })
-          .implement(`${method} /items`, () => ({
+          .implement(
+            `${method} /items`,
             // @ts-expect-error Intentionally writing incorrect TS here
-            message: 123,
-          })),
+            () => ({ message: 123 }),
+          ),
       );
 
       const { status } = await client.request({
@@ -509,6 +510,105 @@ describe('implementations', () => {
   });
 });
 
+describe('using middleware', () => {
+  test('type errors are caught when using middleware', () => {
+    type CustomState = { message: string };
+
+    setup(() =>
+      OneSchemaRouter.create({
+        using: new Router<CustomState>(),
+        introspection: undefined,
+      })
+        .declare({
+          name: 'putItem',
+          route: 'PUT /items/:id',
+          request: z.object({ message: z.string() }),
+          response: z.object({ id: z.string(), message: z.string() }),
+        })
+        .implement(
+          'PUT /items/:id',
+          // @ts-expect-error
+          async (ctx, next) => {
+            ctx.state.message = ctx.request.body.message + '-response';
+            await next();
+          },
+          // We're leaving out the id value here, which should cause the TS error.
+          (ctx) => ({ message: ctx.state.message }),
+        )
+        .implement(
+          'PUT /items/:id',
+          async (ctx, next) => {
+            ctx.params.id;
+
+            ctx.request.body.message;
+
+            // @ts-expect-error The params should be well-typed.
+            ctx.params.bogus;
+
+            // @ts-expect-error The body should be well-typed.
+            ctx.request.body.bogus;
+
+            await next();
+          },
+          (ctx) => ({ id: 'test-id', message: ctx.state.message }),
+        )
+        .implement(
+          'PUT /items/:id',
+          (ctx, next) => {
+            // This call implicitly tests that `message` is a string.
+            ctx.state.message.endsWith('');
+
+            // @ts-expect-error The state should be well-typed.
+            ctx.state.bogus;
+
+            return next();
+          },
+          (ctx) => ({ id: 'test-id', message: ctx.state.message }),
+        ),
+    );
+  });
+
+  test('middlewares are actually executed', async () => {
+    const mock = jest.fn();
+    const { typed: client } = setup((router) =>
+      router
+        .declare({
+          name: 'putItem',
+          route: 'PUT /items/:id',
+          request: z.object({ message: z.string() }),
+          response: z.object({ id: z.string(), message: z.string() }),
+        })
+        .implement(
+          'PUT /items/:id',
+          async (ctx, next) => {
+            mock('middleware 1', ctx.state.message);
+            ctx.state.message = 'message 1';
+            await next();
+          },
+          (ctx, next) => {
+            mock('middleware 2', ctx.state.message);
+            ctx.state.message = 'message 2';
+            return next();
+          },
+          (ctx) => ({ id: ctx.params.id, message: ctx.state.message }),
+        ),
+    );
+
+    const { status, data } = await client.putItem({
+      id: 'test-id-bleh',
+      message: 'test-message',
+    });
+
+    expect(status).toStrictEqual(200);
+    expect(data).toStrictEqual({ id: 'test-id-bleh', message: 'message 2' });
+
+    expect(mock.mock.calls).toEqual([
+      ['middleware 1', undefined],
+      ['middleware 2', 'message 1'],
+    ]);
+  });
+});
+
 describe('introspection', () => {
   test('introspecting + generating a client', async () => {
     const { client } = serve(

src/router.ts

@@ -4,8 +4,10 @@ import { fromZodError } from 'zod-validation-error';
 import zodToJsonSchema from 'zod-to-json-schema';
 import compose = require('koa-compose');
 import {
+  ContextOfEndpoint,
   EndpointImplementation,
   implementRoute,
+  OneSchemaRouterMiddleware,
   PathParamsOf,
 } from './koa-utils';
 import { AxiosInstance, AxiosRequestConfig, AxiosResponse } from 'axios';
@@ -97,15 +99,24 @@ export class OneSchemaRouter<
 
   implement<Route extends keyof Schema & string>(
     route: Route,
-    implementation: EndpointImplementation<
-      Route,
-      z.output<Schema[Route]['request']>,
-      z.infer<Schema[Route]['response']>,
-      R
-    >,
-  ) {
+    ...middlewares: [
+      ...OneSchemaRouterMiddleware<
+        ContextOfEndpoint<Route, z.output<Schema[Route]['request']>, R>
+      >[],
+      EndpointImplementation<
+        Route,
+        z.output<Schema[Route]['request']>,
+        z.infer<Schema[Route]['response']>,
+        R
+      >,
+    ]
+  ): this {
     const endpoint = this.schema[route];
 
+    const mws = middlewares.slice(0, -1) as any[];
+
+    const implementation = middlewares.at(-1) as any;
+
     implementRoute(
       route,
       this.router,
@@ -119,6 +130,8 @@ export class OneSchemaRouter<
         }
         return res.data;
       },
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+      mws,
       async (ctx) => {
         const result = await implementation(ctx);
         const res = endpoint.response.safeParse(result);