From 142071253d0dc50a1c4d78f5a17f0b75fe49b636 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 31 Jan 2021 14:25:18 +0800
Subject: [PATCH] add a method to update the packer's max packet size

---
 packet_packer.go      |  7 +++++++
 packet_packer_test.go | 31 +++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/packet_packer.go b/packet_packer.go
index 00445472..4eef7497 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -862,6 +862,13 @@ func (p *packetPacker) SetToken(token []byte) {
 	p.token = token
 }
 
+// When a higher MTU is discovered, use it.
+func (p *packetPacker) SetMaxPacketSize(s protocol.ByteCount) {
+	p.maxPacketSize = s
+}
+
+// If the peer sets a max_packet_size that's smaller than the size we're currently using,
+// we need to reduce the size of packets we send.
 func (p *packetPacker) HandleTransportParameters(params *wire.TransportParameters) {
 	if params.MaxUDPPayloadSize != 0 {
 		p.maxPacketSize = utils.MinByteCount(p.maxPacketSize, params.MaxUDPPayloadSize)
diff --git a/packet_packer_test.go b/packet_packer_test.go
index 847480c7..a8341ac8 100644
--- a/packet_packer_test.go
+++ b/packet_packer_test.go
@@ -802,8 +802,8 @@ var _ = Describe("Packet packer", func() {
 				})
 			})
 
-			Context("max packet size", func() {
-				It("sets the maximum packet size", func() {
+			Context("handling transport parameters", func() {
+				It("lowers the maximum packet size", func() {
 					pnManager.EXPECT().PeekPacketNumber(protocol.Encryption1RTT).Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2).Times(2)
 					sealingManager.EXPECT().Get1RTTSealer().Return(getSealer(), nil).Times(2)
 					framer.EXPECT().HasData().Return(true).Times(2)
@@ -855,6 +855,33 @@ var _ = Describe("Packet packer", func() {
 					Expect(err).ToNot(HaveOccurred())
 				})
 			})
+
+			Context("max packet size", func() {
+				It("increases the max packet size", func() {
+					pnManager.EXPECT().PeekPacketNumber(protocol.Encryption1RTT).Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2).Times(2)
+					sealingManager.EXPECT().Get1RTTSealer().Return(getSealer(), nil).Times(2)
+					framer.EXPECT().HasData().Return(true).Times(2)
+					ackFramer.EXPECT().GetAckFrame(protocol.Encryption1RTT, false).Times(2)
+					var initialMaxPacketSize protocol.ByteCount
+					framer.EXPECT().AppendControlFrames(gomock.Any(), gomock.Any()).Do(func(_ []ackhandler.Frame, maxLen protocol.ByteCount) ([]ackhandler.Frame, protocol.ByteCount) {
+						initialMaxPacketSize = maxLen
+						return nil, 0
+					})
+					expectAppendStreamFrames()
+					_, err := packer.PackPacket()
+					Expect(err).ToNot(HaveOccurred())
+					// now reduce the maxPacketSize
+					const packetSizeIncrease = 50
+					packer.SetMaxPacketSize(maxPacketSize + packetSizeIncrease)
+					framer.EXPECT().AppendControlFrames(gomock.Any(), gomock.Any()).Do(func(_ []ackhandler.Frame, maxLen protocol.ByteCount) ([]ackhandler.Frame, protocol.ByteCount) {
+						Expect(maxLen).To(Equal(initialMaxPacketSize + packetSizeIncrease))
+						return nil, 0
+					})
+					expectAppendStreamFrames()
+					_, err = packer.PackPacket()
+					Expect(err).ToNot(HaveOccurred())
+				})
+			})
 		})
 
 		Context("packing crypto packets", func() {