feat: Don't send outbound messages that exceed the max size (#3590)

* feat: Don't send outbound messages that exceed the max size

This adds a check to all outbound messages to ensure they fit within the
MAX_MESSAGE_SIZE. This ensures a node is not sending invalid messages,
intentionally or not.

* fix: re-add connection-specific metrics

* fix: add exhaustive enum checking for ConnectionType

* fix: lint and bind instance method

* add test for happy path

ironfish/src/assert.ts

@@ -5,8 +5,8 @@
 import { Constructor } from './utils/types'
 
 export class Assert {
-  static isUnreachable(x: never): never {
-    throw new Error(x)
+  static isUnreachable(x: never, message?: string): never {
+    throw new Error(message || `Expected ${String(x)} to be unreachable`)
   }
 
   static isUndefined<T>(x: undefined | T, message?: string): asserts x is undefined {

ironfish/src/network/peers/connections/connection.test.ts

@@ -0,0 +1,81 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+import ws from 'ws'
+import { createRootLogger } from '../../../logger'
+import { IdentifyMessage } from '../../messages/identify'
+import { defaultFeatures } from '../peerFeatures'
+import { ConnectionDirection } from './connection'
+import { WebSocketConnection } from './webSocketConnection'
+
+jest.mock('../../version', () => {
+  const moduleMock = jest.requireActual<typeof import('../../version')>('../../version')
+  return {
+    ...moduleMock,
+    MAX_MESSAGE_SIZE: 256,
+  }
+})
+
+jest.mock('ws')
+
+describe('Connection', () => {
+  afterAll(() => {
+    jest.unmock('ws')
+  })
+
+  describe('send', () => {
+    it('should send a message that is an acceptable size', () => {
+      const connection = new WebSocketConnection(
+        new ws(''),
+        ConnectionDirection.Outbound,
+        createRootLogger(),
+      )
+
+      const message = new IdentifyMessage({
+        agent: '',
+        head: Buffer.alloc(32, 0),
+        identity: Buffer.alloc(32, 'identity').toString('base64'),
+        port: 9033,
+        sequence: 1,
+        version: 0,
+        work: BigInt(0),
+        networkId: 0,
+        genesisBlockHash: Buffer.alloc(32, 0),
+        features: defaultFeatures(),
+      })
+
+      const _sendSpy = jest.spyOn(connection, '_send').mockImplementationOnce(() => true)
+
+      expect(connection.send(message)).toBe(true)
+      expect(_sendSpy).toHaveBeenCalled()
+      connection.close()
+    })
+
+    it('should not send a message that exceeds the maximum size', () => {
+      const connection = new WebSocketConnection(
+        new ws(''),
+        ConnectionDirection.Outbound,
+        createRootLogger(),
+      )
+
+      const message = new IdentifyMessage({
+        agent: Buffer.alloc(256, 0).toString(),
+        head: Buffer.alloc(32, 0),
+        identity: Buffer.alloc(32, 'identity').toString('base64'),
+        port: 9033,
+        sequence: 1,
+        version: 0,
+        work: BigInt(0),
+        networkId: 0,
+        genesisBlockHash: Buffer.alloc(32, 0),
+        features: defaultFeatures(),
+      })
+
+      const _sendSpy = jest.spyOn(connection, '_send').mockImplementationOnce(jest.fn())
+
+      expect(connection.send(message)).toBe(false)
+      expect(_sendSpy).not.toHaveBeenCalled()
+      connection.close()
+    })
+  })
+})

ironfish/src/network/peers/connections/connection.ts

@@ -3,13 +3,15 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 import type { Logger } from '../../../logger'
 import colors from 'colors/safe'
+import { Assert } from '../../../assert'
 import { Event } from '../../../event'
 import { MetricsMonitor } from '../../../metrics'
-import { SetTimeoutToken } from '../../../utils'
+import { ErrorUtils, SetTimeoutToken } from '../../../utils'
 import { Identity } from '../../identity'
-import { NetworkMessage } from '../../messages/networkMessage'
+import { displayNetworkMessageType, NetworkMessage } from '../../messages/networkMessage'
 import { RPC_TIMEOUT_MILLIS } from '../../messages/rpcNetworkMessage'
 import { NetworkMessageType } from '../../types'
+import { MAX_MESSAGE_SIZE } from '../../version'
 import { HandshakeTimeoutError } from './errors'
 
 /**
@@ -91,7 +93,7 @@ export abstract class Connection {
   /**
    * Send a message into this connection.
    */
-  abstract send: (object: NetworkMessage) => boolean
+  abstract _send: (data: Buffer) => boolean
 
   /**
    * Shutdown the connection, if possible
@@ -114,6 +116,54 @@ export abstract class Connection {
     this.simulateLatencyQueue = []
   }
 
+  send(object: NetworkMessage): boolean {
+    const data = object.serializeWithMetadata()
+    const byteCount = data.byteLength
+
+    if (byteCount >= MAX_MESSAGE_SIZE) {
+      this.logger.warn(
+        `Attempted to send a message that exceeds the maximum size. ${object.type} (${byteCount})`,
+      )
+      return false
+    }
+
+    if (this.shouldLogMessageType(object.type)) {
+      this.logger.debug(
+        `${colors.yellow('SEND')} ${this.displayName}: ${displayNetworkMessageType(
+          object.type,
+        )}`,
+      )
+    }
+
+    let sendResult
+    try {
+      sendResult = this._send(data)
+    } catch (e) {
+      this.logger.debug(
+        `Error occurred while sending ${displayNetworkMessageType(
+          object.type,
+        )} message in state ${this.state.type} ${ErrorUtils.renderError(e)}`,
+      )
+      this.close(e)
+      return false
+    }
+
+    if (sendResult) {
+      this.metrics?.p2p_OutboundTraffic.add(byteCount)
+      this.metrics?.p2p_OutboundTrafficByMessage.get(object.type)?.add(byteCount)
+
+      if (this.type === ConnectionType.WebRtc) {
+        this.metrics?.p2p_OutboundTraffic_WebRTC.add(byteCount)
+      } else if (this.type === ConnectionType.WebSocket) {
+        this.metrics?.p2p_OutboundTraffic_WS.add(byteCount)
+      } else {
+        Assert.isUnreachable(this.type)
+      }
+    }
+
+    return sendResult
+  }
+
   setState(state: Readonly<ConnectionState>): void {
     const prevState = this._state
     this._state = state
@@ -161,7 +211,7 @@ export abstract class Connection {
     if (!this.simulateLatency) {
       return
     }
-    const originalSend = this.send
+    const originalSend = this.send.bind(this)
 
     const wrapper = (
       ...args: Parameters<typeof originalSend>

ironfish/src/network/peers/connections/webRtcConnection.test.ts

@@ -15,7 +15,7 @@ describe('WebRtcConnection', () => {
         const message = new IdentifyMessage({
           agent: '',
           head: Buffer.alloc(32, 0),
-          identity: 'identity',
+          identity: Buffer.alloc(32, 'identity').toString('base64'),
           port: 9033,
           sequence: 1,
           version: 0,

ironfish/src/network/peers/connections/webRtcConnection.ts

@@ -10,7 +10,7 @@ import { Event } from '../../../event'
 import { MetricsMonitor } from '../../../metrics'
 import { ErrorUtils } from '../../../utils'
 import { parseNetworkMessage } from '../../messageRegistry'
-import { displayNetworkMessageType, NetworkMessage } from '../../messages/networkMessage'
+import { displayNetworkMessageType } from '../../messages/networkMessage'
 import { MAX_MESSAGE_SIZE } from '../../version'
 import { Connection, ConnectionDirection, ConnectionType } from './connection'
 import { NetworkError } from './errors'
@@ -197,42 +197,18 @@ export class WebRtcConnection extends Connection {
     }
   }
 
-  send = (message: NetworkMessage): boolean => {
+  _send = (data: Buffer): boolean => {
     if (!this.datachannel) {
       return false
     }
 
-    if (this.shouldLogMessageType(message.type)) {
-      this.logger.debug(
-        `${colors.yellow('SEND')} ${this.displayName}: ${displayNetworkMessageType(
-          message.type,
-        )}`,
-      )
-    }
-
     if (!this.datachannel.isOpen()) {
       this.logger.debug('Datachannel no longer open, closing connection')
       this.close()
       return false
     }
 
-    const data = message.serializeWithMetadata()
-    try {
-      this.datachannel.sendMessageBinary(data)
-    } catch (e) {
-      this.logger.debug(
-        `Error occurred while sending ${displayNetworkMessageType(
-          message.type,
-        )} message in state ${this.state.type} ${ErrorUtils.renderError(e)}`,
-      )
-      this.close(e)
-      return false
-    }
-
-    const byteCount = data.byteLength
-    this.metrics?.p2p_OutboundTraffic.add(byteCount)
-    this.metrics?.p2p_OutboundTraffic_WebRTC.add(byteCount)
-    this.metrics?.p2p_OutboundTrafficByMessage.get(message.type)?.add(byteCount)
+    this.datachannel.sendMessageBinary(data)
 
     return true
   }

ironfish/src/network/peers/connections/webSocketConnection.ts

@@ -6,7 +6,7 @@ import type { Logger } from '../../../logger'
 import colors from 'colors/safe'
 import { MetricsMonitor } from '../../../metrics'
 import { parseNetworkMessage } from '../../messageRegistry'
-import { displayNetworkMessageType, NetworkMessage } from '../../messages/networkMessage'
+import { displayNetworkMessageType } from '../../messages/networkMessage'
 import { IsomorphicWebSocket, IsomorphicWebSocketErrorEvent } from '../../types'
 import { Connection, ConnectionDirection, ConnectionType } from './connection'
 import { NetworkError } from './errors'
@@ -113,23 +113,9 @@ export class WebSocketConnection extends Connection {
     }
   }
 
-  send = (message: NetworkMessage): boolean => {
-    if (this.shouldLogMessageType(message.type)) {
-      this.logger.debug(
-        `${colors.yellow('SEND')} ${this.displayName}: ${displayNetworkMessageType(
-          message.type,
-        )}`,
-      )
-    }
-
-    const data = message.serializeWithMetadata()
+  _send = (data: Buffer): boolean => {
     this.socket.send(data)
 
-    const byteCount = data.byteLength
-    this.metrics?.p2p_OutboundTraffic.add(byteCount)
-    this.metrics?.p2p_OutboundTraffic_WS.add(byteCount)
-    this.metrics?.p2p_OutboundTrafficByMessage.get(message.type)?.add(byteCount)
-
     return true
   }