From 9d18d3093126d2b06f4d6f69e47c93459e512434 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 10 Mar 2018 14:27:07 +0700
Subject: [PATCH] use the peer's max_packet_size when packing packets

---
 packet_packer.go      |  5 +++++
 packet_packer_test.go | 35 ++++++++++++++++++++++++++++++++++-
 session.go            |  3 +++
 session_test.go       |  2 ++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/packet_packer.go b/packet_packer.go
index 22a0adc3..72b5b9a0 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -10,6 +10,7 @@ import (
 	"github.com/lucas-clemente/quic-go/internal/ackhandler"
 	"github.com/lucas-clemente/quic-go/internal/handshake"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
+	"github.com/lucas-clemente/quic-go/internal/utils"
 	"github.com/lucas-clemente/quic-go/internal/wire"
 )
 
@@ -518,3 +519,7 @@ func (p *packetPacker) canSendData(encLevel protocol.EncryptionLevel) bool {
 func (p *packetPacker) SetOmitConnectionID() {
 	p.omitConnectionID = true
 }
+
+func (p *packetPacker) SetMaxPacketSize(size protocol.ByteCount) {
+	p.maxPacketSize = utils.MinByteCount(p.maxPacketSize, size)
+}
diff --git a/packet_packer_test.go b/packet_packer_test.go
index 7ff37ba6..8c6bea9d 100644
--- a/packet_packer_test.go
+++ b/packet_packer_test.go
@@ -895,7 +895,6 @@ var _ = Describe("Packet packer", func() {
 			Expect(sf2.StreamID).To(Equal(protocol.StreamID(5)))
 			Expect(sf2.DataLenPresent).To(BeFalse())
 		})
-
 	})
 
 	Context("packing ACK packets", func() {
@@ -917,4 +916,38 @@ var _ = Describe("Packet packer", func() {
 			}))
 		})
 	})
+
+	Context("max packet size", func() {
+		It("sets the maximum packet size", func() {
+			for i := 0; i < 10*int(maxPacketSize); i++ {
+				packer.QueueControlFrame(&wire.PingFrame{})
+			}
+			mockStreamFramer.EXPECT().HasCryptoStreamData().AnyTimes()
+			mockStreamFramer.EXPECT().PopStreamFrames(gomock.Any()).AnyTimes()
+			p, err := packer.PackPacket()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(p.raw).To(HaveLen(int(maxPacketSize)))
+			// now reduce the maxPacketSize
+			packer.SetMaxPacketSize(maxPacketSize - 10)
+			p, err = packer.PackPacket()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(p.raw).To(HaveLen(int(maxPacketSize) - 10))
+		})
+
+		It("doesn't increase the max packet size", func() {
+			for i := 0; i < 10*int(maxPacketSize); i++ {
+				packer.QueueControlFrame(&wire.PingFrame{})
+			}
+			mockStreamFramer.EXPECT().HasCryptoStreamData().AnyTimes()
+			mockStreamFramer.EXPECT().PopStreamFrames(gomock.Any()).AnyTimes()
+			p, err := packer.PackPacket()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(p.raw).To(HaveLen(int(maxPacketSize)))
+			// now try to increase the maxPacketSize
+			packer.SetMaxPacketSize(maxPacketSize + 10)
+			p, err = packer.PackPacket()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(p.raw).To(HaveLen(int(maxPacketSize)))
+		})
+	})
 })
diff --git a/session.go b/session.go
index d9e5bf13..e2b83c1d 100644
--- a/session.go
+++ b/session.go
@@ -759,6 +759,9 @@ func (s *session) processTransportParameters(params *handshake.TransportParamete
 	if params.OmitConnectionID {
 		s.packer.SetOmitConnectionID()
 	}
+	if params.MaxPacketSize != 0 {
+		s.packer.SetMaxPacketSize(params.MaxPacketSize)
+	}
 	s.connFlowController.UpdateSendWindow(params.ConnectionFlowControlWindow)
 	// the crypto stream is the only open stream at this moment
 	// so we don't need to update stream flow control windows
diff --git a/session_test.go b/session_test.go
index 019f6562..0e16d1a5 100644
--- a/session_test.go
+++ b/session_test.go
@@ -1361,11 +1361,13 @@ var _ = Describe("Session", func() {
 			StreamFlowControlWindow:     0x5000,
 			ConnectionFlowControlWindow: 0x5000,
 			OmitConnectionID:            true,
+			MaxPacketSize:               0x42,
 		}
 		streamManager.EXPECT().UpdateLimits(&params)
 		paramsChan <- params
 		Eventually(func() *handshake.TransportParameters { return sess.peerParams }).Should(Equal(&params))
 		Eventually(func() bool { return sess.packer.omitConnectionID }).Should(BeTrue())
+		Eventually(func() protocol.ByteCount { return sess.packer.maxPacketSize }).Should(Equal(protocol.ByteCount(0x42)))
 		// make the go routine return
 		streamManager.EXPECT().CloseWithError(gomock.Any())
 		Expect(sess.Close(nil)).To(Succeed())