From cfc6933e78a1c36713454ec2a1c69c52b6746c1a Mon Sep 17 00:00:00 2001
From: Jerko Steiner <jerko.steiner@gmail.com>
Date: Wed, 31 Jul 2019 01:10:27 +0800
Subject: [PATCH] jsonrpc: Add ability to specify context via high-order
 functions

---
 packages/jsonrpc/src/server.test.ts | 22 +++++++++++++++++++---
 packages/jsonrpc/src/server.ts      |  6 +++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/packages/jsonrpc/src/server.test.ts b/packages/jsonrpc/src/server.test.ts
index d0a4c6a..3097e0a 100644
--- a/packages/jsonrpc/src/server.test.ts
+++ b/packages/jsonrpc/src/server.test.ts
@@ -5,12 +5,18 @@ import bodyParser from 'body-parser'
 
 describe('jsonrpc', () => {
 
+  interface IContext {
+    userId: number
+  }
+
   interface IService {
     add(a: number, b: number): number
     delay(): Promise<void>
     syncError(message: string): void
     asyncError(message: string): Promise<void>
     httpError(statusCode: number, message: string): Promise<void>
+
+    addWithContext(a: number, b: number): (ctx: IContext) => number
   }
 
   class Service implements IService {
@@ -40,6 +46,9 @@ describe('jsonrpc', () => {
       }]
       throw err
     }
+    addWithContext = (a: number, b: number) => (ctx: IContext): number => {
+      return a + b + ctx.userId
+    }
   }
 
   function createApp() {
@@ -51,14 +60,17 @@ describe('jsonrpc', () => {
       'syncError',
       'asyncError',
       'httpError',
-    ]))
+      'addWithContext',
+    ], req => ({userId: 1000})))
     return app
   }
 
-  type ArgumentTypes<T> = T extends (...args: infer U) => infer R ? U: never
+  type ArgumentTypes<T> = T extends (...args: infer U) => infer R ? U : never
   type RetType<T> = T extends (...args: any[]) => infer R ? R : never
+  type UnwrapHOC<T> = T extends (...args: any[]) => infer R ? R : T
   type RetProm<T> = T extends Promise<any> ? T : Promise<T>
-  type PromisifyReturnType<T> = (...a: ArgumentTypes<T>) => RetProm<RetType<T>>
+  type PromisifyReturnType<T> = (...a: ArgumentTypes<T>) =>
+    RetProm<UnwrapHOC<RetType<T>>>
   type Asyncified<T> = {
     [K in keyof T]: PromisifyReturnType<T[K]>
   }
@@ -157,6 +169,10 @@ describe('jsonrpc', () => {
       const response = await client.delay()
       expect(response).toEqual(undefined)
     })
+    it('can use context', async () => {
+      const response = await client.addWithContext(5, 7)
+      expect(response).toEqual(1000 + 5 + 7)
+    })
     it('handles synchronous notifications', async () => {
       await request(createApp())
       .post('/myService')
diff --git a/packages/jsonrpc/src/server.ts b/packages/jsonrpc/src/server.ts
index 89f005f..e6e3528 100644
--- a/packages/jsonrpc/src/server.ts
+++ b/packages/jsonrpc/src/server.ts
@@ -95,9 +95,10 @@ function isPromise(value: any): value is Promise<unknown> {
  * Adds middleware for handling JSON RPC requests. Expects JSON middleware to
  * already be configured.
  */
-export function jsonrpc<T, F extends FunctionPropertyNames<T>>(
+export function jsonrpc<T, F extends FunctionPropertyNames<T>, Ctx>(
   service: T,
   functions: F[],
+  getContext: (req: Request) => Ctx,
 ) {
   const rpcService = pick(service, functions)
 
@@ -135,6 +136,9 @@ export function jsonrpc<T, F extends FunctionPropertyNames<T>>(
     let retValue
     try {
       retValue = (rpcService as any)[method](...params)
+      if (typeof retValue === 'function') {
+        retValue = retValue(getContext(req))
+      }
     } catch (err) {
       return handleError(err, req, res, next)
     }