From 681e2bfa2f0546738fa569dfc2aca2a731a9acb9 Mon Sep 17 00:00:00 2001
From: cthuang <chungting@cloudflare.com>
Date: Mon, 29 Jul 2024 11:39:10 -0500
Subject: [PATCH] Allow server to specify custom max UDP payload size

---
 connection.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/connection.go b/connection.go
index d65f7c2ef12..0dbfa574019 100644
--- a/connection.go
+++ b/connection.go
@@ -8,7 +8,9 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"os"
 	"reflect"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -288,6 +290,16 @@ var newConnection = func(
 		s.logger,
 	)
 	s.maxPayloadSizeEstimate.Store(uint32(estimateMaxPayloadSize(protocol.ByteCount(s.config.InitialPacketSize))))
+	// Allow server to define custom MaxUDPPayloadSize
+	maxUDPPayloadSize := protocol.MaxPacketBufferSize
+	if maxPacketSize := os.Getenv("TUNNEL_MAX_QUIC_PACKET_SIZE"); maxPacketSize != "" {
+		if customMaxPacketSize, err := strconv.ParseUint(maxPacketSize, 10, 64); err == nil {
+			maxUDPPayloadSize = int(customMaxPacketSize)
+		} else {
+			utils.DefaultLogger.Errorf("failed to parse TUNNEL_MAX_QUIC_PACKET_SIZE: %v", err)
+		}
+	}
+
 	params := &wire.TransportParameters{
 		InitialMaxStreamDataBidiLocal:   protocol.ByteCount(s.config.InitialStreamReceiveWindow),
 		InitialMaxStreamDataBidiRemote:  protocol.ByteCount(s.config.InitialStreamReceiveWindow),
@@ -298,7 +310,7 @@ var newConnection = func(
 		MaxUniStreamNum:                 protocol.StreamNum(s.config.MaxIncomingUniStreams),
 		MaxAckDelay:                     protocol.MaxAckDelayInclGranularity,
 		AckDelayExponent:                protocol.AckDelayExponent,
-		MaxUDPPayloadSize:               protocol.MaxPacketBufferSize,
+		MaxUDPPayloadSize:               protocol.ByteCount(maxUDPPayloadSize),
 		DisableActiveMigration:          true,
 		StatelessResetToken:             &statelessResetToken,
 		OriginalDestinationConnectionID: origDestConnID,