From 9603c4f83dd572f9319c5a8b0059409ea7c33ceb Mon Sep 17 00:00:00 2001
From: Swain Molster <swain.molster@lifeomic.com>
Date: Wed, 27 Jul 2022 13:23:29 -0400
Subject: [PATCH] fix: make body typing more deterministic

---
 src/koa.test.ts | 20 +++++++++++++++++
 src/koa.ts      | 59 +++++++++++++++++++++++++++++++------------------
 2 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/src/koa.test.ts b/src/koa.test.ts
index a36cf55..e7ba779 100644
--- a/src/koa.test.ts
+++ b/src/koa.test.ts
@@ -51,10 +51,30 @@ test('setting a 200-level response code overrides the response', async () => {
       implementation: {
         'DELETE /post/:id': (ctx) => {
           ctx.response.status = 204;
+
+          // This comment essentially serves as a "test" that the `body` property
+          // has been removed from the `request` object.
+          // @ts-expect-error
+          ctx.request.body;
+
+          // This line serves as an implicit test that the `query` property
+          // is present on the `request` object.
+          ctx.request.query;
+
           return {};
         },
         'POST /posts': (ctx) => {
           ctx.response.status = 301;
+
+          // This line serves as an implicit test that the `body` property
+          // is present on the `request` object.
+          ctx.request.body;
+
+          // This comment essentially serves as a "test" that the `query` property
+          // has been removed from the `request` object.
+          // @ts-expect-error
+          ctx.request.query;
+
           return {};
         },
       },
diff --git a/src/koa.ts b/src/koa.ts
index 0814b29..b643137 100644
--- a/src/koa.ts
+++ b/src/koa.ts
@@ -4,34 +4,49 @@ import type Router from 'koa-router';
 import type { EndpointsOf, IntrospectionResponse, OneSchema } from './types';
 
 /**
- * This declare is required to override the "declare" statements common in Koa
- * bodyparser libraries, including:
- *   - @types/koa-bodyparser (as of v4.3.7)
- *   - koa-body (as of v5.0.0)
+ * We use this type to very cleanly remove these fields from the Koa context, so
+ * that we can replace the fields with our strict types from the generated schema.
  *
- * Without this declare, installing either of these^ libraries will cause the
- * typings from one-schema to be overriden and collapsed into "any".
+ * - `params`
+ * - `request.body`
+ * - `request.query`
  *
- * Unfortunately, adding this declare causes it to conflict with those libraries,
- * and thus requires consumers to use `skipLibCheck: true`.
+ * This is primarily important for `request.body` -- the others are included in the
+ * same why for simplicity.
+ *
+ * Why it's important for `request.body`: most popular Koa body parser middlewares use a
+ * global `declare` statement to mark the `body` field as being present on the request.
+ *
+ * Most of these libraries declare the field as `body: any`. With this type declaration,
+ * it becomes impossible to augment the type, since any more stringent types are just
+ * "collapsed" into the `any` type, resulting in a final type of `any`.
+ *
+ * By explicitly removing the fields from the context, then re-adding them, we can be
+ * sure they are typed correctly.
  */
-declare module 'koa' {
-  interface Request {
-    body?: unknown;
-  }
-}
+type WithTypedFieldsRemoved<T> =
+  // Omit params and request
+  Omit<T, 'params' | 'request'> & {
+    // Re-add request, but without the "body" or "query" fields.
+    request: Omit<T, 'body' | 'query'>;
+  };
 
 export type ImplementationOf<Schema extends OneSchema<any>, State, Context> = {
   [Name in keyof EndpointsOf<Schema>]: (
-    context: ParameterizedContext<
-      State,
-      Context & {
-        params: EndpointsOf<Schema>[Name]['PathParams'];
-        request: Name extends `${'GET' | 'DELETE'} ${string}`
-          ? { query: EndpointsOf<Schema>[Name]['Request'] }
-          : { body: EndpointsOf<Schema>[Name]['Request'] };
-      }
-    >,
+    // It's important that we remove our "typed" fields from the root `ParameterizedContext`,
+    // and not from the "inner" `Context` type.
+    //
+    // If we remove from the "inner" `Context` type, the `ParameterizedContext` will
+    // effectively just "re-add" the fields we removed.
+    //
+    // Basically, the "inner" Context can only be used for _extending_ the context,
+    // but not for _restricting_ it.
+    context: WithTypedFieldsRemoved<ParameterizedContext<State, Context>> & {
+      params: EndpointsOf<Schema>[Name]['PathParams'];
+      request: Name extends `${'GET' | 'DELETE'} ${string}`
+        ? { query: EndpointsOf<Schema>[Name]['Request'] }
+        : { body: EndpointsOf<Schema>[Name]['Request'] };
+    },
   ) =>
     | EndpointsOf<Schema>[Name]['Response']
     | Promise<EndpointsOf<Schema>[Name]['Response']>;