From b4213ca09ee0645662324b1cd9d6a33df24dd41e Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 6 Sep 2023 11:18:55 +0700
Subject: [PATCH] fix panics when processing post-handshake messages

See https://go-review.googlesource.com/c/go/+/522595.
---
 quic.go      | 10 +++++--
 quic_test.go | 80 +++++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 74 insertions(+), 16 deletions(-)

diff --git a/quic.go b/quic.go
index 6be1e91..f146688 100644
--- a/quic.go
+++ b/quic.go
@@ -230,16 +230,22 @@ func (q *QUICConn) HandleData(level QUICEncryptionLevel, data []byte) error {
 		return nil
 	}
 	// The handshake goroutine has exited.
+	c.handshakeMutex.Lock()
+	defer c.handshakeMutex.Unlock()
 	c.hand.Write(c.quic.readbuf)
 	c.quic.readbuf = nil
 	for q.conn.hand.Len() >= 4 && q.conn.handshakeErr == nil {
 		b := q.conn.hand.Bytes()
 		n := int(b[1])<<16 | int(b[2])<<8 | int(b[3])
-		if 4+n < len(b) {
+		if n > maxHandshake {
+			q.conn.handshakeErr = fmt.Errorf("tls: handshake message of length %d bytes exceeds maximum of %d bytes", n, maxHandshake)
+			break
+		}
+		if len(b) < 4+n {
 			return nil
 		}
 		if err := q.conn.handlePostHandshakeMessage(); err != nil {
-			return quicError(err)
+			q.conn.handshakeErr = err
 		}
 	}
 	if q.conn.handshakeErr != nil {
diff --git a/quic_test.go b/quic_test.go
index 2563053..a02cd5a 100644
--- a/quic_test.go
+++ b/quic_test.go
@@ -85,7 +85,7 @@ func (q *testQUICConn) setWriteSecret(level QUICEncryptionLevel, suite uint16, s
 
 var errTransportParametersRequired = errors.New("transport parameters required")
 
-func runTestQUICConnection(ctx context.Context, cli, srv *testQUICConn, onHandleCryptoData func()) error {
+func runTestQUICConnection(ctx context.Context, cli, srv *testQUICConn, onEvent func(e QUICEvent, src, dst *testQUICConn) bool) error {
 	a, b := cli, srv
 	for _, c := range []*testQUICConn{a, b} {
 		if !c.conn.conn.quic.started {
@@ -97,6 +97,9 @@ func runTestQUICConnection(ctx context.Context, cli, srv *testQUICConn, onHandle
 	idleCount := 0
 	for {
 		e := a.conn.NextEvent()
+		if onEvent != nil && onEvent(e, a, b) {
+			continue
+		}
 		switch e.Kind {
 		case QUICNoEvent:
 			idleCount++
@@ -210,6 +213,37 @@ func TestQUICSessionResumption(t *testing.T) {
 	}
 }
 
+func TestQUICFragmentaryData(t *testing.T) {
+	clientConfig := testConfig.Clone()
+	clientConfig.MinVersion = VersionTLS13
+	clientConfig.ClientSessionCache = NewLRUClientSessionCache(1)
+	clientConfig.ServerName = "example.go.dev"
+
+	serverConfig := testConfig.Clone()
+	serverConfig.MinVersion = VersionTLS13
+
+	cli := newTestQUICClient(t, clientConfig)
+	cli.conn.SetTransportParameters(nil)
+	srv := newTestQUICServer(t, serverConfig)
+	srv.conn.SetTransportParameters(nil)
+	onEvent := func(e QUICEvent, src, dst *testQUICConn) bool {
+		if e.Kind == QUICWriteData {
+			// Provide the data one byte at a time.
+			for i := range e.Data {
+				if err := dst.conn.HandleData(e.Level, e.Data[i:i+1]); err != nil {
+					t.Errorf("HandleData: %v", err)
+					break
+				}
+			}
+			return true
+		}
+		return false
+	}
+	if err := runTestQUICConnection(context.Background(), cli, srv, onEvent); err != nil {
+		t.Fatalf("error during first connection handshake: %v", err)
+	}
+}
+
 func TestQUICPostHandshakeClientAuthentication(t *testing.T) {
 	// RFC 9001, Section 4.4.
 	config := testConfig.Clone()
@@ -263,6 +297,28 @@ func TestQUICPostHandshakeKeyUpdate(t *testing.T) {
 	}
 }
 
+func TestQUICPostHandshakeMessageTooLarge(t *testing.T) {
+	config := testConfig.Clone()
+	config.MinVersion = VersionTLS13
+	cli := newTestQUICClient(t, config)
+	cli.conn.SetTransportParameters(nil)
+	srv := newTestQUICServer(t, config)
+	srv.conn.SetTransportParameters(nil)
+	if err := runTestQUICConnection(context.Background(), cli, srv, nil); err != nil {
+		t.Fatalf("error during connection handshake: %v", err)
+	}
+
+	size := maxHandshake + 1
+	if err := cli.conn.HandleData(QUICEncryptionLevelApplication, []byte{
+		byte(typeNewSessionTicket),
+		byte(size >> 16),
+		byte(size >> 8),
+		byte(size),
+	}); err == nil {
+		t.Fatalf("%v-byte post-handshake message: got no error, want one", size)
+	}
+}
+
 func TestQUICHandshakeError(t *testing.T) {
 	clientConfig := testConfig.Clone()
 	clientConfig.MinVersion = VersionTLS13
@@ -297,26 +353,22 @@ func TestQUICConnectionState(t *testing.T) {
 	cli.conn.SetTransportParameters(nil)
 	srv := newTestQUICServer(t, config)
 	srv.conn.SetTransportParameters(nil)
-	onHandleCryptoData := func() {
+	onEvent := func(e QUICEvent, src, dst *testQUICConn) bool {
 		cliCS := cli.conn.ConnectionState()
-		cliWantALPN := ""
 		if _, ok := cli.readSecret[QUICEncryptionLevelApplication]; ok {
-			cliWantALPN = "h3"
-		}
-		if want, got := cliCS.NegotiatedProtocol, cliWantALPN; want != got {
-			t.Errorf("cli.ConnectionState().NegotiatedProtocol = %q, want %q", want, got)
+			if want, got := cliCS.NegotiatedProtocol, "h3"; want != got {
+				t.Errorf("cli.ConnectionState().NegotiatedProtocol = %q, want %q", want, got)
+			}
 		}
-
 		srvCS := srv.conn.ConnectionState()
-		srvWantALPN := ""
 		if _, ok := srv.readSecret[QUICEncryptionLevelHandshake]; ok {
-			srvWantALPN = "h3"
-		}
-		if want, got := srvCS.NegotiatedProtocol, srvWantALPN; want != got {
-			t.Errorf("srv.ConnectionState().NegotiatedProtocol = %q, want %q", want, got)
+			if want, got := srvCS.NegotiatedProtocol, "h3"; want != got {
+				t.Errorf("srv.ConnectionState().NegotiatedProtocol = %q, want %q", want, got)
+			}
 		}
+		return false
 	}
-	if err := runTestQUICConnection(context.Background(), cli, srv, onHandleCryptoData); err != nil {
+	if err := runTestQUICConnection(context.Background(), cli, srv, onEvent); err != nil {
 		t.Fatalf("error during connection handshake: %v", err)
 	}
 }