diff --git a/apis/websocket/config.json b/apis/websocket/config.json index 9a07cd7..2b929d4 100755 --- a/apis/websocket/config.json +++ b/apis/websocket/config.json @@ -4,6 +4,5 @@ "address": "127.0.0.1", "port": 3000, "ocrConcurrentQueues": 1, - "clientHeartbeatInterval": 5000, - "consoleLogLevel": "log" + "logLevel": "debug" } diff --git a/apis/websocket/config.revanced.json b/apis/websocket/config.revanced.json index ef47645..ffb38a0 100755 --- a/apis/websocket/config.revanced.json +++ b/apis/websocket/config.revanced.json @@ -4,6 +4,5 @@ "address": "127.0.0.1", "port": 3000, "ocrConcurrentQueues": 3, - "clientHeartbeatInterval": 5000, - "consoleLogLevel": "log" + "logLevel": "log" } diff --git a/apis/websocket/config.schema.json b/apis/websocket/config.schema.json index 51965c1..a6343ef 100755 --- a/apis/websocket/config.schema.json +++ b/apis/websocket/config.schema.json @@ -17,12 +17,7 @@ "type": "integer", "default": 1 }, - "clientHeartbeatInterval": { - "description": "Time in milliseconds to wait for a client to send a heartbeat packet, if no packet is received, the server will wait for `clientHeartbeatExtraTime` milliseconds before disconnecting the client", - "type": "integer", - "default": 60000 - }, - "consoleLogLevel": { + "logLevel": { "description": "The log level to print to console", "type": "string", "enum": ["error", "warn", "info", "log", "debug", "trace", "none"], diff --git a/apis/websocket/docs/1_configuration.md b/apis/websocket/docs/1_configuration.md index aead5c5..54999e5 100644 --- a/apis/websocket/docs/1_configuration.md +++ b/apis/websocket/docs/1_configuration.md @@ -7,7 +7,6 @@ This is the default configuration (provided in [config.json](../config.json)): "address": "127.0.0.1", "port": 3000, "ocrConcurrentQueues": 1, - "clientHeartbeatInterval": 60000, "consoleLogLevel": "log" } ``` @@ -25,22 +24,18 @@ Amount of concurrent queues that can be run at a time. > [!WARNING] > Setting this too high may cause performance issues. -### `config.clientHeartbeatInterval` - -Heartbeat interval for clients. See [**💓 Heartbeating**](./3_packets.md#💓-heartbeating). - -### `config.consoleLogLevel` +### `config.logLevel` The level of logs to print to console. If the level is more important or equally important to set level, it will be forwarded to the console. The possible levels (sorted by their importance descendingly) are: +- `none` (no messages) - `fatal` - `error` - `warn` - `info` - `log` -- `trace` - `debug` ## ⏭️ What's next diff --git a/apis/websocket/docs/2_running.md b/apis/websocket/docs/2_running.md index 27d0909..d96f449 100644 --- a/apis/websocket/docs/2_running.md +++ b/apis/websocket/docs/2_running.md @@ -36,7 +36,7 @@ bun bundle ``` The files will be placed in the `dist` directory. **Configurations and `.env` files will NOT be copied automatically.** -You can run these files after using a runtime, eg. `bun run .` or `node .`. +You can run these files using the command `bun run index.js`. ## ⏭️ What's next diff --git a/apis/websocket/docs/3_packets.md b/apis/websocket/docs/3_packets.md index 49eb54e..e2d64b1 100644 --- a/apis/websocket/docs/3_packets.md +++ b/apis/websocket/docs/3_packets.md @@ -19,15 +19,13 @@ Operation codes are numbers that communicate an action. Data fields include additional information for the server to process. They are **either an object with specific fields or just `null`**. +### `packet.s` (server packets) + +A sequence number, exclusively for server packets. The WebSocket server contacts other APIs and they may not be reliable at all times, this makes race conditions. A sequence number cleanly solves this issue by letting the client know what the next packet sequence number would be by giving the current number. + #### 📦 Schemas and constants Schemas for packets and their respective data[^1], and the list of possible operation codes[^2] can be found in the `@revanced/bot-shared` package, with typings as well. [^1]: [`@revanced/bot-shared/src/schemas/Packet.ts`](../../../packages/shared/src/schemas/Packet.ts) [^2]: [`@revanced/bot-shared/src/constants/Operation`](../../../packages/shared/src/constants/Operation.ts) - -## 💓 Heartbeating - -Heartbeating is a process where the client regularly send each other signals to confirm that they are still connected and functioning. If the server doesn't receive a heartbeat from the client within a specified timeframe, it assume the client has disconnected and closes the socket. - -You can configure the interval in the configuration file. See [**📝 Configuration > `config.clientHeartbeatInterval`**](./1_configuration.md#configclientheartbeatinterval). diff --git a/apis/websocket/package.json b/apis/websocket/package.json index fe8eef9..e803cc2 100755 --- a/apis/websocket/package.json +++ b/apis/websocket/package.json @@ -1,40 +1,40 @@ { - "name": "@revanced/bot-websocket-api", - "type": "module", - "private": true, - "version": "0.1.0", - "description": "🧦 WebSocket API server for bots assisting ReVanced", - "main": "dist/index.js", - "scripts": { - "bundle": "bun build src/index.ts --outdir=dist --target=bun --minify --sourcemap=external", - "dev": "bun run src/index.ts --watch", - "build": "bun bundle", - "watch": "bun dev" - }, - "repository": { - "type": "git", - "url": "git+https://github.com/revanced/revanced-helper.git", - "directory": "apis/websocket" - }, - "author": "Palm (https://github.com/PalmDevs)", - "contributors": [ - "Palm (https://github.com/PalmDevs)", - "ReVanced (https://github.com/revanced)" - ], - "license": "GPL-3.0-or-later", - "bugs": { - "url": "https://github.com/revanced/revanced-helper/issues" - }, - "homepage": "https://github.com/revanced/revanced-helper#readme", - "dependencies": { - "@revanced/bot-shared": "workspace:*", - "@sapphire/async-queue": "^1.5.1", - "chalk": "^5.3.0", - "node-wit": "^6.6.0", - "tesseract.js": "^5.0.4" - }, - "devDependencies": { - "@types/node-wit": "^6.0.3", - "typed-emitter": "^2.1.0" - } + "name": "@revanced/bot-websocket-api", + "type": "module", + "private": true, + "version": "0.1.0", + "description": "🧦 WebSocket API server for bots assisting ReVanced", + "main": "dist/index.js", + "scripts": { + "bundle": "bun build src/index.ts --outdir=dist --target=bun", + "dev": "bun run src/index.ts --watch", + "build": "bun bundle", + "watch": "bun dev" + }, + "repository": { + "type": "git", + "url": "git+https://github.com/revanced/revanced-helper.git", + "directory": "apis/websocket" + }, + "author": "Palm (https://github.com/PalmDevs)", + "contributors": [ + "Palm (https://github.com/PalmDevs)", + "ReVanced (https://github.com/revanced)" + ], + "license": "GPL-3.0-or-later", + "bugs": { + "url": "https://github.com/revanced/revanced-helper/issues" + }, + "homepage": "https://github.com/revanced/revanced-helper#readme", + "dependencies": { + "@revanced/bot-shared": "workspace:*", + "@sapphire/async-queue": "^1.5.2", + "chalk": "^5.3.0", + "node-wit": "^6.6.0", + "tesseract.js": "^5.0.5" + }, + "devDependencies": { + "@types/node-wit": "^6.1.0", + "typed-emitter": "^2.1.0" + } } diff --git a/apis/websocket/src/classes/Client.ts b/apis/websocket/src/classes/Client.ts index 8025c5b..8eda488 100755 --- a/apis/websocket/src/classes/Client.ts +++ b/apis/websocket/src/classes/Client.ts @@ -2,7 +2,7 @@ import { EventEmitter } from 'events' import { ClientOperation, DisconnectReason, - Packet, + type Packet, ServerOperation, deserializePacket, isClientPacket, @@ -17,38 +17,30 @@ export default class Client { id: string disconnected: DisconnectReason | false = false ready = false + currentSequence = 0 - lastHeartbeat: number = null! - heartbeatInterval: number - - #hbTimeout: NodeJS.Timeout = null! #emitter = new EventEmitter() as TypedEmitter #socket: WebSocket constructor(options: ClientOptions) { this.#socket = options.socket - this.heartbeatInterval = options.heartbeatInterval ?? 60000 this.id = options.id - this.#socket.on('error', () => this.forceDisconnect()) - this.#socket.on('close', () => this.forceDisconnect()) - this.#socket.on('unexpected-response', () => this.forceDisconnect()) + this.#socket.on('error', () => this.disconnect(DisconnectReason.ServerError)) + this.#socket.on('close', code => this._handleDisconnect(code)) + this.#socket.on('unexpected-response', () => this.disconnect(DisconnectReason.InvalidPacket)) this.send({ op: ServerOperation.Hello, - d: { - heartbeatInterval: this.heartbeatInterval, - }, + d: null, }) .then(() => { - this.#listen() - this.#listenHeartbeat() + this._listen() this.ready = true this.#emitter.emit('ready') }) .catch(() => { - if (this.disconnected === false) this.disconnect(DisconnectReason.ServerError) - else this.forceDisconnect(DisconnectReason.ServerError) + this.disconnect(DisconnectReason.ServerError) }) } @@ -64,54 +56,40 @@ export default class Client { this.#emitter.off(name, handler) } - send(packet: Packet) { + send(packet: Omit, 's'>, sequence?: number) { return new Promise((resolve, reject) => { - try { - this.#throwIfDisconnected('Cannot send packet to client that has already disconnected') - this.#socket.send(serializePacket(packet)) - resolve() - } catch (e) { - reject(e) - } + this.#throwIfDisconnected('Cannot send packet to client that has already disconnected') + this.#socket.send( + serializePacket({ ...packet, s: sequence ?? this.currentSequence++ } as Packet), + err => (err ? reject(err) : resolve()), + ) }) } - async disconnect(reason: DisconnectReason = DisconnectReason.Generic) { + async disconnect(reason: DisconnectReason | number = DisconnectReason.Generic) { this.#throwIfDisconnected('Cannot disconnect client that has already disconnected') - try { - await this.send({ op: ServerOperation.Disconnect, d: { reason } }) - } catch (err) { - throw new Error(`Cannot send disconnect reason to client ${this.id}: ${err}`) - } finally { - this.forceDisconnect(reason) - } - } - - forceDisconnect(reason: DisconnectReason = DisconnectReason.Generic) { - if (this.disconnected !== false) return - - // It's so weird because if I moved this down a few lines - // it would just fire the disconnect event twice because of a race condition - this.disconnected = reason - this.ready = false - - if (this.#hbTimeout) clearTimeout(this.#hbTimeout) - this.#socket.close() - - this.#emitter.emit('disconnect', reason) + this.#socket.close(reason) + this._handleDisconnect(reason) } #throwIfDisconnected(errorMessage: string) { if (this.disconnected !== false) throw new Error(errorMessage) if (this.#socket.readyState !== this.#socket.OPEN) { - this.forceDisconnect(DisconnectReason.Generic) + this.#socket.close(DisconnectReason.NoOpenSocket) throw new Error(errorMessage) } } - #listen() { + protected _handleDisconnect(code: number) { + this.disconnected = code + this.ready = false + + this.#emitter.emit('disconnect', code) + } + + protected _listen() { this.#socket.on('message', data => { this.#emitter.emit('message', data) try { @@ -136,38 +114,6 @@ export default class Client { }) } - #listenHeartbeat() { - this.lastHeartbeat = Date.now() - this.#startHeartbeatTimeout() - - this.on('heartbeat', () => { - this.lastHeartbeat = Date.now() - this.#hbTimeout.refresh() - - this.send({ - op: ServerOperation.HeartbeatAck, - d: { - nextHeartbeat: this.lastHeartbeat + this.heartbeatInterval, - }, - }).catch(() => {}) - }) - } - - #startHeartbeatTimeout() { - this.#hbTimeout = setTimeout(() => { - if (Date.now() - this.lastHeartbeat > 0) { - // TODO: put into config - // 5000 is extra time to account for latency - const interval = setTimeout(() => this.disconnect(DisconnectReason.TimedOut), 5000) - - this.once('heartbeat', () => clearTimeout(interval)) - // This should never happen but it did in my testing so I'm adding this just in case - this.once('disconnect', () => clearTimeout(interval)) - // Technically we don't have to do this, but JUST IN CASE! - } else this.#hbTimeout.refresh() - }, this.heartbeatInterval) - } - protected _toBuffer(data: RawData) { if (data instanceof Buffer) return data if (data instanceof ArrayBuffer) return Buffer.from(data) @@ -178,7 +124,6 @@ export default class Client { export interface ClientOptions { id: string socket: WebSocket - heartbeatInterval?: number } export type ClientPacketObject = Packet & { diff --git a/apis/websocket/src/events/index.ts b/apis/websocket/src/events/index.ts index 6e81e88..53392b1 100755 --- a/apis/websocket/src/events/index.ts +++ b/apis/websocket/src/events/index.ts @@ -1,20 +1,33 @@ import type { ClientOperation } from '@revanced/bot-shared' import type { Logger } from '@revanced/bot-shared' -import type { Wit } from 'node-wit' import type { Worker as TesseractWorker } from 'tesseract.js' -import { ClientPacketObject } from '../classes/Client' -import type { Config } from '../utils/getConfig' +import type { ClientPacketObject } from '../classes/Client' +import type { Config } from '../utils/config' export { default as parseTextEventHandler } from './parseText' export { default as parseImageEventHandler } from './parseImage' +export { default as trainMessageEventHandler } from './trainMessage' export type EventHandler = ( packet: ClientPacketObject, context: EventContext, ) => void | Promise + export type EventContext = { - witClient: Wit - tesseractWorker: TesseractWorker + wit: { + train(text: string, label: string): Promise + message(text: string): Promise + } + tesseract: TesseractWorker logger: Logger config: Config } + +export interface WitMessageResponse { + text: string + intents: Array<{ + id: string + name: string + confidence: number + }> +} diff --git a/apis/websocket/src/events/parseImage.ts b/apis/websocket/src/events/parseImage.ts index b48bf09..a94f689 100755 --- a/apis/websocket/src/events/parseImage.ts +++ b/apis/websocket/src/events/parseImage.ts @@ -1,19 +1,21 @@ -import { ClientOperation, ServerOperation } from '@revanced/bot-shared' +import { type ClientOperation, ServerOperation } from '@revanced/bot-shared' import { AsyncQueue } from '@sapphire/async-queue' -import type { EventHandler } from './index' +import type { EventHandler } from '.' const queue = new AsyncQueue() const parseImageEventHandler: EventHandler = async ( packet, - { tesseractWorker, logger, config }, + { tesseract, logger, config }, ) => { const { client, - d: { image_url: imageUrl, id }, + d: { image_url: imageUrl }, } = packet + const nextSeq = client.currentSequence++ + logger.debug(`Client ${client.id} requested to parse image from URL:`, imageUrl) logger.debug(`Queue currently has ${queue.remaining}/${config.ocrConcurrentQueues} items in it`) @@ -23,24 +25,27 @@ const parseImageEventHandler: EventHandler = async ( try { logger.debug(`Recognizing image from URL for client ${client.id}`) - const { data, jobId } = await tesseractWorker.recognize(imageUrl) + const { data, jobId } = await tesseract.recognize(imageUrl) logger.debug(`Recognized image from URL for client ${client.id} (job ${jobId}):`, data.text) - await client.send({ - op: ServerOperation.ParsedImage, - d: { - id, - text: data.text, + await client.send( + { + op: ServerOperation.ParsedImage, + d: { + text: data.text, + }, }, - }) + nextSeq, + ) } catch { logger.error(`Failed to parse image from URL for client ${client.id}:`, imageUrl) - await client.send({ - op: ServerOperation.ParseImageFailed, - d: { - id, + await client.send( + { + op: ServerOperation.ParseImageFailed, + d: null, }, - }) + nextSeq, + ) } finally { queue.shift() logger.debug( diff --git a/apis/websocket/src/events/parseText.ts b/apis/websocket/src/events/parseText.ts index 690dbb5..46ace12 100755 --- a/apis/websocket/src/events/parseText.ts +++ b/apis/websocket/src/events/parseText.ts @@ -1,35 +1,41 @@ -import { ClientOperation, ServerOperation } from '@revanced/bot-shared' +import { type ClientOperation, ServerOperation } from '@revanced/bot-shared' import { inspect as inspectObject } from 'util' -import type { EventHandler } from './index' +import type { EventHandler } from '.' -const parseTextEventHandler: EventHandler = async (packet, { witClient, logger }) => { +const parseTextEventHandler: EventHandler = async (packet, { wit, logger }) => { const { client, - d: { text, id }, + d: { text }, } = packet - logger.debug(`Client ${client.id} requested to parse text:`, text) + const nextSeq = client.currentSequence++ + const actualText = text.slice(0, 279) + + logger.debug(`Client ${client.id} requested to parse text:`, actualText) try { - const { intents } = await witClient.message(text, {}) + const { intents } = await wit.message(actualText) const intentsWithoutIds = intents.map(({ id, ...rest }) => rest) - await client.send({ - op: ServerOperation.ParsedText, - d: { - id, - labels: intentsWithoutIds, + await client.send( + { + op: ServerOperation.ParsedText, + d: { + labels: intentsWithoutIds, + }, }, - }) + nextSeq, + ) } catch (e) { - await client.send({ - op: ServerOperation.ParseTextFailed, - d: { - id, + await client.send( + { + op: ServerOperation.ParseTextFailed, + d: null, }, - }) + nextSeq, + ) if (e instanceof Error) logger.error(e.stack ?? e.message) else logger.error(inspectObject(e)) diff --git a/apis/websocket/src/events/trainMessage.ts b/apis/websocket/src/events/trainMessage.ts new file mode 100644 index 0000000..b028660 --- /dev/null +++ b/apis/websocket/src/events/trainMessage.ts @@ -0,0 +1,43 @@ +import { type ClientOperation, ServerOperation } from '@revanced/bot-shared' + +import { inspect as inspectObject } from 'util' + +import type { EventHandler } from '.' + +const trainMessageEventHandler: EventHandler = async (packet, { wit, logger }) => { + const { + client, + d: { text, label }, + } = packet + + const nextSeq = client.currentSequence++ + const actualText = text.slice(0, 279) + + logger.debug(`Client ${client.id} requested to train label ${label} with:`, actualText) + + try { + await wit.train(actualText, label) + await client.send( + { + op: ServerOperation.TrainedMessage, + d: null, + }, + nextSeq, + ) + + logger.debug(`Trained label ${label} with:`, actualText) + } catch (e) { + await client.send( + { + op: ServerOperation.TrainMessageFailed, + d: null, + }, + nextSeq, + ) + + if (e instanceof Error) logger.error(e.stack ?? e.message) + else logger.error(inspectObject(e)) + } +} + +export default trainMessageEventHandler diff --git a/apis/websocket/src/index.ts b/apis/websocket/src/index.ts index 638b9bc..110c82e 100755 --- a/apis/websocket/src/index.ts +++ b/apis/websocket/src/index.ts @@ -1,43 +1,93 @@ -import witPkg from 'node-wit' import { createWorker as createTesseractWorker } from 'tesseract.js' -const { Wit } = witPkg import { inspect as inspectObject } from 'util' import Client from './classes/Client' -import { EventContext, parseImageEventHandler, parseTextEventHandler } from './events/index' +import { + type EventContext, + type WitMessageResponse, + parseImageEventHandler, + parseTextEventHandler, + trainMessageEventHandler, +} from './events' import { DisconnectReason, HumanizedDisconnectReason, createLogger } from '@revanced/bot-shared' -import { checkEnvironment, getConfig } from './utils/index' +import { getConfig } from './utils/config' import { createServer } from 'http' -import { WebSocket, WebSocketServer } from 'ws' +import { type WebSocket, WebSocketServer } from 'ws' // Load config, init logger, check environment const config = getConfig() const logger = createLogger({ - level: config['consoleLogLevel'] === 'none' ? Infinity : config['consoleLogLevel'], + level: config.logLevel === 'none' ? Number.MAX_SAFE_INTEGER : config.logLevel, }) -checkEnvironment(logger) +if (!process.env['NODE_ENV']) logger.warn('NODE_ENV not set, defaulting to `development`') +const environment = (process.env['NODE_ENV'] ?? 'development') as NodeEnvironment + +if (!['development', 'production'].includes(environment)) { + logger.error('NODE_ENV is neither `development` nor `production`, unable to determine environment') + logger.info('Set NODE_ENV to blank to use `development` mode') + process.exit(1) +} + +logger.info(`Running in ${environment} mode...`) + +if (environment === 'production' && process.env['IS_USING_DOT_ENV']) { + logger.warn('You seem to be using .env files, this is generally not a good idea in production...') +} + +if (!process.env['WIT_AI_TOKEN']) { + logger.error('WIT_AI_TOKEN is not defined in the environment variables') + process.exit(1) +} // Workers and API clients -const tesseractWorker = await createTesseractWorker('eng') -const witClient = new Wit({ - accessToken: process.env['WIT_AI_TOKEN']!, -}) +const tesseract = await createTesseractWorker('eng') +const wit = { + token: process.env['WIT_AI_TOKEN']!, + async fetch(route: string, options?: RequestInit) { + const res = await fetch(`https://api.wit.ai${route}`, { + headers: { + Authorization: `Bearer ${this.token}`, + 'Content-Type': 'application/json', + }, + ...options, + }) + + if (!res.ok) throw new Error(`Failed to fetch from Wit.ai: ${res.statusText} (${res.status})`) + + return await res.json() + }, + message(text: string) { + return this.fetch(`/message?q=${encodeURIComponent(text)}&n=8`) as Promise + }, + async train(text: string, label: string) { + await this.fetch('/utterances', { + body: JSON.stringify([ + { + text, + intent: label, + entities: [], + traits: [], + }, + ]), + method: 'POST', + }) + }, +} as const // Server logic -const clients = new Set() -const clientSocketMap = new WeakMap() +const clientMap = new WeakMap() const eventContext: EventContext = { - tesseractWorker, + tesseract, logger, - witClient, + wit, config, } @@ -61,25 +111,24 @@ wss.on('connection', async (socket, request) => { const client = new Client({ socket, id: `${request.socket.remoteAddress}:${request.socket.remotePort}`, - heartbeatInterval: config['clientHeartbeatInterval'], }) - clientSocketMap.set(socket, client) - clients.add(client) + clientMap.set(socket, client) logger.debug(`Client ${client.id}'s instance has been added`) - logger.info(`New client connected (now ${clients.size} clients) with ID:`, client.id) + logger.info(`New client connected with ID: ${client.id}`) client.on('disconnect', reason => { - clients.delete(client) - logger.info(`Client ${client.id} disconnected because client ${HumanizedDisconnectReason[reason]}`) + logger.info( + `Client ${client.id} disconnected because client ${HumanizedDisconnectReason[reason]} (${reason})`, + ) }) - client.on('parseText', async packet => parseTextEventHandler(packet, eventContext)) + client.on('parseText', packet => parseTextEventHandler(packet, eventContext)) + client.on('parseImage', packet => parseImageEventHandler(packet, eventContext)) + client.on('trainMessage', packet => trainMessageEventHandler(packet, eventContext)) - client.on('parseImage', async packet => parseImageEventHandler(packet, eventContext)) - - if (['debug', 'trace'].includes(config['consoleLogLevel'])) { + if (['debug', 'trace'].includes(config.logLevel)) { logger.debug('Debug logs enabled, attaching debug events...') client.on('packet', ({ client, ...rawPacket }) => @@ -87,14 +136,12 @@ wss.on('connection', async (socket, request) => { ) client.on('message', d => logger.debug(`Message from client ${client.id}:`, d)) - - client.on('heartbeat', () => logger.debug('Heartbeat received from client', client.id)) } } catch (e) { if (e instanceof Error) logger.error(e.stack ?? e.message) else logger.error(inspectObject(e)) - const client = clientSocketMap.get(socket) + const client = clientMap.get(socket) if (!client) { logger.error( @@ -104,9 +151,6 @@ wss.on('connection', async (socket, request) => { } if (client.disconnected === false) client.disconnect(DisconnectReason.ServerError) - else client.forceDisconnect() - - clients.delete(client) logger.debug(`Client ${client.id} disconnected because of an internal error`) } @@ -114,7 +158,7 @@ wss.on('connection', async (socket, request) => { // Start the server -server.listen(config['port'], config['address']) +server.listen(config.port, config.address) logger.debug(`Starting with these configurations: ${inspectObject(config)}`) diff --git a/apis/websocket/src/utils/checkEnvironment.ts b/apis/websocket/src/utils/checkEnvironment.ts deleted file mode 100755 index 803d94e..0000000 --- a/apis/websocket/src/utils/checkEnvironment.ts +++ /dev/null @@ -1,23 +0,0 @@ -import type { Logger } from '@revanced/bot-shared' - -export default function checkEnvironment(logger: Logger) { - if (!process.env['NODE_ENV']) logger.warn('NODE_ENV not set, defaulting to `development`') - const environment = (process.env['NODE_ENV'] ?? 'development') as NodeEnvironment - - if (!['development', 'production'].includes(environment)) { - logger.error('NODE_ENV is neither `development` nor `production`, unable to determine environment') - logger.info('Set NODE_ENV to blank to use `development` mode') - process.exit(1) - } - - logger.info(`Running in ${environment} mode...`) - - if (environment === 'production' && process.env['IS_USING_DOT_ENV']) { - logger.warn('You seem to be using .env files, this is generally not a good idea in production...') - } - - if (!process.env['WIT_AI_TOKEN']) { - logger.error('WIT_AI_TOKEN is not defined in the environment variables') - process.exit(1) - } -} diff --git a/apis/websocket/src/utils/getConfig.ts b/apis/websocket/src/utils/config.ts similarity index 87% rename from apis/websocket/src/utils/getConfig.ts rename to apis/websocket/src/utils/config.ts index 90bf638..7beb09b 100755 --- a/apis/websocket/src/utils/getConfig.ts +++ b/apis/websocket/src/utils/config.ts @@ -7,7 +7,7 @@ const configPath = resolvePath(process.cwd(), 'config.json') const userConfig: Partial = existsSync(configPath) ? ( await import(pathToFileURL(configPath).href, { - assert: { + with: { type: 'json', }, }) @@ -28,10 +28,9 @@ export const defaultConfig: Config = { address: '127.0.0.1', port: 8080, ocrConcurrentQueues: 1, - clientHeartbeatInterval: 60000, - consoleLogLevel: 'info', + logLevel: 'info', } -export default function getConfig() { +export function getConfig() { return Object.assign(defaultConfig, userConfig) satisfies Config } diff --git a/apis/websocket/src/utils/index.ts b/apis/websocket/src/utils/index.ts deleted file mode 100755 index 24a2c7c..0000000 --- a/apis/websocket/src/utils/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export { default as getConfig } from './getConfig' -export { default as checkEnvironment } from './checkEnvironment' diff --git a/apis/websocket/tsconfig.json b/apis/websocket/tsconfig.json index 29a7c65..b80117e 100755 --- a/apis/websocket/tsconfig.json +++ b/apis/websocket/tsconfig.json @@ -1,11 +1,14 @@ { - "extends": "../../tsconfig.json", - "compilerOptions": { - "baseUrl": ".", - "outDir": "dist", - "module": "ESNext", - "composite": false - }, - "exclude": ["node_modules", "dist"], - "include": ["./*.json", "src/**/*.ts"] + "extends": "../../tsconfig.json", + "compilerOptions": { + "baseUrl": ".", + "outDir": "dist", + "module": "ESNext", + "target": "ESNext", + "lib": ["ESNext"], + "composite": false, + "skipLibCheck": true + }, + "exclude": ["node_modules", "dist"], + "include": ["./*.json", "src/**/*.ts"] } diff --git a/packages/api/package.json b/packages/api/package.json index e30e8b4..ca31c12 100755 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -8,7 +8,7 @@ "scripts": { "build": "bun bundle && bun types", "watch": "bunx conc --raw \"bun bundle:watch\" \"bun types:watch\"", - "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify", + "bundle": "bun build src/index.ts --outdir=dist --target=bun", "bundle:watch": "bun run bundle --watch", "types": "tsc --declaration --emitDeclarationOnly", "types:watch": "bun types --watch --preserveWatchOutput", diff --git a/packages/api/src/classes/Client.ts b/packages/api/src/classes/Client.ts index 3db2488..c645aa5 100755 --- a/packages/api/src/classes/Client.ts +++ b/packages/api/src/classes/Client.ts @@ -1,5 +1,10 @@ -import { ClientOperation, Packet, ServerOperation } from '@revanced/bot-shared' -import { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOptions } from './ClientWebSocket' +import { ClientOperation, ServerOperation } from '@revanced/bot-shared' +import { awaitPacket } from 'src/utils/packets' +import { + type ClientWebSocketEvents, + ClientWebSocketManager, + type ClientWebSocketManagerOptions, +} from './ClientWebSocket' /** * The client that connects to the API. @@ -7,7 +12,6 @@ import { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOp export default class Client { ready = false ws: ClientWebSocketManager - #parseId = 0 constructor(options: ClientOptions) { this.ws = new ClientWebSocketManager(options.api.websocket) @@ -15,7 +19,7 @@ export default class Client { this.ready = true }) this.ws.on('disconnect', () => { - + this.ready = false }) } @@ -35,36 +39,34 @@ export default class Client { async parseText(text: string) { this.#throwIfNotReady() - const currentId = (this.#parseId++).toString() + return await this.ws + .send({ + op: ClientOperation.ParseText, + d: { + text, + }, + }) + .then(() => { + // Since we don't have heartbeats anymore, this is fine. + // But if we add anything similar, this will cause another race condition + // To fix this, we can try adding a instanced function that would return the currentSequence + // and it would be updated every time a "heartbeat ack" packet is received + const expectedNextSeq = this.ws.currentSequence + 1 + const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) => + awaitPacket(this.ws, op, expectedNextSeq, timeout) - this.ws.send({ - op: ClientOperation.ParseText, - d: { - text, - id: currentId, - }, - }) - - type CorrectPacket = Packet - - const promise = new Promise((rs, rj) => { - const parsedTextListener = (packet: CorrectPacket) => { - if (packet.d.id !== currentId) return - this.ws.off('parsedText', parsedTextListener) - rs(packet.d) - } - - const parseTextFailedListener = (packet: Packet) => { - if (packet.d.id !== currentId) return - this.ws.off('parseTextFailed', parseTextFailedListener) - rj() - } - - this.ws.on('parsedText', parsedTextListener) - this.ws.on('parseTextFailed', parseTextFailedListener) - }) - - return await promise + return Promise.race([ + awaitPkt(ServerOperation.ParsedText), + awaitPkt(ServerOperation.ParseTextFailed, this.ws.timeout + 5000), + ]) + .then(pkt => { + if (pkt.op === ServerOperation.ParsedText) return pkt.d + throw new Error('Failed to parse text, the API encountered an error') + }) + .catch(() => { + throw new Error('Failed to parse text, the API did not respond in time') + }) + }) } /** @@ -75,36 +77,62 @@ export default class Client { async parseImage(url: string) { this.#throwIfNotReady() - const currentId = (this.#parseId++).toString() + return await this.ws + .send({ + op: ClientOperation.ParseImage, + d: { + image_url: url, + }, + }) + .then(() => { + // See line 48 + const expectedNextSeq = this.ws.currentSequence + 1 + const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) => + awaitPacket(this.ws, op, expectedNextSeq, timeout) - this.ws.send({ - op: ClientOperation.ParseImage, - d: { - image_url: url, - id: currentId, - }, - }) + return Promise.race([ + awaitPkt(ServerOperation.ParsedImage), + awaitPkt(ServerOperation.ParseImageFailed, this.ws.timeout + 5000), + ]) + .then(pkt => { + if (pkt.op === ServerOperation.ParsedImage) return pkt.d + throw new Error('Failed to parse image, the API encountered an error') + }) + .catch(() => { + throw new Error('Failed to parse image, the API did not respond in time') + }) + }) + } - type CorrectPacket = Packet + async trainMessage(text: string, label: string) { + this.#throwIfNotReady() - const promise = new Promise((rs, rj) => { - const parsedImageListener = (packet: CorrectPacket) => { - if (packet.d.id !== currentId) return - this.ws.off('parsedImage', parsedImageListener) - rs(packet.d) - } + return await this.ws + .send({ + op: ClientOperation.TrainMessage, + d: { + label, + text, + }, + }) + .then(() => { + // See line 48 + const expectedNextSeq = this.ws.currentSequence + 1 + const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) => + awaitPacket(this.ws, op, expectedNextSeq, timeout) - const parseImageFailedListener = (packet: Packet) => { - if (packet.d.id !== currentId) return - this.ws.off('parseImageFailed', parseImageFailedListener) - rj() - } - - this.ws.on('parsedImage', parsedImageListener) - this.ws.on('parseImageFailed', parseImageFailedListener) - }) - - return await promise + return Promise.race([ + awaitPkt(ServerOperation.TrainedMessage), + awaitPkt(ServerOperation.TrainMessageFailed, this.ws.timeout + 5000), + ]) + .then(pkt => { + if (pkt.op === ServerOperation.TrainedMessage) return + throw new Error('Failed to train message, the API encountered an error') + }) + .catch(() => { + throw new Error('Failed to train message, the API did not respond in time') + }) + }) } /** @@ -135,14 +163,18 @@ export default class Client { * @param handler The event handler * @returns The event handler function */ - once( - name: TOpName, - handler: ClientWebSocketEvents[TOpName], - ) { + once(name: TOpName, handler: ClientWebSocketEvents[TOpName]) { this.ws.once(name, handler) return handler } + /** + * Disconnects the client from the API + */ + disconnect() { + this.ws.disconnect() + } + #throwIfNotReady() { if (!this.isReady()) throw new Error('Client is not ready') } diff --git a/packages/api/src/classes/ClientWebSocket.ts b/packages/api/src/classes/ClientWebSocket.ts index 2d71c9c..34847f6 100755 --- a/packages/api/src/classes/ClientWebSocket.ts +++ b/packages/api/src/classes/ClientWebSocket.ts @@ -1,8 +1,8 @@ import { EventEmitter } from 'events' import { - ClientOperation, + type ClientOperation, DisconnectReason, - Packet, + type Packet, ServerOperation, deserializePacket, isServerPacket, @@ -10,7 +10,7 @@ import { uncapitalize, } from '@revanced/bot-shared' import type TypedEmitter from 'typed-emitter' -import { RawData, WebSocket } from 'ws' +import { type RawData, WebSocket } from 'ws' /** * The class that handles the WebSocket connection to the server. @@ -21,10 +21,9 @@ export class ClientWebSocketManager { timeout: number ready = false - disconnected: boolean | DisconnectReason = DisconnectReason.NeverConnected - config: Readonly['d']> | null = null + disconnected: false | DisconnectReason = false + currentSequence = 0 - #hbTimeout: NodeJS.Timeout = null! #socket: WebSocket = null! #emitter = new EventEmitter() as TypedEmitter @@ -42,26 +41,27 @@ export class ClientWebSocketManager { try { this.#socket = new WebSocket(this.url) - setTimeout(() => { - if (!this.ready) throw new Error('WebSocket connection timed out') - this.#socket.close() + const timeout = setTimeout(() => { + if (!this.ready) { + this.#socket?.close(DisconnectReason.TooSlow) + throw new Error('WebSocket connection was not readied in time') + } }, this.timeout) this.#socket.on('open', () => { - this.disconnected = false + clearTimeout(timeout) this.#listen() - this.ready = true - this.#emitter.emit('ready') rs() }) - this.#socket.on('error', (err) => { + this.#socket.on('error', err => { + clearTimeout(timeout) throw err }) this.#socket.on('close', (code, reason) => { - if (code === 1006) throw new Error(`Failed to connect to WebSocket server: ${reason}`) - this.#handleDisconnect(DisconnectReason.Generic) + clearTimeout(timeout) + this._handleDisconnect(code, reason.toString()) }) } catch (e) { rj(e) @@ -75,10 +75,7 @@ export class ClientWebSocketManager { * @param handler The event handler * @returns The event handler function */ - on( - name: TOpName, - handler: ClientWebSocketEvents[typeof name], - ) { + on(name: TOpName, handler: ClientWebSocketEvents[typeof name]) { this.#emitter.on(name, handler) } @@ -88,10 +85,7 @@ export class ClientWebSocketManager { * @param handler The event handler to remove * @returns The removed event handler function */ - off( - name: TOpName, - handler: ClientWebSocketEvents[typeof name], - ) { + off(name: TOpName, handler: ClientWebSocketEvents[typeof name]) { this.#emitter.off(name, handler) } @@ -101,10 +95,7 @@ export class ClientWebSocketManager { * @param handler The event handler * @returns The event handler function */ - once( - name: TOpName, - handler: ClientWebSocketEvents[typeof name], - ) { + once(name: TOpName, handler: ClientWebSocketEvents[typeof name]) { this.#emitter.once(name, handler) } @@ -126,7 +117,7 @@ export class ClientWebSocketManager { */ disconnect() { this.#throwIfDisconnected('Cannot disconnect when already disconnected from the server') - this.#handleDisconnect(DisconnectReason.PlannedDisconnect) + this._handleDisconnect(DisconnectReason.PlannedDisconnect) } /** @@ -143,22 +134,22 @@ export class ClientWebSocketManager { if (!isServerPacket(packet)) return this.#emitter.emit('invalidPacket', packet) + this.currentSequence = packet.s this.#emitter.emit('packet', packet) switch (packet.op) { case ServerOperation.Hello: { - const data = Object.freeze((packet as Packet).d) - this.config = data - this.#emitter.emit('hello', data) - this.#startHeartbeating() + this.#emitter.emit('hello') + this.ready = true + this.#emitter.emit('ready') break } case ServerOperation.Disconnect: - return this.#handleDisconnect((packet as Packet).d.reason) + return this._handleDisconnect((packet as Packet).d.reason) default: return this.#emitter.emit( uncapitalize(ServerOperation[packet.op] as ClientWebSocketEventName), - // @ts-expect-error TypeScript doesn't know that the lines above negate the type enough + // @ts-expect-error: TS at it again packet, ) } @@ -170,30 +161,12 @@ export class ClientWebSocketManager { if (this.#socket.readyState !== this.#socket.OPEN) throw new Error(errorMessage) } - #handleDisconnect(reason: DisconnectReason) { - clearTimeout(this.#hbTimeout) - this.disconnected = reason - this.#socket.close() + protected _handleDisconnect(reason: DisconnectReason | number, message?: string) { + this.disconnected = reason in DisconnectReason ? reason : DisconnectReason.Generic + this.#socket?.close(reason) this.#socket = null! - this.#emitter.emit('disconnect', reason) - } - - #startHeartbeating() { - this.on('heartbeatAck', packet => { - this.#hbTimeout = setTimeout(() => { - this.send({ - op: ClientOperation.Heartbeat, - d: null, - }) - }, packet.d.nextHeartbeat - Date.now()) - }) - - // Immediately send a heartbeat so we can get when to send the next one - this.send({ - op: ClientOperation.Heartbeat, - d: null, - }) + this.#emitter.emit('disconnect', reason, message) } protected _toBuffer(data: RawData) { @@ -217,16 +190,18 @@ export interface ClientWebSocketManagerOptions { export type ClientWebSocketEventName = keyof typeof ServerOperation -export type ClientWebSocketEvents = { - [K in Uncapitalize]: ( - packet: Packet<(typeof ServerOperation)[Capitalize]>, - ) => Promise | void -} & { - hello: (config: NonNullable) => Promise | void +type ClientWebSocketPredefinedEvents = { + hello: () => Promise | void ready: () => Promise | void packet: (packet: Packet) => Promise | void invalidPacket: (packet: Packet) => Promise | void - disconnect: (reason: DisconnectReason) => Promise | void + disconnect: (reason: DisconnectReason | number, message?: string) => Promise | void } +export type ClientWebSocketEvents = { + [K in Exclude, keyof ClientWebSocketPredefinedEvents>]: ( + packet: Packet<(typeof ServerOperation)[Capitalize]>, + ) => Promise | void +} & ClientWebSocketPredefinedEvents + export type ReadiedClientWebSocketManager = RequiredProperty> diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index 63e4228..f50cdd7 100755 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -1 +1 @@ -export * from './classes/index' +export * from './classes' diff --git a/packages/api/src/utils/packets.ts b/packages/api/src/utils/packets.ts new file mode 100644 index 0000000..70e14de --- /dev/null +++ b/packages/api/src/utils/packets.ts @@ -0,0 +1,26 @@ +import type { Packet, ServerOperation } from '@revanced/bot-shared' +import type { ClientWebSocketManager } from 'src/classes' + +export function awaitPacket( + ws: ClientWebSocketManager, + op: TOp, + expectedSeq: number, + timeout = 10000, +): Promise> { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + ws.off('packet', handler) + reject('Awaiting packet timed out') + }, timeout) + + function handler(packet: Packet) { + if (packet.op === op && packet.s === expectedSeq) { + clearTimeout(timer) + ws.off('packet', handler) + resolve(packet as Packet) + } + } + + ws.on('packet', handler) + }) +} diff --git a/packages/api/tsconfig.json b/packages/api/tsconfig.json index 01f29d1..ef4e094 100755 --- a/packages/api/tsconfig.json +++ b/packages/api/tsconfig.json @@ -1,12 +1,12 @@ { - "extends": "../../tsconfig.packages.json", - "compilerOptions": { - "baseUrl": ".", - "rootDir": "./src", - "outDir": "dist", - "module": "ESNext", - "composite": true, - "noEmit": false - }, - "exclude": ["node_modules", "dist"] + "extends": "../../tsconfig.packages.json", + "compilerOptions": { + "baseUrl": ".", + "rootDir": "./src", + "outDir": "dist", + "module": "ESNext", + "composite": true, + "noEmit": false + }, + "exclude": ["node_modules", "dist"] } diff --git a/packages/shared/package.json b/packages/shared/package.json index 9d6141d..f1c9023 100755 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -1,40 +1,40 @@ { - "name": "@revanced/bot-shared", - "type": "module", - "version": "0.1.0", - "description": "🙌🏻 Shared components for bots assisting ReVanced", - "main": "dist/index.js", - "types": "dist/index.d.ts", - "scripts": { - "build": "bun bundle && bun types", - "watch": "conc --raw \"bun bundle:watch\" \"bun types:watch\"", - "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify", - "bundle:watch": "bun run bundle --watch", - "types": "tsc --declaration --emitDeclarationOnly", - "types:watch": "bun types --watch --preserveWatchOutput", - "types:clean": "bun types --build --clean" - }, - "repository": { - "type": "git", - "url": "git+https://github.com/revanced/revanced-helper.git", - "directory": "packages/shared" - }, - "author": "Palm (https://github.com/PalmDevs)", - "contributors": [ - "Palm (https://github.com/PalmDevs)", - "ReVanced (https://github.com/revanced)" - ], - "license": "GPL-3.0-or-later", - "bugs": { - "url": "https://github.com/revanced/revanced-helper/issues" - }, - "homepage": "https://github.com/revanced/revanced-helper#readme", - "dependencies": { - "bson": "^6.2.0", - "chalk": "^5.3.0", - "supports-color": "^9.4.0", - "tracer": "^1.3.0", - "valibot": "^0.21.0", - "zod": "^3.22.4" - } + "name": "@revanced/bot-shared", + "type": "module", + "version": "0.1.0", + "description": "🙌🏻 Shared components for bots assisting ReVanced", + "main": "dist/index.js", + "types": "dist/index.d.ts", + "scripts": { + "build": "bun bundle && bun types", + "watch": "conc --raw \"bun bundle:watch\" \"bun types:watch\"", + "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify", + "bundle:watch": "bun run bundle --watch", + "types": "tsc --declaration --emitDeclarationOnly", + "types:watch": "bun types --watch --preserveWatchOutput", + "types:clean": "bun types --build --clean" + }, + "repository": { + "type": "git", + "url": "git+https://github.com/revanced/revanced-helper.git", + "directory": "packages/shared" + }, + "author": "Palm (https://github.com/PalmDevs)", + "contributors": [ + "Palm (https://github.com/PalmDevs)", + "ReVanced (https://github.com/revanced)" + ], + "license": "GPL-3.0-or-later", + "bugs": { + "url": "https://github.com/revanced/revanced-helper/issues" + }, + "homepage": "https://github.com/revanced/revanced-helper#readme", + "dependencies": { + "bson": "^6.5.0", + "chalk": "^5.3.0", + "supports-color": "^9.4.0", + "tracer": "^1.3.0", + "valibot": "^0.30.0", + "zod": "^3.22.4" + } } diff --git a/packages/shared/src/constants/DisconnectReason.ts b/packages/shared/src/constants/DisconnectReason.ts index 37588ba..048f49b 100755 --- a/packages/shared/src/constants/DisconnectReason.ts +++ b/packages/shared/src/constants/DisconnectReason.ts @@ -3,29 +3,33 @@ */ enum DisconnectReason { /** - * Unknown reason + * The client disconnected on its own (**CLIENT-ONLY**) */ - Generic = 1, - /** - * The client did not respond in time - */ - TimedOut = 2, + PlannedDisconnect = 1000, /** * The client sent an invalid packet (unserializable or invalid JSON) */ - InvalidPacket = 3, + InvalidPacket = 1007, /** * The server has encountered an internal error */ - ServerError = 4, + ServerError = 1011, /** - * The client had never connected to the server (**CLIENT-ONLY**) + * Unknown reason */ - NeverConnected = 5, + Generic = 4000, /** - * The client disconnected on its own (**CLIENT-ONLY**) + * The client did not respond with a heartbeat in time */ - PlannedDisconnect = 6, + TimedOut = 4001, + /** + * The receiving end didn't have an open socket + */ + NoOpenSocket = 4003, + /** + * The client was not ready in time (**CLIENT-ONLY**) + */ + TooSlow = 4002, } export default DisconnectReason diff --git a/packages/shared/src/constants/HumanizedDisconnectReason.ts b/packages/shared/src/constants/HumanizedDisconnectReason.ts index 0f3c536..e3d18b6 100755 --- a/packages/shared/src/constants/HumanizedDisconnectReason.ts +++ b/packages/shared/src/constants/HumanizedDisconnectReason.ts @@ -4,12 +4,14 @@ import DisconnectReason from './DisconnectReason' * Humanized disconnect reasons for logs */ const HumanizedDisconnectReason = { - [DisconnectReason.InvalidPacket]: 'has sent invalid packet', - [DisconnectReason.Generic]: 'has been disconnected for unknown reasons', - [DisconnectReason.TimedOut]: 'has timed out', - [DisconnectReason.ServerError]: 'has been disconnected due to an internal server error', - [DisconnectReason.NeverConnected]: 'had never connected to the server', - [DisconnectReason.PlannedDisconnect]: 'has disconnected on its own', -} as const satisfies Record + [1006]: 'the receiving end had unexpectedly closed the connection', + [DisconnectReason.InvalidPacket]: 'the client has sent invalid packet', + [DisconnectReason.Generic]: '(unknown reason)', + [DisconnectReason.TimedOut]: 'the client did not respond with a heartbeat in time', + [DisconnectReason.ServerError]: 'the server had an internal server error', + [DisconnectReason.TooSlow]: 'the client was not ready in time', + [DisconnectReason.PlannedDisconnect]: 'the client has disconnected on its own', + [DisconnectReason.NoOpenSocket]: 'the receiving end did not have an open socket', +} as const satisfies Record export default HumanizedDisconnectReason diff --git a/packages/shared/src/constants/Operation.ts b/packages/shared/src/constants/Operation.ts index 9c4ee4e..d3104ad 100755 --- a/packages/shared/src/constants/Operation.ts +++ b/packages/shared/src/constants/Operation.ts @@ -2,33 +2,28 @@ * Client operation codes for the gateway */ export enum ClientOperation { - /** - * Client's heartbeat (to check if the connection is dead or not) - */ - Heartbeat = 100, - /** * Client's request to parse text */ - ParseText = 110, + ParseText = 100, /** * Client's request to parse image */ - ParseImage = 111, + ParseImage = 101, + /** + * Client's request to train a message + */ + TrainMessage = 102, } /** * Server operation codes for the gateway */ export enum ServerOperation { - /** - * Server's acknowledgement of a client's heartbeat - */ - HeartbeatAck = 1, /** * Server's initial response to a client's connection */ - Hello = 2, + Hello = 1, /** * Server's response to client's request to parse text @@ -46,6 +41,14 @@ export enum ServerOperation { * Server's failure response to client's request to parse image */ ParseImageFailed = 13, + /** + * Server's response to client's request to train a message + */ + TrainedMessage = 14, + /** + * Server's failure response to client's request to train a message + */ + TrainMessageFailed = 15, /** * Server's disconnect message diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index ab1d8b0..8d84460 100755 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,3 +1,3 @@ -export * from './constants/index' -export * from './schemas/index' -export * from './utils/index' +export * from './constants' +export * from './schemas' +export * from './utils' diff --git a/packages/shared/src/schemas/Packet.ts b/packages/shared/src/schemas/Packet.ts index e1f93f9..2e91901 100755 --- a/packages/shared/src/schemas/Packet.ts +++ b/packages/shared/src/schemas/Packet.ts @@ -1,13 +1,12 @@ import { url, - AnySchema, - NullSchema, - ObjectSchema, - Output, + type AnySchema, + type NullSchema, + type ObjectSchema, + type Output, array, enum_, null_, - number, object, parse, special, @@ -30,6 +29,8 @@ export const PacketSchema = special(input => { 'd' in input && typeof input.d === 'object' ) { + if (input.op in ServerOperation && !('s' in input && typeof input.s === 'number')) return false + try { parse(PacketDataSchemas[input.op as Operation], input.d) return true @@ -44,14 +45,8 @@ export const PacketSchema = special(input => { * Schema to validate packet data for each possible operations */ export const PacketDataSchemas = { - [ServerOperation.Hello]: object({ - heartbeatInterval: number(), - }), - [ServerOperation.HeartbeatAck]: object({ - nextHeartbeat: number(), - }), + [ServerOperation.Hello]: null_(), [ServerOperation.ParsedText]: object({ - id: string(), labels: array( object({ name: string(), @@ -60,35 +55,38 @@ export const PacketDataSchemas = { ), }), [ServerOperation.ParsedImage]: object({ - id: string(), text: string(), }), - [ServerOperation.ParseTextFailed]: object({ - id: string(), - }), - [ServerOperation.ParseImageFailed]: object({ - id: string(), - }), + [ServerOperation.ParseTextFailed]: null_(), + [ServerOperation.ParseImageFailed]: null_(), [ServerOperation.Disconnect]: object({ reason: enum_(DisconnectReason), }), + [ServerOperation.TrainedMessage]: null_(), + [ServerOperation.TrainMessageFailed]: null_(), - [ClientOperation.Heartbeat]: null_(), [ClientOperation.ParseText]: object({ - id: string(), text: string(), }), [ClientOperation.ParseImage]: object({ - id: string(), image_url: string([url()]), }), + [ClientOperation.TrainMessage]: object({ + text: string(), + label: string(), + }), } as const satisfies Record< Operation, // biome-ignore lint/suspicious/noExplicitAny: This is a schema, it's not possible to type it ObjectSchema | AnySchema | NullSchema > -export type Packet = { +export type Packet = TOp extends ServerOperation + ? PacketWithSequenceNumber + : Omit, 's'> + +type PacketWithSequenceNumber = { op: TOp d: Output<(typeof PacketDataSchemas)[TOp]> + s: number } diff --git a/packages/shared/src/utils/guard.ts b/packages/shared/src/utils/guard.ts index 9b92cba..91825cb 100755 --- a/packages/shared/src/utils/guard.ts +++ b/packages/shared/src/utils/guard.ts @@ -1,5 +1,5 @@ -import { ClientOperation, Operation, ServerOperation } from '../constants/Operation' -import { Packet } from '../schemas/Packet' +import { ClientOperation, type Operation, ServerOperation } from '../constants/Operation' +import type { Packet } from '../schemas/Packet' /** * Checks whether a packet is trying to do the given operation @@ -21,7 +21,7 @@ export function isClientPacket(packet: Packet): packet is Packet