From 46a0a064d0e6cebe7e17321d4e7f3650d7221019 Mon Sep 17 00:00:00 2001
From: Jerko Steiner <jerko.steiner@gmail.com>
Date: Tue, 27 Aug 2019 21:09:39 +0700
Subject: [PATCH] Make @ensure work with classes

---
 packages/jsonrpc/src/ensure.test.ts | 66 +++++++++++++++++++++++++++++
 packages/jsonrpc/src/ensure.ts      | 46 ++++++++++++++++++++
 packages/jsonrpc/src/jsonrpc.ts     | 42 ++++++++++++------
 3 files changed, 140 insertions(+), 14 deletions(-)
 create mode 100644 packages/jsonrpc/src/ensure.test.ts

diff --git a/packages/jsonrpc/src/ensure.test.ts b/packages/jsonrpc/src/ensure.test.ts
new file mode 100644
index 0000000..1d6cc23
--- /dev/null
+++ b/packages/jsonrpc/src/ensure.test.ts
@@ -0,0 +1,66 @@
+import {
+  ensure,
+  getValidatorsForMethod,
+  getValidatorsForInstance,
+  Validate,
+} from './ensure'
+
+describe('ensure', () => {
+
+  interface IContext {
+    userId: number
+  }
+
+  const validate: Validate<IContext> = c => c.userId > 0
+
+  it('decorates class methods', () => {
+    class Service {
+      @ensure<IContext>(validate)
+      // @ensureMethod<IContext>(validate)
+      fetchData() {
+        return 1
+      }
+      fetchData2() {
+        return 2
+      }
+    }
+    const s = new Service()
+    const validators1 = getValidatorsForMethod(s, 'fetchData')
+    const validators2 = getValidatorsForMethod(s, 'fetchData2')
+    expect(validators1).toEqual([ validate ])
+    expect(validators2).toEqual([])
+  })
+
+  it('decorates classes', () => {
+    @ensure(validate)
+    class Service {
+      fetchData() {
+        return 1
+      }
+      fetchData2() {
+        return 2
+      }
+    }
+    const s = new Service()
+    const validators = getValidatorsForInstance(s)
+    expect(validators).toEqual([ validate ])
+  })
+
+  it('works with subclasses', () => {
+    @ensure(validate)
+    class Service1 {
+      fetchData() {
+        return 1
+      }
+    }
+    class Service2 extends Service1 {
+      fetchData2() {
+        return 2
+      }
+    }
+    const s = new Service2()
+    const validators = getValidatorsForInstance(s)
+    expect(validators).toEqual([ validate ])
+  })
+
+})
diff --git a/packages/jsonrpc/src/ensure.ts b/packages/jsonrpc/src/ensure.ts
index 3dcec9f..177fec2 100644
--- a/packages/jsonrpc/src/ensure.ts
+++ b/packages/jsonrpc/src/ensure.ts
@@ -1,12 +1,46 @@
 import 'reflect-metadata'
 
 export const ensureKey = Symbol('ensure')
+export const ensureClassKey = Symbol('ensureClass')
 
 export type Validate<Context> = (context: Context) => boolean
 
 export function ensure<Context>(
   validate: Validate<Context>,
   message: string = 'Validation failed',
+) {
+  return function ensureImpl(
+    target: any, key?: string, descriptor?: PropertyDescriptor) {
+    switch (arguments.length) {
+      case 1:
+        return ensureClass(validate, message).apply(null, arguments as any)
+      case 3:
+        return ensureMethod(validate, message).apply(null, arguments as any)
+      default:
+        throw new Error('Unexpected number of arguments for @ensure. ' +
+          'It can only be used as as class or method decorator')
+    }
+  }
+}
+
+function ensureClass<Context>(
+  validate: Validate<Context>,
+  message: string = 'Validation failed',
+) {
+  // tslint:disable-next-line
+  return (target: any) => {
+    const validators: Array<Validate<Context>> =
+      getValidatorsForClass<Context>(target)
+
+    validators.push(validate)
+
+    Reflect.defineMetadata(ensureKey, validators, target)
+  }
+}
+
+function ensureMethod<Context>(
+  validate: Validate<Context>,
+  message: string = 'Validation failed',
 ) {
   return (
     target: any,
@@ -23,6 +57,18 @@ export function ensure<Context>(
   }
 }
 
+export function getValidatorsForInstance<Context>(
+  target: any,
+): Array<Validate<Context>> {
+  return getValidatorsForClass(target.__proto__.constructor)
+}
+
+export function getValidatorsForClass<Context>(
+  target: any,
+): Array<Validate<Context>> {
+  return Reflect.getMetadata(ensureKey, target) || []
+}
+
 export function getValidatorsForMethod<Context>(
   target: any,
   method: string,
diff --git a/packages/jsonrpc/src/jsonrpc.ts b/packages/jsonrpc/src/jsonrpc.ts
index 12c189c..fb54326 100644
--- a/packages/jsonrpc/src/jsonrpc.ts
+++ b/packages/jsonrpc/src/jsonrpc.ts
@@ -2,7 +2,8 @@ export type TId = number | string
 import {ArgumentTypes, FunctionPropertyNames, RetType} from './types'
 import {isPromise} from './isPromise'
 import {createError, IErrorResponse, IErrorWithData} from './error'
-import {getValidatorsForMethod} from './ensure'
+import {getValidatorsForMethod, getValidatorsForInstance} from './ensure'
+import {Validate} from './ensure'
 
 export const ERROR_PARSE = {
   code: -32700,
@@ -70,6 +71,31 @@ export function createSuccessResponse<T>(id: number | string, result: T)
   }
 }
 
+function validateServiceContext<T, M extends FunctionPropertyNames<T>, Context>(
+  id: string | number | null,
+  service: T,
+  method: M,
+  context: Context,
+) {
+
+  function doValidate(validate: Validate<Context>) {
+    const success = validate(context)
+    if (!success) {
+      throw createError(ERROR_INVALID_REQUEST, {
+        id,
+        data: null,
+        statusCode: 400,
+      })
+    }
+  }
+
+  getValidatorsForInstance<Context>(service)
+  .forEach(doValidate)
+
+  getValidatorsForMethod<Context>(service, method)
+  .forEach(doValidate)
+}
+
 export const createRpcService = <T, M extends FunctionPropertyNames<T>>(
   service: T,
   methods: M[],
@@ -107,19 +133,7 @@ export const createRpcService = <T, M extends FunctionPropertyNames<T>>(
         })
       }
 
-      const validators = getValidatorsForMethod<Context>(
-        (service as any), method)
-
-      validators.forEach(validate => {
-        const success = validate(context)
-        if (!success) {
-          throw createError(ERROR_INVALID_REQUEST, {
-            id,
-            data: null,
-            statusCode: 400,
-          })
-        }
-      })
+      validateServiceContext(id, service, method, context)
 
       let retValue = (rpcService[method] as any)(...params, context)