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.
iron-fish · Feb 28, 2023 · 2be999b · 2be999b
1 parent 41c9f43
commit 2be999b
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 45 deletions.
diff --git a/ironfish/src/network/peers/connections/connection.test.ts b/ironfish/src/network/peers/connections/connection.test.ts
@@ -0,0 +1,54 @@
+/* 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('WebSocketConnection', () => {
+  afterAll(() => {
+    jest.unmock('ws')
+  })
+
+  describe('send', () => {
+    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()
+    })
+  })
+})
diff --git a/ironfish/src/network/peers/connections/connection.ts b/ironfish/src/network/peers/connections/connection.ts
@@ -5,11 +5,12 @@ import type { Logger } from '../../../logger'
 import colors from 'colors/safe'
 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 +92,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 +115,47 @@ 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_OutboundTraffic_WebRTC.add(byteCount)
+      this.metrics?.p2p_OutboundTrafficByMessage.get(object.type)?.add(byteCount)
+    }
+
+    return sendResult
+  }
+
   setState(state: Readonly<ConnectionState>): void {
     const prevState = this._state
     this._state = state

diff --git a/ironfish/src/network/peers/connections/webRtcConnection.test.ts b/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,

diff --git a/ironfish/src/network/peers/connections/webRtcConnection.ts b/ironfish/src/network/peers/connections/webRtcConnection.ts
@@ -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
   }

diff --git a/ironfish/src/network/peers/connections/webSocketConnection.ts b/ironfish/src/network/peers/connections/webSocketConnection.ts
@@ -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
   }