rpclibrary/src/backend/RPCSocketServer.ts

import http = require('http');
import bsock = require('bsock');

import * as uuid from "uuid/v4"
import { Socket } from "./RPCSocketServer"

type rpcType = 'hook' | 'unhook' | 'call'
export type Outcome = "Success" | "Error"
export type Visibility = "127.0.0.1" | "0.0.0.0"

/* Responses */
export class Response{
    constructor(
        public message?:string
    ){}
}

export class SuccessResponse extends Response{
    result:Outcome = "Success"

    constructor(
        message?:string
    ){
        super(message)
    }
}

export class ErrorResponse extends Response{
    result:Outcome = "Error"

    constructor(
        message: string = "Unknown error"
    ){
        super(message)
    } 
}

export class SubscriptionResponse extends SuccessResponse{
    constructor(
        public uid: string,
        message?:string
    ){
        super(message)
    }
}
export type UnhookFunction = (uid:string) => Promise<SuccessResponse | ErrorResponse>
export type callbackFunction = (...args) => Promise<SubscriptionResponse | ErrorResponse>
export type AsyncFunction = (...args) => Promise<any>

export interface RPCExporter{
    name: string
    exportRPCs() : socketioRPC[]
    exportPublicRPCs() : socketioRPC[]
}

type baseRPC = { 
    type: rpcType
    name: string
}

type hookRPC = baseRPC & {
    type: 'hook'
    func: callbackFunction
    unhook: UnhookFunction
}

type unhookRPC = baseRPC & {
    type: 'unhook'
    func: UnhookFunction
}

type callRPC = baseRPC & {
    type: 'call'
    func: (...args) => Promise<any>
}

export type socketioRPC = callRPC | unhookRPC | hookRPC

export type baseInfo = {
    owner: string,
    argNames: string[],
}

type HookInfo = baseRPC & baseInfo & { 
    type: 'hook', 
    generator: (socket) => callbackFunction
    unhook: UnhookFunction
}

type UnhookInfo = baseRPC & baseInfo & {
    type: 'unhook',
    func: UnhookFunction
}

type CallInfo = baseRPC & baseInfo & {
    type: 'call',
    func: AsyncFunction
}

type RpcInfo = HookInfo | UnhookInfo | CallInfo

export type ExtendedRpcInfo = RpcInfo & { uniqueName: string } 

export const rpcToRpcinfo = (rpc : socketioRPC, owner: string):RpcInfo => {
    switch(rpc.type){
        case "call" :
            return {
                owner: owner,
                argNames: extractArgs(rpc.func),
                type: rpc.type,
                name: rpc.name,
                func: rpc.func,
            }
        case "unhook" :
            return {
                owner: owner,
                argNames: extractArgs(rpc.func),
                type: rpc.type,
                name: rpc.name,
                func: rpc.func,
            }
        case "hook" :
            const generator = hookGenerator(rpc)
            return {
                owner: owner,
                argNames: extractArgs(generator(undefined)),
                type: rpc.type,
                name: rpc.name,
                unhook: rpc.unhook,
                generator: generator,
            }
    }
}


function rpcHooker(socket: Socket, exporter:RPCExporter, makeUnique = true):ExtendedRpcInfo[]{
    const owner = exporter.name
    const RPCs = [...exporter.exportPublicRPCs(), ...exporter.exportRPCs()]
    const suffix = makeUnique?"-"+uuid().substr(0,4):""
    return RPCs.map(rpc => rpcToRpcinfo(rpc, owner))
    .map(info => {
        const ret:any = info
        ret.uniqueName = info.name+suffix

        switch(info.type){
            case "hook":
                socket.hook(ret.uniqueName, info.generator(socket))
                break;
            default:
                socket.hook(ret.uniqueName, info.func)
        }
        socket.on('close', () => socket.unhook(info.name))
        return ret
    })
}

const hookGenerator = (rpc:hookRPC): HookInfo['generator'] => { 
    const argsArr = extractArgs(rpc.func)
    argsArr.pop()
    const args = argsArr.join(',')

    return eval(`(socket) => async (`+args+`) => { 
        const res = await rpc.func(`+args+(args.length!==0?',':'')+` (x) => {
            socket.call(res.uid, x)
        })
        if(res.result == 'Success'){
            socket.on('close', async () => {
                const unhookRes = await rpc.unhook(res.uid)
                console.log("Specific close handler for", rpc.name, res.uid, unhookRes)
            })

        }
        return res
    }`)
}

const extractArgs = (f:Function):string[] => {
    let fn = String(f)
    let args = fn.substr(0, fn.indexOf(")"))
    args = args.substr(fn.indexOf("(")+1)
    let ret = args.split(",")
    return ret
}


type OnFunction = (type: 'error' | 'close', f: (e?:any)=>void) => Socket

export interface Socket {
    port: number
    hook: (rpcname: string, ...args: any[]) => Socket
    unhook: (rpcname:string) => Socket
    call: (rpcname:string, ...args: any[]) => Promise<any>
    fire: (rpcname:string, ...args: any[]) => Promise<any>
    on: OnFunction
    destroy: ()=>void
    close: ()=>void
}

export type RPCSocketConf = {
    connectionHandler: (socket:Socket) => void
    errorHandler: (socket:Socket) => (error:any) => void
    closeHandler: (socket:Socket) => () =>  void
}

export class RPCSocketServer{
    
    private io = bsock.createServer()
    private wsServer = http.createServer()

    constructor(
        private port:number,
        private rpcExporters: RPCExporter[] = [],
        private visibility: Visibility = "127.0.0.1",
        private conf: RPCSocketConf = {
            errorHandler: (socket:Socket) => (error:any) => { socket.destroy(); console.error(error) },
            closeHandler: (socket:Socket) => () => { console.log("Socket closing") },
            connectionHandler: (socket:Socket) => { console.log("New websocket connection in port "+socket.port) }
        }
    ){
        this.startWebsocket()
    }

    private startWebsocket(){
        try{
            this.io.attach(this.wsServer)
            this.io.on('socket', (socket:Socket) => {
                socket.on('error', this.conf.errorHandler(socket))
                socket.on('close', this.conf.closeHandler(socket))
                if(this.visibility === "127.0.0.1")
                    this.initRPCs(socket)
                else
                    this.initPublicRPCs(socket)
            })
            this.wsServer.listen(this.port, this.visibility)
        }catch(e){
            //@ts-ignore
            this.errorHandler(undefined)("Unable to connect to socket")
        }
    }

    protected initRPCs(socket:Socket){
        socket.hook('info', () => rpcInfos)

        const rpcInfos:ExtendedRpcInfo[] = [
            ...this.rpcExporters.flatMap(exporter => rpcHooker(socket, exporter))
        ]
    }

    protected initPublicRPCs(socket:Socket){
        socket.hook('info', () => rpcInfos)

        const rpcInfos:ExtendedRpcInfo[] = [
            ...this.rpcExporters.flatMap(exporter => rpcHooker(socket, exporter))
        ]
    }
}