feat: add tests for RPC protected methods and enhance RPCSession method protection

__tests__/e2e/rpc-disconnected.test.ts

@@ -0,0 +1,40 @@
+import { RPCError, RPCErrorCode } from "@/core/RPCError";
+import { RPCHandler } from "@/index"
+import { getRandomAvailablePort } from "@/utils/utils";
+
+type serverProvider = { test: () => Promise<string> };
+
+describe('Rpc disconnected test', () => {
+    test('main', async () => {
+        const port = await getRandomAvailablePort();
+
+        const server = new RPCHandler();
+        server.setProvider<serverProvider>({
+            test() {
+                return new Promise<string>((resolve) => setTimeout(() => resolve('ok'), 1000))
+            },
+        })
+        await server.listen({
+            port,
+        });
+
+        const client = new RPCHandler();
+        const session = await client.connect({
+            url: `http://localhost:${port}`,
+        });
+        const api = session.getAPI<serverProvider>();
+
+        const callPromise = api.test()
+
+        await session.connection.close();
+        await expect(api.test()).rejects.toMatchObject(
+            expect.objectContaining({
+                constructor: RPCError,
+                errorCode: RPCErrorCode.CONNECTION_DISCONNECTED
+            })
+        );
+        await expect(callPromise).rejects.toBeInstanceOf(RPCError);
+        await expect(callPromise).rejects
+            .toHaveProperty('errorCode', RPCErrorCode.CONNECTION_DISCONNECTED);
+    })
+})

__tests__/e2e/rpc-protected-method.test.ts

@@ -0,0 +1,63 @@
+import { isPublicMethod, markAsPublicMethod, publicMethod, RPCErrorCode, RPCHandler } from "@/index"
+
+describe('Rpc protected method test', () => {
+    test('disabled protection', async () => {
+        class Methods {
+            @publicMethod
+            allow() {
+                return 0;
+            }
+            disabled() { }
+        }
+
+        const classMethods = new Methods();
+        const provider = {
+            classMethods,
+            normal: markAsPublicMethod(() => 0),
+            normal2: markAsPublicMethod(function () { return 0 }),
+            normalProtected: function () { },
+            shallowObj: markAsPublicMethod({
+                fn1: () => 0,
+                l1: {
+                    fn1: () => 0,
+                }
+            }),
+            deepObj: markAsPublicMethod({
+                fn1: () => 0,
+                l1: {
+                    fn1: () => 0,
+                }
+            }, { deep: true }),
+        }
+
+        const server = new RPCHandler({
+            enableMethodProtection: true,
+        });
+        server.setProvider(provider);
+        await server.listen({
+            port: 5210
+        });
+
+
+        const client = new RPCHandler();
+        const session = await client.connect({
+            url: 'http://localhost:5210'
+        });
+        const api = session.getAPI<typeof provider>();
+        await expect(api.classMethods.allow()).resolves.toBe(0);
+        await expect(api.classMethods.disabled()).rejects
+            .toHaveProperty('errorCode', RPCErrorCode.METHOD_PROTECTED);
+        await expect(api.normal()).resolves.toBe(0);
+        await expect(api.normal2()).resolves.toBe(0);
+        await expect(api.normalProtected()).rejects
+            .toHaveProperty('errorCode', RPCErrorCode.METHOD_PROTECTED);
+
+        await expect(api.shallowObj.fn1()).resolves.toBe(0);
+        await expect(api.shallowObj.l1.fn1()).rejects
+            .toHaveProperty('errorCode', RPCErrorCode.METHOD_PROTECTED);
+
+        await expect(api.deepObj.fn1()).resolves.toBe(0);
+        await expect(api.deepObj.l1.fn1()).resolves.toBe(0);
+    })
+
+})

__tests__/unit/utils/utils.test.ts

