From 2be999b8dc531d8a8c712e2480c5d836afd3c856 Mon Sep 17 00:00:00 2001 From: Mat Date: Tue, 28 Feb 2023 10:56:26 -0700 Subject: [PATCH] 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. --- .../peers/connections/connection.test.ts | 54 +++++++++++++++++++ .../network/peers/connections/connection.ts | 48 +++++++++++++++-- .../connections/webRtcConnection.test.ts | 2 +- .../peers/connections/webRtcConnection.ts | 28 +--------- .../peers/connections/webSocketConnection.ts | 16 +----- 5 files changed, 103 insertions(+), 45 deletions(-) create mode 100644 ironfish/src/network/peers/connections/connection.test.ts diff --git a/ironfish/src/network/peers/connections/connection.test.ts b/ironfish/src/network/peers/connections/connection.test.ts new file mode 100644 index 0000000000..670e530c2e --- /dev/null +++ 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('../../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 index 5f6b5433e5..10bbda5132 100644 --- 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): 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 index 5961b4ce41..38352b04fd 100644 --- 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 index 074630fe78..6c5f8128c4 100644 --- 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 index a6cb5e1ffc..7d9d3d30ab 100644 --- 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 }