sprinklers3/server/sprinklersRpc/WebSocketConnection.ts

import { autorun } from "mobx";
import { serialize } from "serializr";
import * as WebSocket from "ws";

import { ErrorCode } from "@common/ErrorCode";
import * as rpc from "@common/jsonRpc";
import log from "@common/logger";
import { RpcError } from "@common/sprinklersRpc";
import * as deviceRequests from "@common/sprinklersRpc/deviceRequests";
import * as schema from "@common/sprinklersRpc/schema";
import * as ws from "@common/sprinklersRpc/websocketData";
import { AccessToken } from "@common/TokenClaims";
import { verifyToken } from "@server/authentication";
import { User } from "@server/entities";

import { WebSocketApi } from "./WebSocketApi";

type Disposer = () => void;

export class WebSocketConnection {
    api: WebSocketApi;
    socket: WebSocket;

    disposers: Array<() => void> = [];
    // map of device id to disposer function
    deviceSubscriptions: Map<string, Disposer> = new Map();

    /// This shall be the user id if the client has been authenticated, null otherwise
    userId: number | null = null;
    user: User | null = null;

    private requestHandlers: ws.ClientRequestHandlers = new WebSocketRequestHandlers();

    get state() {
        return this.api.state;
    }

    constructor(api: WebSocketApi, socket: WebSocket) {
        this.api = api;
        this.socket = socket;

        this.socket.on("message", this.handleSocketMessage);
        this.socket.on("close", this.onClose);
    }

    stop = () => {
        this.socket.close();
    }

    onClose = (code: number, reason: string) => {
        log.debug({ code, reason }, "WebSocketConnection closing");
        this.disposers.forEach((disposer) => disposer());
        this.deviceSubscriptions.forEach((disposer) => disposer());
        this.api.removeClient(this);
    }

    subscribeBrokerConnection() {
        this.disposers.push(autorun(() => {
            const updateData: ws.IBrokerConnectionUpdate = {
                brokerConnected: this.state.mqttClient.connected,
            };
            this.sendNotification("brokerConnectionUpdate", updateData);
        }));
    }

    checkAuthorization() {
        if (!this.userId || !this.user) {
            throw new RpcError("this WebSocket session has not been authenticated",
                ErrorCode.Unauthorized);
        }
    }

    checkDevice(devId: string) {
        const userDevice = this.user!.devices!.find((dev) => dev.deviceId === devId);
        if (userDevice == null) {
            throw new RpcError("you do not have permission to subscribe to device",
                ErrorCode.NoPermission, { id: devId });
        }
        const deviceId = userDevice.deviceId;
        if (!deviceId) {
            throw new RpcError("device has no associated device prefix", ErrorCode.Internal);
        }
        return userDevice;
    }

    sendMessage(data: ws.ServerMessage) {
        this.socket.send(JSON.stringify(data));
    }

    sendNotification<Method extends ws.ServerNotificationMethod>(
        method: Method,
        data: ws.IServerNotificationTypes[Method]) {
        this.sendMessage({ type: "notification", method, data });
    }

    sendResponse<Method extends ws.ClientRequestMethods>(
        method: Method,
        id: number,
        data: ws.ServerResponseData<Method>) {
        this.sendMessage({ type: "response", method, id, ...data });
    }

    handleSocketMessage = (socketData: WebSocket.Data) => {
        this.doHandleSocketMessage(socketData)
            .catch((err) => {
                this.onError({ err }, "unhandled error on handling socket message");
            });
    }

    async doDeviceCallRequest(requestData: ws.IDeviceCallRequest): Promise<deviceRequests.Response> {
        const userDevice = this.checkDevice(requestData.deviceId);
        const deviceId = userDevice.deviceId!;
        const device = this.state.mqttClient.acquireDevice(deviceId);
        try {
            const request = schema.requests.deserializeRequest(requestData.data);
            return await device.makeRequest(request);
        } finally {
            device.release();
        }
    }

    private async doHandleSocketMessage(socketData: WebSocket.Data) {
        if (typeof socketData !== "string") {
            return this.onError({ type: typeof socketData },
                "received invalid socket data type from client", ErrorCode.Parse);
        }
        let data: ws.ClientMessage;
        try {
            data = JSON.parse(socketData);
        } catch (err) {
            return this.onError({ socketData, err }, "received invalid websocket message from client",
                ErrorCode.Parse);
        }
        switch (data.type) {
            case "request":
                await this.handleRequest(data);
                break;
            default:
                return this.onError({ data }, "received invalid message type from client",
                    ErrorCode.BadRequest);
        }
    }

