From cee3c519e18232f5e39ed5f8c5679274989d3459 Mon Sep 17 00:00:00 2001
From: achingbrain <alex@achingbrain.net>
Date: Wed, 11 Dec 2024 18:43:18 +0000
Subject: [PATCH] fix: increase max autonat streams and limit incoming message
 size

Allows more outgoing streams to be open at the same time and limits
the allowed incoming message size.
---
 packages/protocol-autonat/src/autonat.ts   | 11 ++++++-----
 packages/protocol-autonat/src/constants.ts |  5 +++--
 packages/protocol-autonat/src/index.ts     |  8 ++++++++
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/packages/protocol-autonat/src/autonat.ts b/packages/protocol-autonat/src/autonat.ts
index 99e65cf378..bb2e5bc8f9 100644
--- a/packages/protocol-autonat/src/autonat.ts
+++ b/packages/protocol-autonat/src/autonat.ts
@@ -10,10 +10,7 @@ import { multiaddr, protocols } from '@multiformats/multiaddr'
 import { anySignal } from 'any-signal'
 import { pbStream } from 'it-protobuf-stream'
 import * as Digest from 'multiformats/hashes/digest'
-import {
-  DEFAULT_CONNECTION_THRESHOLD,
-  MAX_INBOUND_STREAMS, MAX_OUTBOUND_STREAMS, PROTOCOL_NAME, PROTOCOL_PREFIX, PROTOCOL_VERSION, TIMEOUT
-} from './constants.js'
+import { DEFAULT_CONNECTION_THRESHOLD, MAX_INBOUND_STREAMS, MAX_MESSAGE_SIZE, MAX_OUTBOUND_STREAMS, PROTOCOL_NAME, PROTOCOL_PREFIX, PROTOCOL_VERSION, TIMEOUT } from './constants.js'
 import { Message } from './pb/index.js'
 import type { AutoNATComponents, AutoNATServiceInit } from './index.js'
 import type { Logger, Connection, PeerId, Startable, AbortOptions } from '@libp2p/interface'
@@ -89,6 +86,7 @@ export class AutoNATService implements Startable {
   private readonly timeout: number
   private readonly maxInboundStreams: number
   private readonly maxOutboundStreams: number
+  private readonly maxMessageSize: number
   private started: boolean
   private readonly log: Logger
   private topologyId?: string
@@ -106,6 +104,7 @@ export class AutoNATService implements Startable {
     this.maxInboundStreams = init.maxInboundStreams ?? MAX_INBOUND_STREAMS
     this.maxOutboundStreams = init.maxOutboundStreams ?? MAX_OUTBOUND_STREAMS
     this.connectionThreshold = init.connectionThreshold ?? DEFAULT_CONNECTION_THRESHOLD
+    this.maxMessageSize = init.maxMessageSize ?? MAX_MESSAGE_SIZE
     this.dialResults = new Map()
     this.findPeers = repeatingTask(this.findRandomPeers.bind(this), 60_000)
     this.addressFilter = createScalableCuckooFilter(1024)
@@ -229,7 +228,9 @@ export class AutoNATService implements Startable {
     const signal = AbortSignal.timeout(this.timeout)
     setMaxListeners(Infinity, signal)
 
-    const messages = pbStream(data.stream).pb(Message)
+    const messages = pbStream(data.stream, {
+      maxDataLength: this.maxMessageSize
+    }).pb(Message)
 
     try {
       const request = await messages.read({
diff --git a/packages/protocol-autonat/src/constants.ts b/packages/protocol-autonat/src/constants.ts
index 2a437b39b2..2b3b476a54 100644
--- a/packages/protocol-autonat/src/constants.ts
+++ b/packages/protocol-autonat/src/constants.ts
@@ -13,6 +13,7 @@ export const PROTOCOL_NAME = 'autonat'
  */
 export const PROTOCOL_VERSION = '1.0.0'
 export const TIMEOUT = 30000
-export const MAX_INBOUND_STREAMS = 1
-export const MAX_OUTBOUND_STREAMS = 1
+export const MAX_INBOUND_STREAMS = 2
+export const MAX_OUTBOUND_STREAMS = 20
 export const DEFAULT_CONNECTION_THRESHOLD = 80
+export const MAX_MESSAGE_SIZE = 8192
diff --git a/packages/protocol-autonat/src/index.ts b/packages/protocol-autonat/src/index.ts
index 0f852ed639..31d7f1687e 100644
--- a/packages/protocol-autonat/src/index.ts
+++ b/packages/protocol-autonat/src/index.ts
@@ -75,6 +75,14 @@ export interface AutoNATServiceInit {
    * @default 80
    */
   connectionThreshold?: number
+
+  /**
+   * How large incoming autonat messages are allowed to be in bytes. If messages
+   * larger than this are received the stream will be reset.
+   *
+   * @default 8192
+   */
+  maxMessageSize?: number
 }
 
 export interface AutoNATComponents {