fix: add timeouts to rpcs

Author: adiwajshing

proto/api.proto

@@ -50,6 +50,8 @@ enum ErrorCode {
 	ERROR_BGP_ANNOUNCEMENT_OVERLAP = 8;
 	// authentication failed
 	ERROR_AUTHENTICATION_FAILED = 9;
+	// timed out waiting for the server to respond
+	ERROR_TIMEOUT = 10;
 }
 
 enum ZKProofEngine {

src/client/utils/client-socket.ts

@@ -1,8 +1,8 @@
 import { base64 } from 'ethers/lib/utils'
-import { DEFAULT_METADATA } from 'src/config'
+import { DEFAULT_METADATA, DEFAULT_RPC_TIMEOUT_MS } from 'src/config'
 import { InitResponse, RPCMessages } from 'src/proto/api'
 import { IAttestorClient, IAttestorClientCreateOpts, RPCEvent, RPCRequestData, RPCResponseData, RPCType } from 'src/types'
-import { AttestorError, getRpcRequestType, logger as LOGGER, packRpcMessages } from 'src/utils'
+import { AttestorError, generateRpcMessageId, getRpcRequestType, logger as LOGGER, packRpcMessages } from 'src/utils'
 import { AttestorSocket } from 'src/utils/socket-base'
 import { makeWebSocket as defaultMakeWebSocket } from 'src/utils/ws'
 
@@ -40,7 +40,7 @@ export class AttestorClient extends AttestorSocket implements IAttestorClient {
 
 		const initReqId = msg.messages[0].id
 		this.waitForInitPromise = this
-			.waitForResponse<'init'>(initReqId)
+			.waitForResponse<'init'>(initReqId, DEFAULT_RPC_TIMEOUT_MS)
 			.then(res => {
 				logger.info('client initialised')
 				this.isInitialised = true
@@ -58,15 +58,15 @@ export class AttestorClient extends AttestorSocket implements IAttestorClient {
 
 	async rpc<T extends RPCType>(
 		type: T,
-		request: Partial<RPCRequestData<T>>
+		request: Partial<RPCRequestData<T>>,
+		timeoutMs: number = DEFAULT_RPC_TIMEOUT_MS
 	) {
 		this.logger.debug({ type }, 'sending rpc request')
 		const now = Date.now()
 		try {
-			const {
-				messages: [{ id }]
-			} = await this.sendMessage({ [getRpcRequestType(type)]: request })
-			const rslt = await this.waitForResponse<T>(id)
+			const msgId = generateRpcMessageId()
+			const rslt = await this.waitForResponse<T>(msgId, timeoutMs)
+			await this.sendMessage({ id: msgId, [getRpcRequestType(type)]: request })
 
 			return rslt
 		} finally {
@@ -75,7 +75,10 @@ export class AttestorClient extends AttestorSocket implements IAttestorClient {
 		}
 	}
 
-	waitForResponse<T extends RPCType>(id: number) {
+	waitForResponse<T extends RPCType>(
+		id: number,
+		timeoutMs: number = DEFAULT_RPC_TIMEOUT_MS
+	) {
 		if(this.isClosed) {
 			throw new AttestorError(
 				'ERROR_NETWORK_ERROR',
@@ -118,7 +121,19 @@ export class AttestorClient extends AttestorSocket implements IAttestorClient {
 				reject(event.data)
 			}
 
+			const timeout = setTimeout(() => {
+				removeHandlers()
+				reject(
+					new AttestorError(
+						'ERROR_TIMEOUT',
+						`RPC request timed out after ${timeoutMs}ms`,
+						{ id }
+					)
+				)
+			}, timeoutMs)
+
 			const removeHandlers = () => {
+				clearTimeout(timeout)
 				this.removeEventListener('rpc-response', handler)
 				this.removeEventListener('connection-terminated', terminateHandler)
 			}

src/client/utils/message-handler.ts

@@ -12,6 +12,7 @@ export async function wsMessageHandler(this: IAttestorSocket, data: unknown) {
 }
 
 export function handleMessage(this: IAttestorSocket, msg: RPCMessage) {
+	this.logger?.trace({ msg }, 'received message')
 	// handle connection termination alert
 	if(msg.connectionTerminationAlert) {
 		const err = AttestorError.fromProto(

src/config/index.ts

@@ -47,6 +47,8 @@ export const MAX_PAYLOAD_SIZE = 512 * 1024 * 1024 // 512MB
 
 export const DEFAULT_AUTH_EXPIRY_S = 15 * 60 // 15m
 
+export const DEFAULT_RPC_TIMEOUT_MS = 90_000
+
 export const TOPRF_DOMAIN_SEPARATOR = 'reclaim-toprf'
 
 export const BGP_WS_URL = 'wss://ris-live.ripe.net/v1/ws/?client=reclaim-hijack-detector'

src/proto/api.ts

@@ -1,4 +1,5 @@
 import { handleMessage } from 'src/client/utils/message-handler'
+import { DEFAULT_RPC_TIMEOUT_MS } from 'src/config'
 import { TunnelMessage } from 'src/proto/api'
 import { HANDLERS } from 'src/server/handlers'
 import { getApm } from 'src/server/utils/apm'
@@ -102,24 +103,15 @@ async function handleTunnelMessage(
 		const tunnel = this.getTunnel(tunnelId)
 		await tunnel.write(message)
 	} catch(err) {
-		this.logger?.error(
-			{
-				err,
-				tunnelId,
-			},
-			'error writing to tunnel'
-		)
+		this.logger?.error({ err, tunnelId }, 'error writing to tunnel')
 	}
 }
 
 async function handleRpcRequest(
 	this: IAttestorServerSocket,
 	{ data: { data, requestId, respond, type } }: RPCEvent<'rpc-request'>
 ) {
-	const logger = this.logger.child({
-		rpc: type,
-		requestId
-	})
+	const logger = this.logger.child({ rpc: type, requestId })
 
 	const apm = getApm()
 	const tx = apm?.startTransaction(type)
@@ -131,6 +123,10 @@ async function handleRpcRequest(
 		tx?.setLabel('authUserId', userId)
 	}
 
+	const timeout = setTimeout(() => {
+		logger.warn({ type, requestId }, 'RPC took too long to respond')
+	}, DEFAULT_RPC_TIMEOUT_MS)
+
 	try {
 		logger.debug({ data }, 'handling RPC request')
 
@@ -147,6 +143,7 @@ async function handleRpcRequest(
 
 		apm?.captureError(err, { parent: tx })
 	} finally {
+		clearTimeout(timeout)
 		tx?.end()
 	}
 }

src/server/socket.ts

@@ -167,14 +167,16 @@ export declare class IAttestorClient extends IAttestorSocket {
 	 * the promise will reject.
 	 */
 	waitForResponse<T extends RPCType>(
-		id: number
+		id: number,
+		timeoutMs?: number
 	): Promise<RPCResponseData<T>>
 	/**
 	 * Make an RPC request to the other end of the WebSocket.
 	 */
 	rpc<T extends RPCType>(
 		type: T,
-		request: Partial<RPCRequestData<T>>
+		request: Partial<RPCRequestData<T>>,
+		timeoutMs?: number
 	): Promise<RPCResponseData<T>>
 	/**
 	 * Waits for the "init" request to be responded to

src/types/client.ts

@@ -75,6 +75,8 @@ export class AttestorSocket implements IAttestorSocket {
 		const msg = packRpcMessages(...msgs)
 		const bytes = RPCMessages.encode(msg).finish()
 
+		this.logger.trace({ msg }, 'sending messages')
+
 		if('sendPromise' in this.socket && this.socket.sendPromise) {
 			await this.socket.sendPromise(bytes)
 		} else {