Merge pull request #371 from TaloDev/maps-logging

Maps and logging

Author: tudddorrr

src/socket/index.ts

@@ -6,6 +6,7 @@ import authenticateSocket from './authenticateSocket'
 import SocketConnection from './socketConnection'
 import SocketRouter from './router/socketRouter'
 import { sendMessage } from './messages/socketMessage'
+import { logConnection, logConnectionClosed } from './messages/socketLogger'
 
 type CloseConnectionOptions = {
   code?: number
@@ -16,7 +17,7 @@ type CloseConnectionOptions = {
 
 export default class Socket {
   private readonly wss: WebSocketServer
-  private connections: SocketConnection[] = []
+  private connections: Map<WebSocket, SocketConnection> = new Map()
   private router: SocketRouter
 
   constructor(server: Server, private readonly em: EntityManager) {
@@ -39,32 +40,35 @@ export default class Socket {
     return this.wss
   }
 
+  /* v8 ignore start */
   heartbeat(): void {
     const interval = setInterval(() => {
-      this.connections.forEach((conn) => {
-        /* v8 ignore start */
+      for (const [ws, conn] of this.connections.entries()) {
         if (!conn.alive) {
-          this.closeConnection(conn.ws, { terminate: true })
-          return
+          this.closeConnection(ws, { terminate: true })
+          continue
         }
 
         conn.alive = false
-        conn.ws.ping()
-        /* v8 ignore end */
-      })
+        ws.ping()
+      }
     }, 30_000)
 
     this.wss.on('close', () => {
       clearInterval(interval)
     })
   }
+  /* v8 ignore stop */
 
   async handleConnection(ws: WebSocket, req: IncomingMessage): Promise<void> {
+    logConnection(req)
+
     await RequestContext.create(this.em, async () => {
       const key = await authenticateSocket(req.headers?.authorization ?? '')
       if (key) {
-        this.connections.push(new SocketConnection(ws, key, req))
-        sendMessage(this.connections.at(-1), 'v1.connected', {})
+        const connection = new SocketConnection(ws, key, req)
+        this.connections.set(ws, connection)
+        sendMessage(connection, 'v1.connected', {})
       } else {
         this.closeConnection(ws)
       }
@@ -73,48 +77,54 @@ export default class Socket {
 
   async handleMessage(ws: WebSocket, data: RawData): Promise<void> {
     await RequestContext.create(this.em, async () => {
-      await this.router.handleMessage(this.findConnectionBySocket(ws), data)
+      const connection = this.connections.get(ws)
+      if (connection) {
+        await this.router.handleMessage(connection, data)
+      /* v8 ignore next 3 */
+      } else {
+        this.closeConnection(ws)
+      }
     })
   }
 
   /* v8 ignore start */
   handlePong(ws: WebSocket): void {
-    const connection = this.findConnectionBySocket(ws)
+    const connection = this.findConnection(ws)
     if (!connection) return
 
     connection.alive = true
     if (connection.rateLimitWarnings > 0) {
       connection.rateLimitWarnings--
     }
   }
-  /* v8 ignore end */
+  /* v8 ignore stop */
 
   closeConnection(ws: WebSocket, options: CloseConnectionOptions = {}): void {
     const terminate = options.terminate ?? false
     const preclosed = options.preclosed ?? false
+    const code = options.code ?? 3000
 
+    /* v8 ignore next 2 */
     if (terminate) {
       ws.terminate()
     } else if (!preclosed) {
-      ws.close(options.code ?? 3000, options.reason)
+      ws.close(code, options.reason)
     }
 
-    this.connections = this.connections.filter((conn) => conn.ws !== ws)
-  }
+    const connection = this.findConnection(ws)
+    /* v8 ignore next */
+    if (!connection) return
 
-  findConnectionBySocket(ws: WebSocket): SocketConnection | undefined {
-    const connection = this.connections.find((conn) => conn.ws === ws)
-    /* v8 ignore start */
-    if (!connection) {
-      this.closeConnection(ws)
-      return
-    }
-    /* v8 ignore end */
+    logConnectionClosed(connection, preclosed, code, options.reason)
+
+    this.connections.delete(ws)
+  }
 
-    return connection
+  findConnection(ws: WebSocket): SocketConnection | undefined {
+    return this.connections.get(ws)
   }
 
   findConnections(filter: (conn: SocketConnection) => boolean): SocketConnection[] {
-    return this.connections.filter(filter)
+    return Array.from(this.connections.values()).filter(filter)
   }
 }

src/socket/messages/socketLogger.ts

@@ -0,0 +1,35 @@
+import SocketConnection from '../socketConnection'
+import { SocketMessageResponse } from './socketMessage'
+import { IncomingMessage } from 'http'
+
+function getSocketUrl(conn: SocketConnection | undefined): string {
+  if (!conn) {
+    return 'WSS /'
+  }
+  return `WSS /games/${conn.game.id}/${conn.playerAliasId ? `aliases/${conn.playerAliasId}/` : ''}`
+}
+
+function getSize(message: string): string {
+  return `${Buffer.byteLength(message)}b`
+}
+
+export function logRequest(conn: SocketConnection, message: string) {
+  const req = JSON.parse(message)?.req
+  console.log(`  <-- ${getSocketUrl(conn)}{${req}} ${conn.getRemoteAddress()} ${getSize(message)}`)
+}
+
+export function logResponse(conn: SocketConnection, res: SocketMessageResponse, message: string) {
+  console.log(`  --> ${getSocketUrl(conn)}{${res}} ${conn.getRemoteAddress()} ${getSize(message)}`)
+}
+
+export function logConnection(req: IncomingMessage) {
+  console.log(`  <-- WSS /open ${req.socket.remoteAddress}`)
+}
+
+export function logConnectionClosed(conn: SocketConnection | undefined, preclosed: boolean, code: number, reason?: string) {
+  const direction = preclosed ? '<--' : '-->'
+  const ip = conn?.getRemoteAddress() ?? 'unknown'
+  const displayCode = preclosed ? '' : code
+  const displayReason = reason ?? ''
+  console.log(`  ${direction} ${getSocketUrl(conn)}close ${ip} ${displayCode} ${displayReason}`.trimEnd())
+}

src/socket/messages/socketMessage.ts

@@ -1,4 +1,5 @@
 import SocketConnection from '../socketConnection'
+import { logResponse } from './socketLogger'
 
 export const requests = [
   'v1.players.identify',
@@ -20,10 +21,14 @@ export type SocketMessageResponse = typeof responses[number]
 
 export function sendMessage<T>(conn: SocketConnection, res: SocketMessageResponse, data: T) {
   if (conn.ws.readyState === conn.ws.OPEN) {
-    conn.ws.send(JSON.stringify({
+    const message = JSON.stringify({
       res,
       data
-    }))
+    })
+
+    logResponse(conn, res, message)
+
+    conn.ws.send(message)
   }
 }
 

src/socket/router/socketRouter.ts

@@ -9,6 +9,7 @@ import SocketError, { sendError } from '../messages/socketError'
 import { APIKeyScope } from '../../entities/api-key'
 import playerListeners from '../listeners/playerListeners'
 import gameChannelListeners from '../listeners/gameChannelListeners'
+import { logRequest } from '../messages/socketLogger'
 
 const socketMessageValidator = z.object({
   req: z.enum(requests),
@@ -26,6 +27,8 @@ export default class SocketRouter {
   constructor(readonly socket: Socket) {}
 
   async handleMessage(conn: SocketConnection, rawData: RawData): Promise<void> {
+    logRequest(conn, rawData.toString())
+
     addBreadcrumb({
       category: 'message',
       message: rawData.toString(),

src/socket/socketConnection.ts

@@ -16,6 +16,7 @@ export default class SocketConnection {
   game: Game | null = null
   scopes: APIKeyScope[] = []
   private headers: IncomingHttpHeaders = {}
+  private remoteAddress: string = 'unknown'
 
   rateLimitKey: string = v4()
   rateLimitWarnings: number = 0
@@ -24,6 +25,7 @@ export default class SocketConnection {
     this.game = apiKey.game
     this.scopes = apiKey.scopes
     this.headers = req.headers
+    this.remoteAddress = req.socket.remoteAddress
   }
 
   async getPlayerAlias(): Promise<PlayerAlias | null> {
@@ -67,4 +69,8 @@ export default class SocketConnection {
 
     return rateLimitExceeded
   }
+
+  getRemoteAddress(): string {
+    return this.remoteAddress
+  }
 }

tests/socket/messages/socketLogger.test.ts

@@ -0,0 +1,126 @@
+import { WebSocket } from 'ws'
+import TaloSocket from '../../../src/socket'
+import SocketConnection from '../../../src/socket/socketConnection'
+import createAPIKeyAndToken from '../../utils/createAPIKeyAndToken'
+import { IncomingMessage } from 'http'
+import { Socket } from 'net'
+import { logConnection, logConnectionClosed, logRequest, logResponse } from '../../../src/socket/messages/socketLogger'
+import { EntityManager } from '@mikro-orm/mysql'
+
+describe('Socket logger', () => {
+  const consoleMock = vi.spyOn(console, 'log').mockImplementation(() => undefined)
+
+  afterEach(() => {
+    consoleMock.mockReset()
+  })
+
+  async function createSocketConnection(): Promise<[TaloSocket, SocketConnection, () => void]>{
+    const [apiKey] = await createAPIKeyAndToken([])
+    await (<EntityManager>global.em).persistAndFlush(apiKey)
+
+    const socket = new TaloSocket(global.server, global.em)
+    const conn = new SocketConnection(new WebSocket(null, [], {}), apiKey, new IncomingMessage(new Socket()))
+    vi.spyOn(conn, 'getRemoteAddress').mockReturnValue('0.0.0.0')
+
+    return [
+      socket,
+      conn,
+      () => {
+        socket.getServer().close()
+      }
+    ]
+  }
+
+  it('should log requests', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+
+    logRequest(conn, JSON.stringify({ req: 'v1.fake', data: {} }))
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  <-- WSS /games/${conn.game.id}/{v1.fake} 0.0.0.0 27b`)
+
+    cleanup()
+  })
+
+  it('should log requests with aliases', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+    conn.playerAliasId = 2
+
+    logRequest(conn, JSON.stringify({ req: 'v1.fake', data: {} }))
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  <-- WSS /games/${conn.game.id}/aliases/2/{v1.fake} 0.0.0.0 27b`)
+
+    cleanup()
+  })
+
+  it('should log responses', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+
+    logResponse(conn, 'v1.players.identify.success', JSON.stringify({ res: 'v1.players.identify.success', data: {} }))
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  --> WSS /games/${conn.game.id}/{v1.players.identify.success} 0.0.0.0 47b`)
+
+    cleanup()
+  })
+
+  it('should log connections', async () => {
+    logConnection(new IncomingMessage(new Socket()))
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith('  <-- WSS /open undefined')
+  })
+
+  it('should log pre-closed connections', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+
+    logConnectionClosed(conn, true, 3000)
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  <-- WSS /games/${conn.game.id}/close 0.0.0.0`)
+
+    cleanup()
+  })
+
+  it('should log manually-closed connections', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+
+    logConnectionClosed(conn, false, 3000, 'Unauthorised')
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  --> WSS /games/${conn.game.id}/close 0.0.0.0 3000 Unauthorised`)
+
+    cleanup()
+  })
+
+  it('should log manually-closed connections without a reason', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+
+    logConnectionClosed(conn, false, 3000)
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  --> WSS /games/${conn.game.id}/close 0.0.0.0 3000`)
+
+    cleanup()
+  })
+
+  it('should log manually-closed connections without a SocketConnection', async () => {
+    logConnectionClosed(undefined, false, 3000)
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith('  --> WSS /close unknown 3000')
+  })
+
+  it('should log pre-closed connection with aliases', async () => {
+    const [, conn, cleanup] = await createSocketConnection()
+    conn.playerAliasId = 2
+
+    logConnectionClosed(conn, true, 3000)
+
+    expect(consoleMock).toHaveBeenCalledOnce()
+    expect(consoleMock).toHaveBeenLastCalledWith(`  <-- WSS /games/${conn.game.id}/aliases/2/close 0.0.0.0`)
+
+    cleanup()
+  })
+})