From d7b8447e0e9ab65cb4bc14e276a156a727a7bba1 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 2 May 2025 00:35:11 +0800
Subject: [PATCH] fix dequeuing logic for tiny CRYPTO frames (#5104)

For very small sizes, cryptoStream.PopCryptoStream could
have returned CRYPTO frames larger than the requested size.

Instead, it should return a nil frame.
---
 connection.go         |  4 +++-
 crypto_stream.go      |  3 +++
 crypto_stream_test.go | 14 ++++++++++----
 packet_packer.go      |  7 ++++---
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/connection.go b/connection.go
index 5555cfd1..79506fed 100644
--- a/connection.go
+++ b/connection.go
@@ -826,7 +826,9 @@ func (s *connection) handleHandshakeComplete(now time.Time) error {
 	if ticket != nil { // may be nil if session tickets are disabled via tls.Config.SessionTicketsDisabled
 		s.oneRTTStream.Write(ticket)
 		for s.oneRTTStream.HasData() {
-			s.queueControlFrame(s.oneRTTStream.PopCryptoFrame(protocol.MaxPostHandshakeCryptoFrameSize))
+			if cf := s.oneRTTStream.PopCryptoFrame(protocol.MaxPostHandshakeCryptoFrameSize); cf != nil {
+				s.queueControlFrame(cf)
+			}
 		}
 	}
 	token, err := s.tokenGenerator.NewToken(s.conn.RemoteAddr())
diff --git a/crypto_stream.go b/crypto_stream.go
index 9a387baa..443a07e1 100644
--- a/crypto_stream.go
+++ b/crypto_stream.go
@@ -76,6 +76,9 @@ func (s *cryptoStream) HasData() bool {
 func (s *cryptoStream) PopCryptoFrame(maxLen protocol.ByteCount) *wire.CryptoFrame {
 	f := &wire.CryptoFrame{Offset: s.writeOffset}
 	n := min(f.MaxDataLen(maxLen), protocol.ByteCount(len(s.writeBuf)))
+	if n == 0 {
+		return nil
+	}
 	f.Data = s.writeBuf[:n]
 	s.writeBuf = s.writeBuf[n:]
 	s.writeOffset += n
diff --git a/crypto_stream_test.go b/crypto_stream_test.go
index 6b548882..6702b488 100644
--- a/crypto_stream_test.go
+++ b/crypto_stream_test.go
@@ -91,11 +91,17 @@ func TestCryptoStreamWrite(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, str.HasData())
 
-	f := str.PopCryptoFrame(expectedCryptoFrameLen(0) + 3)
-	require.Equal(t, &wire.CryptoFrame{Data: []byte("foo")}, f)
+	for i := range expectedCryptoFrameLen(0) {
+		require.Nil(t, str.PopCryptoFrame(i))
+	}
+
+	f := str.PopCryptoFrame(expectedCryptoFrameLen(0) + 1)
+	require.Equal(t, &wire.CryptoFrame{Data: []byte("f")}, f)
 	require.True(t, str.HasData())
+	f = str.PopCryptoFrame(expectedCryptoFrameLen(1) + 3)
+	// the three write calls were coalesced into a single frame
+	require.Equal(t, &wire.CryptoFrame{Offset: 1, Data: []byte("oob")}, f)
 	f = str.PopCryptoFrame(protocol.MaxByteCount)
-	// the two write calls were coalesced into a single frame
-	require.Equal(t, &wire.CryptoFrame{Offset: 3, Data: []byte("barbaz")}, f)
+	require.Equal(t, &wire.CryptoFrame{Offset: 4, Data: []byte("arbaz")}, f)
 	require.False(t, str.HasData())
 }
diff --git a/packet_packer.go b/packet_packer.go
index e84a6c02..877ac672 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -561,9 +561,10 @@ func (p *packetPacker) maybeGetCryptoPacket(
 			maxPacketSize -= frameLen
 		}
 	} else if s.HasData() {
-		cf := s.PopCryptoFrame(maxPacketSize)
-		pl.frames = append(pl.frames, ackhandler.Frame{Frame: cf, Handler: handler})
-		pl.length += cf.Length(v)
+		if cf := s.PopCryptoFrame(maxPacketSize); cf != nil {
+			pl.frames = append(pl.frames, ackhandler.Frame{Frame: cf, Handler: handler})
+			pl.length += cf.Length(v)
+		}
 	}
 	return hdr, pl
 }