fix(packages/api)!: handle dead connections better

This commit is contained in:
PalmDevs
2024-01-18 22:43:19 +07:00
parent 4792fde5a3
commit 56e364cedb
5 changed files with 71 additions and 57 deletions

View File

@@ -1,30 +1,22 @@
import { ClientOperation, Packet, ServerOperation } from '@revanced/bot-shared' import { ClientOperation, Packet, ServerOperation } from '@revanced/bot-shared'
import ClientGateway, { ClientGatewayEventHandlers } from './ClientGateway' import { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOptions } from './ClientWebSocket'
/** /**
* The client that connects to the API. * The client that connects to the API.
*/ */
export default class Client { export default class Client {
ready = false ready = false
gateway: ClientGateway ws: ClientWebSocketManager
#parseId = 0 #parseId = 0
constructor(options: ClientOptions) { constructor(options: ClientOptions) {
this.gateway = new ClientGateway({ this.ws = new ClientWebSocketManager(options.api.websocket)
url: options.api.gatewayUrl, this.ws.on('ready', () => {
})
this.gateway.on('ready', () => {
this.ready = true this.ready = true
}) })
} this.ws.on('disconnect', () => {
/** })
* Connects to the WebSocket API
* @returns A promise that resolves when the client is ready
*/
connect() {
return this.gateway.connect()
} }
/** /**
@@ -45,7 +37,7 @@ export default class Client {
const currentId = (this.#parseId++).toString() const currentId = (this.#parseId++).toString()
this.gateway.send({ this.ws.send({
op: ClientOperation.ParseText, op: ClientOperation.ParseText,
d: { d: {
text, text,
@@ -58,18 +50,18 @@ export default class Client {
const promise = new Promise<CorrectPacket['d']>((rs, rj) => { const promise = new Promise<CorrectPacket['d']>((rs, rj) => {
const parsedTextListener = (packet: CorrectPacket) => { const parsedTextListener = (packet: CorrectPacket) => {
if (packet.d.id !== currentId) return if (packet.d.id !== currentId) return
this.gateway.off('parsedText', parsedTextListener) this.ws.off('parsedText', parsedTextListener)
rs(packet.d) rs(packet.d)
} }
const parseTextFailedListener = (packet: Packet<ServerOperation.ParseTextFailed>) => { const parseTextFailedListener = (packet: Packet<ServerOperation.ParseTextFailed>) => {
if (packet.d.id !== currentId) return if (packet.d.id !== currentId) return
this.gateway.off('parseTextFailed', parseTextFailedListener) this.ws.off('parseTextFailed', parseTextFailedListener)
rj() rj()
} }
this.gateway.on('parsedText', parsedTextListener) this.ws.on('parsedText', parsedTextListener)
this.gateway.on('parseTextFailed', parseTextFailedListener) this.ws.on('parseTextFailed', parseTextFailedListener)
}) })
return await promise return await promise
@@ -85,7 +77,7 @@ export default class Client {
const currentId = (this.#parseId++).toString() const currentId = (this.#parseId++).toString()
this.gateway.send({ this.ws.send({
op: ClientOperation.ParseImage, op: ClientOperation.ParseImage,
d: { d: {
image_url: url, image_url: url,
@@ -98,18 +90,18 @@ export default class Client {
const promise = new Promise<CorrectPacket['d']>((rs, rj) => { const promise = new Promise<CorrectPacket['d']>((rs, rj) => {
const parsedImageListener = (packet: CorrectPacket) => { const parsedImageListener = (packet: CorrectPacket) => {
if (packet.d.id !== currentId) return if (packet.d.id !== currentId) return
this.gateway.off('parsedImage', parsedImageListener) this.ws.off('parsedImage', parsedImageListener)
rs(packet.d) rs(packet.d)
} }
const parseImageFailedListener = (packet: Packet<ServerOperation.ParseImageFailed>) => { const parseImageFailedListener = (packet: Packet<ServerOperation.ParseImageFailed>) => {
if (packet.d.id !== currentId) return if (packet.d.id !== currentId) return
this.gateway.off('parseImageFailed', parseImageFailedListener) this.ws.off('parseImageFailed', parseImageFailedListener)
rj() rj()
} }
this.gateway.on('parsedImage', parsedImageListener) this.ws.on('parsedImage', parsedImageListener)
this.gateway.on('parseImageFailed', parseImageFailedListener) this.ws.on('parseImageFailed', parseImageFailedListener)
}) })
return await promise return await promise
@@ -121,8 +113,8 @@ export default class Client {
* @param handler The event handler * @param handler The event handler
* @returns The event handler function * @returns The event handler function
*/ */
on<TOpName extends keyof ClientGatewayEventHandlers>(name: TOpName, handler: ClientGatewayEventHandlers[TOpName]) { on<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[TOpName]) {
this.gateway.on(name, handler) this.ws.on(name, handler)
return handler return handler
} }
@@ -132,8 +124,8 @@ export default class Client {
* @param handler The event handler to remove * @param handler The event handler to remove
* @returns The removed event handler function * @returns The removed event handler function
*/ */
off<TOpName extends keyof ClientGatewayEventHandlers>(name: TOpName, handler: ClientGatewayEventHandlers[TOpName]) { off<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[TOpName]) {
this.gateway.off(name, handler) this.ws.off(name, handler)
return handler return handler
} }
@@ -143,11 +135,11 @@ export default class Client {
* @param handler The event handler * @param handler The event handler
* @returns The event handler function * @returns The event handler function
*/ */
once<TOpName extends keyof ClientGatewayEventHandlers>( once<TOpName extends keyof ClientWebSocketEvents>(
name: TOpName, name: TOpName,
handler: ClientGatewayEventHandlers[TOpName], handler: ClientWebSocketEvents[TOpName],
) { ) {
this.gateway.once(name, handler) this.ws.once(name, handler)
return handler return handler
} }
@@ -163,5 +155,5 @@ export interface ClientOptions {
} }
export interface ClientApiOptions { export interface ClientApiOptions {
gatewayUrl: string websocket: ClientWebSocketManagerOptions
} }

View File

@@ -16,18 +16,21 @@ import { RawData, WebSocket } from 'ws'
* The class that handles the WebSocket connection to the server. * The class that handles the WebSocket connection to the server.
* This is the only relevant class for the time being. But in the future, there may be more classes to handle different protocols of the API. * This is the only relevant class for the time being. But in the future, there may be more classes to handle different protocols of the API.
*/ */
export default class ClientGateway { export class ClientWebSocketManager {
readonly url: string readonly url: string
timeout: number
ready = false ready = false
disconnected: boolean | DisconnectReason = DisconnectReason.NeverConnected disconnected: boolean | DisconnectReason = DisconnectReason.NeverConnected
config: Readonly<Packet<ServerOperation.Hello>['d']> | null = null! config: Readonly<Packet<ServerOperation.Hello>['d']> | null = null
#hbTimeout: NodeJS.Timeout = null! #hbTimeout: NodeJS.Timeout = null!
#socket: WebSocket = null! #socket: WebSocket = null!
#emitter = new EventEmitter() as TypedEmitter<ClientGatewayEventHandlers> #emitter = new EventEmitter() as TypedEmitter<ClientWebSocketEvents>
constructor(options: ClientGatewayOptions) { constructor(options: ClientWebSocketManagerOptions) {
this.url = options.url this.url = options.url
this.timeout = options.timeout ?? 10000
} }
/** /**
@@ -39,6 +42,11 @@ export default class ClientGateway {
try { try {
this.#socket = new WebSocket(this.url) this.#socket = new WebSocket(this.url)
setTimeout(() => {
if (!this.ready) throw new Error('WebSocket connection timed out')
this.#socket.close()
}, this.timeout)
this.#socket.on('open', () => { this.#socket.on('open', () => {
this.disconnected = false this.disconnected = false
this.#listen() this.#listen()
@@ -47,10 +55,14 @@ export default class ClientGateway {
rs() rs()
}) })
const errorHandler = () => this.#handleDisconnect(DisconnectReason.Generic) this.#socket.on('error', (err) => {
throw err
})
this.#socket.on('close', errorHandler) this.#socket.on('close', (code, reason) => {
this.#socket.on('error', errorHandler) if (code === 1006) throw new Error(`Failed to connect to WebSocket server: ${reason}`)
this.#handleDisconnect(DisconnectReason.Generic)
})
} catch (e) { } catch (e) {
rj(e) rj(e)
} }
@@ -63,9 +75,9 @@ export default class ClientGateway {
* @param handler The event handler * @param handler The event handler
* @returns The event handler function * @returns The event handler function
*/ */
on<TOpName extends keyof ClientGatewayEventHandlers>( on<TOpName extends keyof ClientWebSocketEvents>(
name: TOpName, name: TOpName,
handler: ClientGatewayEventHandlers[typeof name], handler: ClientWebSocketEvents[typeof name],
) { ) {
this.#emitter.on(name, handler) this.#emitter.on(name, handler)
} }
@@ -76,9 +88,9 @@ export default class ClientGateway {
* @param handler The event handler to remove * @param handler The event handler to remove
* @returns The removed event handler function * @returns The removed event handler function
*/ */
off<TOpName extends keyof ClientGatewayEventHandlers>( off<TOpName extends keyof ClientWebSocketEvents>(
name: TOpName, name: TOpName,
handler: ClientGatewayEventHandlers[typeof name], handler: ClientWebSocketEvents[typeof name],
) { ) {
this.#emitter.off(name, handler) this.#emitter.off(name, handler)
} }
@@ -89,9 +101,9 @@ export default class ClientGateway {
* @param handler The event handler * @param handler The event handler
* @returns The event handler function * @returns The event handler function
*/ */
once<TOpName extends keyof ClientGatewayEventHandlers>( once<TOpName extends keyof ClientWebSocketEvents>(
name: TOpName, name: TOpName,
handler: ClientGatewayEventHandlers[typeof name], handler: ClientWebSocketEvents[typeof name],
) { ) {
this.#emitter.once(name, handler) this.#emitter.once(name, handler)
} }
@@ -114,14 +126,14 @@ export default class ClientGateway {
*/ */
disconnect() { disconnect() {
this.#throwIfDisconnected('Cannot disconnect when already disconnected from the server') this.#throwIfDisconnected('Cannot disconnect when already disconnected from the server')
this.#handleDisconnect(DisconnectReason.Generic) this.#handleDisconnect(DisconnectReason.PlannedDisconnect)
} }
/** /**
* Checks whether the client is ready * Checks whether the client is ready
* @returns Whether the client is ready * @returns Whether the client is ready
*/ */
isReady(): this is ReadiedClientGateway { isReady(): this is ReadiedClientWebSocketManager {
return this.ready return this.ready
} }
@@ -145,7 +157,7 @@ export default class ClientGateway {
return this.#handleDisconnect((packet as Packet<ServerOperation.Disconnect>).d.reason) return this.#handleDisconnect((packet as Packet<ServerOperation.Disconnect>).d.reason)
default: default:
return this.#emitter.emit( return this.#emitter.emit(
uncapitalize(ServerOperation[packet.op] as ClientGatewayServerEventName), uncapitalize(ServerOperation[packet.op] as ClientWebSocketEventName),
// @ts-expect-error TypeScript doesn't know that the lines above negate the type enough // @ts-expect-error TypeScript doesn't know that the lines above negate the type enough
packet, packet,
) )
@@ -162,6 +174,7 @@ export default class ClientGateway {
clearTimeout(this.#hbTimeout) clearTimeout(this.#hbTimeout)
this.disconnected = reason this.disconnected = reason
this.#socket.close() this.#socket.close()
this.#socket = null!
this.#emitter.emit('disconnect', reason) this.#emitter.emit('disconnect', reason)
} }
@@ -190,25 +203,30 @@ export default class ClientGateway {
} }
} }
export interface ClientGatewayOptions { export interface ClientWebSocketManagerOptions {
/** /**
* The gateway URL to connect to * The URL to connect to
*/ */
url: string url: string
/**
* The timeout for the connection
* @default 10000
*/
timeout?: number
} }
export type ClientGatewayServerEventName = keyof typeof ServerOperation export type ClientWebSocketEventName = keyof typeof ServerOperation
export type ClientGatewayEventHandlers = { export type ClientWebSocketEvents = {
[K in Uncapitalize<ClientGatewayServerEventName>]: ( [K in Uncapitalize<ClientWebSocketEventName>]: (
packet: Packet<(typeof ServerOperation)[Capitalize<K>]>, packet: Packet<(typeof ServerOperation)[Capitalize<K>]>,
) => Promise<void> | void ) => Promise<void> | void
} & { } & {
hello: (config: NonNullable<ClientGateway['config']>) => Promise<void> | void hello: (config: NonNullable<ClientWebSocketManager['config']>) => Promise<void> | void
ready: () => Promise<void> | void ready: () => Promise<void> | void
packet: (packet: Packet<ServerOperation>) => Promise<void> | void packet: (packet: Packet<ServerOperation>) => Promise<void> | void
invalidPacket: (packet: Packet) => Promise<void> | void invalidPacket: (packet: Packet) => Promise<void> | void
disconnect: (reason: DisconnectReason) => Promise<void> | void disconnect: (reason: DisconnectReason) => Promise<void> | void
} }
export type ReadiedClientGateway = RequiredProperty<InstanceType<typeof ClientGateway>> export type ReadiedClientWebSocketManager = RequiredProperty<InstanceType<typeof ClientWebSocketManager>>

View File

@@ -1,4 +1,3 @@
export { default as Client } from './Client' export { default as Client } from './Client'
export * from './Client' export * from './Client'
export { default as ClientGateway } from './ClientGateway' export * from './ClientWebSocket'
export * from './ClientGateway'

View File

@@ -22,6 +22,10 @@ enum DisconnectReason {
* The client had never connected to the server (**CLIENT-ONLY**) * The client had never connected to the server (**CLIENT-ONLY**)
*/ */
NeverConnected = 5, NeverConnected = 5,
/**
* The client disconnected on its own (**CLIENT-ONLY**)
*/
PlannedDisconnect = 6,
} }
export default DisconnectReason export default DisconnectReason

View File

@@ -9,6 +9,7 @@ const HumanizedDisconnectReason = {
[DisconnectReason.TimedOut]: 'has timed out', [DisconnectReason.TimedOut]: 'has timed out',
[DisconnectReason.ServerError]: 'has been disconnected due to an internal server error', [DisconnectReason.ServerError]: 'has been disconnected due to an internal server error',
[DisconnectReason.NeverConnected]: 'had never connected to the server', [DisconnectReason.NeverConnected]: 'had never connected to the server',
[DisconnectReason.PlannedDisconnect]: 'has disconnected on its own',
} as const satisfies Record<DisconnectReason, string> } as const satisfies Record<DisconnectReason, string>
export default HumanizedDisconnectReason export default HumanizedDisconnectReason