@@ -1,4 +1,4 @@
-import { getRandomAvailablePort, isObject, isString, makeId } from "@/utils/utils"
+import { getRandomAvailablePort, isObject, isPublicMethod, isString, makeId, markAsPublicMethod } from "@/utils/utils"
 
 test('makeId', () => {
     const id = makeId();
@@ -26,4 +26,29 @@ test('getRandomAvailablePort', async () => {
     const port = await getRandomAvailablePort();
     expect(port).toBeGreaterThanOrEqual(1);
     expect(port).toBeLessThanOrEqual(65535);
+})
+
+test('markAsPublick', () => {
+    const shallowObj = {
+        fn1() { },
+        l1: {
+            fn1() { },
+        }
+    };
+
+    const deepObj = {
+        fn1() { },
+        l1: {
+            fn1() { },
+        }
+    };
+
+    markAsPublicMethod(shallowObj);
+    markAsPublicMethod(deepObj, { deep: true });
+
+    expect(isPublicMethod(shallowObj.fn1)).toBeTruthy();
+    expect(isPublicMethod(shallowObj.l1.fn1)).toBeUndefined();
+
+    expect(isPublicMethod(deepObj.fn1)).toBeTruthy();
+    expect(isPublicMethod(deepObj.l1.fn1)).toBeTruthy();
 })

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tonecn/typesrpc",
-  "version": "1.0.1",
+  "version": "1.0.2",
   "description": "A lightweight, type-safe RPC framework for TypeScript with deep nested API support",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

src/core/RPCClient.ts

@@ -85,6 +85,7 @@ export class RPCClient {
                     resolve(new RPCSession(
                         new RPCConnection(connection!),
                         this.rpcHandler,
+                        this,
                     ));
                 } else {
                     reject(new Error('Server rejected handshake request'));

src/core/RPCConnection.ts

@@ -1,9 +1,6 @@
 import { EventEmitter } from "@/utils/EventEmitter";
 import { SocketConnection } from "./SocketConnection";
 import { RPCPacket } from "./RPCPacket";
-import { makeCallPacket, makeCallResponsePacket, parseCallPacket, parseCallResponsePacket } from "./RPCCommon";
-import { RPCProvider } from "./RPCProvider";
-import { RPCError, RPCErrorCode } from "./RPCError";
 
 interface RPCConnectionEvents {
     call: RPCPacket;
@@ -14,33 +11,15 @@ interface RPCConnectionEvents {
     closed: void;
 }
 
-class CallResponseEmitter extends EventEmitter<{
-    [id: string]: RPCPacket;
-}> {
-    emitAll(packet: RPCPacket) {
-        this.events.forEach(subscribers => {
-            subscribers.forEach(fn => fn(packet));
-        })
-    }
-}
-
 export class RPCConnection extends EventEmitter<RPCConnectionEvents> {
 
     closed: boolean = false;
 
-    private callResponseEmitter = new CallResponseEmitter();
-
     constructor(public socket: SocketConnection) {
         super();
         socket.on('closed', () => {
-            this.emit('closed');
-            this.callResponseEmitter.emitAll(makeCallResponsePacket({
-                status: 'error',
-                requestPacketId: 'connection error',
-                errorCode: RPCErrorCode.CONNECTION_DISCONNECTED,
-            }));
-            this.callResponseEmitter.removeAllListeners();
             this.closed = true;
+            this.emit('closed');
         });
 
         socket.on('msg', (msg) => {
@@ -68,156 +47,17 @@ export class RPCConnection extends EventEmitter<RPCConnectionEvents> {
 
             this.emit('unknownPacket', packet);
         });
-
-        /** route by packet.id */
-        this.on('callResponse', (packet) => {
-            this.callResponseEmitter.emit(packet.id, packet);
-        })
     }
 
-    /** @throws */
-    public async callRequest(options: {
-        fnPath: string;
-        args: any[];
-        timeout: number;
-    }): Promise<any> {
+    public async close() {
+        return this.socket.close();
+    }
+
+    public async send(data: RPCPacket) {
         if (this.closed) {
-            throw new RPCError({
-                errorCode: RPCErrorCode.CONNECTION_DISCONNECTED,
-            });
+            return;
         }
 
-        const { fnPath, args } = options;
-        const packet = makeCallPacket({
-            fnPath,
-            args
-        });
-
-        let resolve: (data: any) => void;
-        let reject: (data: any) => void;
-        const promise = new Promise((res, rej) => {
-            resolve = res;
-            reject = rej;
-        });
-
-        const cancelTimeoutTimer = (() => {
-            const t = setTimeout(() => {
-                reject(new RPCError({
-                    errorCode: RPCErrorCode.TIMEOUT_ERROR,
-                }))
-            }, options.timeout);
-
-            return () => clearTimeout(t);
-        })();
-
-        promise.finally(() => {
-            this.callResponseEmitter.removeAllListeners(packet.id);
-            cancelTimeoutTimer();
-        })
-
-        const handleCallResponsePacket = (packet: RPCPacket) => {
-            const result = parseCallResponsePacket(packet);
-            if (result === null) {
-                return reject(new RPCError({
-                    errorCode: RPCErrorCode.UNKNOWN_ERROR,
-                }));;
-            }
-
-            const { success, error } = result;
-            if (success) {
-                return resolve(success.data);
-            }
-
-            if (error) {
-                return reject(new RPCError({
-                    errorCode: error.errorCode,
-                    reason: error.reason
-                }));
-            }
-
-            return reject(new RPCError({
-                errorCode: RPCErrorCode.UNKNOWN_ERROR,
-            }));;
-        }
-        this.callResponseEmitter.on(packet.id, handleCallResponsePacket);
-
-        /** send call request */
-        this.socket.send(packet);
-
-        return promise;
-    }
-
-    public onCallRequest(getProvider: () => RPCProvider | undefined) {
-        this.on('call', async (packet) => {
-            const request = parseCallPacket(packet);
-            if (request === null) {
-                return this.socket.send(makeCallResponsePacket({
-                    status: 'error',
-                    requestPacket: packet,
-                    errorCode: RPCErrorCode.CALL_PROTOCOL_ERROR,
-                })).catch(() => { })
-            }
-
-            // call the function
-            const provider = getProvider();
-            if (!provider) {
-                return this.socket.send(makeCallResponsePacket({
-                    status: 'error',
-                    requestPacket: packet,
-                    errorCode: RPCErrorCode.PROVIDER_NOT_AVAILABLE,
-                }))
-            }
-
-            const { fnPath, args } = request;
-            const fn = this.getProviderFunction(provider, fnPath);
-            if (!fn) {
-                return this.socket.send(makeCallResponsePacket({
-                    status: 'error',
-                    requestPacket: packet,
-                    errorCode: RPCErrorCode.METHOD_NOT_FOUND,
-                }))
-            }
-
-            try {
-                const result = await fn(...args);
-                this.socket.send(makeCallResponsePacket({
-                    status: 'success',
-                    requestPacket: packet,
-                    data: result,
-                }))
-            } catch (error) {
-                this.socket.send(makeCallResponsePacket({
-                    status: 'error',
-                    requestPacket: packet,
-                    errorCode: RPCErrorCode.SERVER_ERROR,
-                    ...(error instanceof RPCError ? {
-                        errorCode: error.errorCode,
-                        reason: error.reason,
-                    } : {})
-                }))
-            }
-        })
-    }
-
-    private getProviderFunction(provider: RPCProvider, fnPath: string) {
-        const paths = fnPath.split(':');
-        let fnThis: any = provider;
-        let fn: any = provider;
-        try {
-            while (paths.length) {
-                const path = paths.shift()!;
-                fn = fn[path];
-                if (paths.length !== 0) {
-                    fnThis = fn;
-                }
-            }
-            if (typeof fn === 'function') {
-                return fn.bind(fnThis);
-            }
-
-            throw new Error();
-        } catch (error) {
-            return null;
-        }
+        return this.socket.send(data);
     }
 }

src/core/RPCHandler.ts

@@ -19,28 +19,47 @@ interface RPCHandlerEvents {
     connect: RPCSession;
 }
 
+interface RPCConfig {
+    enableMethodProtection: boolean;
+}
+
+let DefaultRPCConfig: RPCConfig = {
+    enableMethodProtection: false,
+}
+
 export class RPCHandler extends EventEmitter<RPCHandlerEvents> {
 
     private rpcClient?: RPCClient;
     private rpcServer?: RPCServer;
     private provider?: RPCProvider;
     private accessKey?: string;
+    private config: RPCConfig;
 
     constructor(
         args?: {
             rpcClient?: RPCClient;
             rpcServer?: RPCServer;
-        }
+        } & Partial<RPCConfig>
     ) {
         super();
+        const { rpcClient, rpcServer, ...config } = args ?? {};
 
-        if (args?.rpcClient) {
-            this.setRPCProvider(args.rpcClient);
+        if (rpcClient) {
+            this.setRPCProvider(rpcClient);
         }
 
-        if (args?.rpcServer) {
-            this.setRPCProvider(args.rpcServer);
+        if (rpcServer) {
+            this.setRPCProvider(rpcServer);
         }
+
+        this.config = {
+            ...DefaultRPCConfig,
+            ...config,
+        }
+    }
+
+    getConfig() {
+        return this.config;
     }
 
     setProvider<T extends RPCProvider>(provider: T) {
@@ -116,27 +135,4 @@ export class RPCHandler extends EventEmitter<RPCHandlerEvents> {
             throw new Error();
         }
     }
-}
-
-
-// const h = new RPCHandler();
-
-// h.setProvider<{
-//     plus: (a: number, b: number) => number;
-//     math: {
-//         minus: (a: number, b: number) => number;
-//         multiply: (a: number, b: number) => number;
-//     }
-// }>({
-//     plus(a, b) {
-//         return a + b
-//     },
-//     math: {
-//         minus(a, b) {
-//             return a - b;
-//         },
-//         multiply(a, b) {
-//             return a * b;
-//         },
-//     }
-// })
+}

src/core/RPCServer.ts

@@ -78,6 +78,7 @@ export class RPCServer extends EventEmitter<RPCServerEvents> {
                     this.emit('connect', new RPCSession(
                         new RPCConnection(socketConnection),
                         this.rpcHandler,
+                        this,
                     ));
                 }
 

src/core/RPCSession.ts

@@ -1,16 +1,77 @@
-import { ToDeepPromise } from "@/utils/utils";
+import { isPublicMethod, ToDeepPromise } from "@/utils/utils";
 import { RPCConnection } from "./RPCConnection";
 import { RPCHandler } from "./RPCHandler";
 import { RPCProvider } from "./RPCProvider";
+import { RPCClient } from "./RPCClient";
+import { RPCServer } from "./RPCServer";
+import { RPCError, RPCErrorCode } from "./RPCError";
+import { makeCallPacket, makeCallResponsePacket, parseCallPacket, parseCallResponsePacket } from "./RPCCommon";
+import { RPCPacket } from "./RPCPacket";
+import { EventEmitter } from "@/utils/EventEmitter";
 
+function getProviderFunction(provider: RPCProvider, fnPath: string):
+    [(...args: any[]) => Promise<any>, object] | null {
+    const paths = fnPath.split(':');
+    let fnThis: any = provider;
+    let fn: any = provider;
+    try {
+        while (paths.length) {
+            const path = paths.shift()!;
+            fn = fn[path];
+            if (paths.length !== 0) {
+                fnThis = fn;
+            }
+        }
+        if (typeof fn === 'function') {
+            return [fn, fnThis];
+        }
+
+        throw new Error();
+    } catch (error) {
+        return null;
+    }
+}
+
+class CallResponseEmitter extends EventEmitter<{
+    [id: string]: RPCPacket;
+}> {
+    emitAll(packet: RPCPacket) {
+        this.events.forEach(subscribers => {
+            subscribers.forEach(fn => fn(packet));
+        })
+    }
+}
 
 export class RPCSession {
 
+    public callResponseEmitter = new CallResponseEmitter();
+
     constructor(
         public readonly connection: RPCConnection,
         public readonly rpcHandler: RPCHandler,
+        public readonly rpcProvider: RPCClient | RPCServer,
     ) {
-        connection.onCallRequest(rpcHandler.getProvider.bind(rpcHandler));
+        /** route by packet.id */
+        this.connection.on('callResponse', (packet) => {
+            this.callResponseEmitter.emit(packet.id, packet);
+        });
+
+        this.connection.on('closed', () => {
+            this.callResponseEmitter.emitAll(makeCallResponsePacket({
+                status: 'error',
+                requestPacketId: 'connection error',
+                errorCode: RPCErrorCode.CONNECTION_DISCONNECTED,
+            }));
+            this.callResponseEmitter.removeAllListeners();
+        });
+
+        this.connection.on('call', (packet) => {
+            this.onCallRequest(packet).then(res => {
+                this.connection.send(res);
+            }).catch((e) => {
+                console.warn(`${e}`);
+            })
+        });
     }
 
     getAPI<T extends RPCProvider>(): ToDeepPromise<T> {
@@ -23,7 +84,7 @@ export class RPCSession {
                     return createProxy(newPath);
                 },
                 apply: (target, thisArg, args) => {
-                    return this.connection.callRequest({
+                    return this.callRequest({
                         fnPath: path.join(':'),
                         args: args,
                         /** @todo accept from caller */
@@ -37,4 +98,136 @@ export class RPCSession {
 
         return createProxy() as unknown as ToDeepPromise<T>;
     }
+
+    /** @throws */
+    public async callRequest(options: {
+        fnPath: string;
+        args: any[];
+        timeout: number;
+    }): Promise<any> {
+        if (this.connection.closed) {
+            throw new RPCError({
+                errorCode: RPCErrorCode.CONNECTION_DISCONNECTED,
+            });
+        }
+
+        const { fnPath, args } = options;
+        const packet = makeCallPacket({
+            fnPath,
+            args
+        });
+
+        let resolve: (data: any) => void;
+        let reject: (data: any) => void;
+        const promise = new Promise((res, rej) => {
+            resolve = res;
+            reject = rej;
+        });
+
+        const cancelTimeoutTimer = (() => {
+            const t = setTimeout(() => {
+                reject(new RPCError({
+                    errorCode: RPCErrorCode.TIMEOUT_ERROR,
+                }))
+            }, options.timeout);
+
+            return () => clearTimeout(t);
+        })();
+
+        const handleCallResponsePacket = (packet: RPCPacket) => {
+            const result = parseCallResponsePacket(packet);
+            if (result === null) {
+                return reject(new RPCError({
+                    errorCode: RPCErrorCode.UNKNOWN_ERROR,
+                }));;
+            }
+
+            const { success, error } = result;
+            if (success) {
+                return resolve(success.data);
+            }
+
+            if (error) {
+                return reject(new RPCError({
+                    errorCode: error.errorCode,
+                    reason: error.reason
+                }));
+            }
+
+            return reject(new RPCError({
+                errorCode: RPCErrorCode.UNKNOWN_ERROR,
+            }));;
+        }
+        this.callResponseEmitter.once(packet.id, handleCallResponsePacket);
+
+        /** send call request */
+        this.connection.send(packet);
+
+        return promise.finally(() => {
+            this.callResponseEmitter.removeAllListeners(packet.id);
+            cancelTimeoutTimer();
+        });
+    }
+
+    private async onCallRequest(packet: RPCPacket): Promise<RPCPacket> {
+        const request = parseCallPacket(packet);
+        if (request === null) {
+            return makeCallResponsePacket({
+                status: 'error',
+                requestPacket: packet,
+                errorCode: RPCErrorCode.CALL_PROTOCOL_ERROR,
+            });
+        }
+
+        // call the function
+        const provider = this.rpcHandler.getProvider();
+        if (!provider) {
+            return makeCallResponsePacket({
+                status: 'error',
+                requestPacket: packet,
+                errorCode: RPCErrorCode.PROVIDER_NOT_AVAILABLE,
+            });
+        }
+
+        const { fnPath, args } = request;
+        const fnRes = getProviderFunction(provider, fnPath);
+        if (!fnRes) {
+            return makeCallResponsePacket({
+                status: 'error',
+                requestPacket: packet,
+                errorCode: RPCErrorCode.METHOD_NOT_FOUND,
+            })
+        }
+        const [fn, fnThis] = fnRes;
+
+        const { enableMethodProtection } = this.rpcHandler.getConfig();
+        if (enableMethodProtection) {
+            if (!isPublicMethod(fn)) {
+                return makeCallResponsePacket({
+                    status: 'error',
+                    requestPacket: packet,
+                    errorCode: RPCErrorCode.METHOD_PROTECTED,
+                })
+            }
+        }
+
+        try {
+            const result = await fn.bind(fnThis)(...args);
+            return makeCallResponsePacket({
+                status: 'success',
+                requestPacket: packet,
+                data: result,
+            })
+        } catch (error) {
+            return makeCallResponsePacket({
+                status: 'error',
+                requestPacket: packet,
+                errorCode: RPCErrorCode.SERVER_ERROR,
+                ...(error instanceof RPCError ? {
+                    errorCode: error.errorCode,
+                    reason: error.reason,
+                } : {})
+            })
+        }
+    }
 }

src/index.ts

@@ -16,6 +16,8 @@ export { injectSocketClient } from "./core/SocketClient";
 export { injectSocketServer } from "./core/SocketServer";
 import { injectSocketIOImplements } from "./implements/socket.io";
 
+export { publicMethod, isPublicMethod, markAsPublicMethod } from './utils/utils';
+
 injectSocketIOImplements();
 
 export {

src/utils/utils.ts

@@ -4,13 +4,17 @@ export const makeId = () => md5(`${Date.now()}${Math.random()}`);
 
 export const isObject = (v: unknown): v is Record<string, any> => typeof v === 'object' && v !== null;
 
+export const isArray = (v: unknown): v is Array<unknown> => Array.isArray(v);
+
 export const isString = (v: unknown): v is string => typeof v === 'string';
 
+export const isFunction = (v: unknown): v is Function => typeof v === 'function';
+
 export type ObjectType = Record<string, any>;
 
 export type ToDeepPromise<T> = {
     [K in keyof T]: T[K] extends (...args: infer P) => infer R
-    ? (...args: P) => Promise<R>
+    ? (...args: P) => Promise<Awaited<R>>
     : T[K] extends object
     ? ToDeepPromise<T[K]>
     : T[K]
@@ -41,4 +45,40 @@ export async function getRandomAvailablePort() {
 
         server.listen(0);
     })
+}
+
+const publicMethodMap = new WeakMap<Function, boolean>();
+export function publicMethod(target: any, propertyKey: string, descriptor: PropertyDescriptor) {
+    publicMethodMap.set(descriptor.value, true);
+};
+export function isPublicMethod(target: Function) {
+    return publicMethodMap.get(target);
+};
+export function markAsPublicMethod<T extends Function | Record<any, unknown> | unknown>(obj: T, options?: {
+    deep?: boolean
+}): T {
+    const accessed = new Set();
+    function markAs(obj: Function | Record<any, unknown> | unknown) {
+        if (accessed.has(obj)) {
+            return;
+        }
+
+        if (isFunction(obj)) {
+            publicMethodMap.set(obj, true);
+        } else if (isObject(obj)) {
+            accessed.add(obj);
+
+            Object.values(obj).forEach(subObj => {
+                if (isFunction(subObj)) {
+                    publicMethodMap.set(subObj, true);
+                }
+                if (options?.deep && isObject(subObj)) {
+                    markAs(subObj);
+                }
+            });
+        }
+    }
+
+    markAs(obj);
+    return obj;
 }

tsconfig.json

@@ -12,8 +12,8 @@
     "target": "es2016", /* Set the JavaScript language version for emitted JavaScript and include compatible library declarations. */
     // "lib": [],                                        /* Specify a set of bundled library declaration files that describe the target runtime environment. */
     // "jsx": "preserve",                                /* Specify what JSX code is generated. */
-    // "experimentalDecorators": true,                   /* Enable experimental support for legacy experimental decorators. */
-    // "emitDecoratorMetadata": true,                    /* Emit design-type metadata for decorated declarations in source files. */
+    "experimentalDecorators": true,                   /* Enable experimental support for legacy experimental decorators. */
+    "emitDecoratorMetadata": true,                    /* Emit design-type metadata for decorated declarations in source files. */
     // "jsxFactory": "",                                 /* Specify the JSX factory function used when targeting React JSX emit, e.g. 'React.createElement' or 'h'. */
     // "jsxFragmentFactory": "",                         /* Specify the JSX Fragment reference used for fragments when targeting React JSX emit e.g. 'React.Fragment' or 'Fragment'. */
     // "jsxImportSource": "",                            /* Specify module specifier used to import the JSX factory functions when using 'jsx: react-jsx*'. */