    private async handleRequest(request: ws.ClientRequest) {
        let response: ws.ServerResponseData;
        try {
            if (!this.requestHandlers[request.method]) {
                // noinspection ExceptionCaughtLocallyJS
                throw new RpcError("received invalid client request method");
            }
            response = await rpc.handleRequest(this.requestHandlers, request, this);
        } catch (err) {
            if (err instanceof RpcError) {
                log.debug({ err }, "rpc error");
                response = { result: "error", error: err.toJSON() };
            } else {
                log.error({ method: request.method, err }, "unhandled error during processing of client request");
                response = {
                    result: "error", error: {
                        code: ErrorCode.Internal, message: "unhandled error during processing of client request",
                        data: err.toString(),
                    },
                };
            }
        }
        this.sendResponse(request.method, request.id, response);
    }

    private onError(data: any, message: string, code: number = ErrorCode.Internal) {
        log.error(data, message);
        const errorData: ws.IError = { code, message, data };
        this.sendNotification("error", errorData);
    }
}

class WebSocketRequestHandlers implements ws.ClientRequestHandlers {
    async authenticate(this: WebSocketConnection, data: ws.IAuthenticateRequest):
        Promise<ws.ServerResponseData<"authenticate">> {
        if (!data.accessToken) {
            throw new RpcError("no token specified", ErrorCode.BadRequest);
        }
        let claims: AccessToken;
        try {
            claims = await verifyToken<AccessToken>(data.accessToken, "access");
        } catch (e) {
            throw new RpcError("invalid token", ErrorCode.BadToken, e);
        }
        this.userId = claims.aud;
        this.user = await this.state.database.users.
            findById(this.userId, { devices: true }) || null;
        if (!this.user) {
            throw new RpcError("user no longer exists", ErrorCode.BadToken);
        }
        log.debug({ userId: claims.aud, name: claims.name }, "authenticated websocket client");
        this.subscribeBrokerConnection();
        return {
            result: "success",
            data: { authenticated: true, message: "authenticated", user: this.user.toJSON() },
        };
    }

    async deviceSubscribe(this: WebSocketConnection, data: ws.IDeviceSubscribeRequest):
        Promise<ws.ServerResponseData<"deviceSubscribe">> {
        this.checkAuthorization();
        const userDevice = this.checkDevice(data.deviceId);
        const deviceId = userDevice.deviceId!;
        if (!this.deviceSubscriptions.has(deviceId)) {
            const device = this.state.mqttClient.acquireDevice(deviceId);
            log.debug({ deviceId, userId: this.userId }, "websocket client subscribed to device");

            const autorunDisposer = autorun(() => {
                const json = serialize(schema.sprinklersDevice, device);
                log.trace({ device: json });
                const updateData: ws.IDeviceUpdate = { deviceId, data: json };
                this.sendNotification("deviceUpdate", updateData);
            }, { delay: 100 });

            this.deviceSubscriptions.set(deviceId, () => {
                autorunDisposer();
                device.release();
                this.deviceSubscriptions.delete(deviceId);
            });
        }

        const response: ws.IDeviceSubscribeResponse = {
            deviceId,
        };
        return { result: "success", data: response };
    }

    async deviceUnsubscribe(this: WebSocketConnection, data: ws.IDeviceSubscribeRequest):
        Promise<ws.ServerResponseData<"deviceUnsubscribe">> {
        this.checkAuthorization();
        const userDevice = this.checkDevice(data.deviceId);
        const deviceId = userDevice.deviceId!;
        const disposer = this.deviceSubscriptions.get(deviceId);

        if (disposer) {
            disposer();
        }

        const response: ws.IDeviceSubscribeResponse = {
            deviceId,
        };
        return { result: "success", data: response };
    }

    async deviceCall(this: WebSocketConnection, data: ws.IDeviceCallRequest):
        Promise<ws.ServerResponseData<"deviceCall">> {
        this.checkAuthorization();
        try {
            const response = await this.doDeviceCallRequest(data);
            const resData: ws.IDeviceCallResponse = {
                data: response,
            };
            return { result: "success", data: resData };
        } catch (err) {
            const e: deviceRequests.ErrorResponseData = err;
            throw new RpcError(e.message, e.code, e);
        }
    }
}