From e3723a0ef162b38abd4cef7af8abf11ad056b128 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 27 Aug 2022 11:57:00 +0300
Subject: [PATCH] move the check for empty payload to the unpacker

---
 connection.go           |  7 -------
 connection_test.go      | 30 ------------------------------
 packet_unpacker.go      |  8 ++++++++
 packet_unpacker_test.go | 24 +++++++++++++++++++++++-
 4 files changed, 31 insertions(+), 38 deletions(-)

diff --git a/connection.go b/connection.go
index fdde3533..36331c8f 100644
--- a/connection.go
+++ b/connection.go
@@ -1123,13 +1123,6 @@ func (s *connection) handleUnpackedPacket(
 	rcvTime time.Time,
 	packetSize protocol.ByteCount, // only for logging
 ) error {
-	if len(packet.data) == 0 {
-		return &qerr.TransportError{
-			ErrorCode:    qerr.ProtocolViolation,
-			ErrorMessage: "empty packet",
-		}
-	}
-
 	if !s.receivedFirstPacket {
 		s.receivedFirstPacket = true
 		if !s.versionNegotiated && s.tracer != nil {
diff --git a/connection_test.go b/connection_test.go
index bbe49fa0..47e81c0d 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -1013,36 +1013,6 @@ var _ = Describe("Connection", func() {
 			Eventually(conn.Context().Done()).Should(BeClosed())
 		})
 
-		It("rejects packets with empty payload", func() {
-			unpacker.EXPECT().Unpack(gomock.Any(), gomock.Any(), gomock.Any()).Return(&unpackedPacket{
-				hdr:             &wire.ExtendedHeader{},
-				data:            []byte{}, // no payload
-				encryptionLevel: protocol.Encryption1RTT,
-			}, nil)
-			streamManager.EXPECT().CloseWithError(gomock.Any())
-			cryptoSetup.EXPECT().Close()
-			packer.EXPECT().PackConnectionClose(gomock.Any()).Return(&coalescedPacket{buffer: getPacketBuffer()}, nil)
-			done := make(chan struct{})
-			go func() {
-				defer GinkgoRecover()
-				cryptoSetup.EXPECT().RunHandshake().MaxTimes(1)
-				Expect(conn.run()).To(MatchError(&qerr.TransportError{
-					ErrorCode:    qerr.ProtocolViolation,
-					ErrorMessage: "empty packet",
-				}))
-				close(done)
-			}()
-			expectReplaceWithClosed()
-			mconn.EXPECT().Write(gomock.Any())
-			tracer.EXPECT().ClosedConnection(gomock.Any())
-			tracer.EXPECT().Close()
-			conn.handlePacket(getPacket(&wire.ExtendedHeader{
-				Header:          wire.Header{DestConnectionID: srcConnID},
-				PacketNumberLen: protocol.PacketNumberLen1,
-			}, nil))
-			Eventually(done).Should(BeClosed())
-		})
-
 		It("ignores packets with a different source connection ID", func() {
 			hdr1 := &wire.ExtendedHeader{
 				Header: wire.Header{
diff --git a/packet_unpacker.go b/packet_unpacker.go
index 59f6966f..a74b3047 100644
--- a/packet_unpacker.go
+++ b/packet_unpacker.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/lucas-clemente/quic-go/internal/handshake"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
+	"github.com/lucas-clemente/quic-go/internal/qerr"
 	"github.com/lucas-clemente/quic-go/internal/wire"
 )
 
@@ -102,6 +103,13 @@ func (u *packetUnpacker) Unpack(hdr *wire.Header, rcvTime time.Time, data []byte
 		}
 	}
 
+	if len(decrypted) == 0 {
+		return nil, &qerr.TransportError{
+			ErrorCode:    qerr.ProtocolViolation,
+			ErrorMessage: "empty packet",
+		}
+	}
+
 	return &unpackedPacket{
 		hdr:             extHdr,
 		encryptionLevel: encLevel,
diff --git a/packet_unpacker_test.go b/packet_unpacker_test.go
index fd8b50b4..239311be 100644
--- a/packet_unpacker_test.go
+++ b/packet_unpacker_test.go
@@ -155,7 +155,7 @@ var _ = Describe("Packet Unpacker", func() {
 		Expect(packet.data).To(Equal([]byte("decrypted")))
 	})
 
-	It("returns the error when getting the sealer fails", func() {
+	It("returns the error when getting the opener fails", func() {
 		extHdr := &wire.ExtendedHeader{
 			Header:          wire.Header{DestConnectionID: connID},
 			PacketNumber:    0x1337,
@@ -167,6 +167,28 @@ var _ = Describe("Packet Unpacker", func() {
 		Expect(err).To(MatchError(handshake.ErrKeysNotYetAvailable))
 	})
 
+	It("errors on empty packets", func() {
+		extHdr := &wire.ExtendedHeader{
+			Header:          wire.Header{DestConnectionID: connID},
+			KeyPhase:        protocol.KeyPhaseOne,
+			PacketNumberLen: protocol.PacketNumberLen4,
+		}
+		hdr, hdrRaw := getHeader(extHdr)
+		opener := mocks.NewMockShortHeaderOpener(mockCtrl)
+		now := time.Now()
+		gomock.InOrder(
+			cs.EXPECT().Get1RTTOpener().Return(opener, nil),
+			opener.EXPECT().DecryptHeader(gomock.Any(), gomock.Any(), gomock.Any()),
+			opener.EXPECT().DecodePacketNumber(gomock.Any(), gomock.Any()).Return(protocol.PacketNumber(321)),
+			opener.EXPECT().Open(gomock.Any(), payload, now, protocol.PacketNumber(321), protocol.KeyPhaseOne, hdrRaw).Return([]byte(""), nil),
+		)
+		_, err := unpacker.Unpack(hdr, now, append(hdrRaw, payload...))
+		Expect(err).To(MatchError(&qerr.TransportError{
+			ErrorCode:    qerr.ProtocolViolation,
+			ErrorMessage: "empty packet",
+		}))
+	})
+
 	It("returns the error when unpacking fails", func() {
 		extHdr := &wire.ExtendedHeader{
 			Header: wire.Header{