move packet number decoding to the unpacker

packet_unpacker.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/lucas-clemente/quic-go/internal/qerr"
+	"github.com/lucas-clemente/quic-go/internal/utils"
 	"github.com/lucas-clemente/quic-go/internal/wire"
 )
 
@@ -22,7 +23,10 @@ type quicAEAD interface {
 
 // The packetUnpacker unpacks QUIC packets.
 type packetUnpacker struct {
-	aead    quicAEAD
+	aead quicAEAD
+
+	largestRcvdPacketNumber protocol.PacketNumber
+
 	version protocol.VersionNumber
 }
 
@@ -36,6 +40,12 @@ func newPacketUnpacker(aead quicAEAD, version protocol.VersionNumber) unpacker {
 }
 
 func (u *packetUnpacker) Unpack(headerBinary []byte, hdr *wire.ExtendedHeader, data []byte) (*unpackedPacket, error) {
+	hdr.PacketNumber = protocol.DecodePacketNumber(
+		hdr.PacketNumberLen,
+		u.largestRcvdPacketNumber,
+		hdr.PacketNumber,
+	)
+
 	buf := *getPacketBuffer()
 	buf = buf[:0]
 	defer putPacketBuffer(&buf)
@@ -61,6 +71,9 @@ func (u *packetUnpacker) Unpack(headerBinary []byte, hdr *wire.ExtendedHeader, d
 		return nil, qerr.Error(qerr.DecryptionFailure, err.Error())
 	}
 
+	// Only do this after decrypting, so we are sure the packet is not attacker-controlled
+	u.largestRcvdPacketNumber = utils.MaxPacketNumber(u.largestRcvdPacketNumber, hdr.PacketNumber)
+
 	fs, err := u.parseFrames(decrypted)
 	if err != nil {
 		return nil, err

packet_unpacker_test.go

@@ -2,6 +2,7 @@ package quic
 
 import (
 	"bytes"
+	"errors"
 
 	"github.com/golang/mock/gomock"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
@@ -54,6 +55,33 @@ var _ = Describe("Packet Unpacker", func() {
 		Expect(packet.encryptionLevel).To(Equal(protocol.EncryptionHandshake))
 	})
 
+	It("returns the error when unpacking fails", func() {
+		hdr.IsLongHeader = true
+		hdr.Type = protocol.PacketTypeHandshake
+		aead.EXPECT().OpenHandshake(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("test err"))
+		_, err := unpacker.Unpack(hdr.Raw, hdr, nil)
+		Expect(err).To(MatchError(qerr.Error(qerr.DecryptionFailure, "test err")))
+	})
+
+	It("decodes the packet number", func() {
+		firstHdr := &wire.ExtendedHeader{
+			PacketNumber:    0x1337,
+			PacketNumberLen: 2,
+		}
+		aead.EXPECT().Open1RTT(gomock.Any(), gomock.Any(), firstHdr.PacketNumber, gomock.Any()).Return([]byte{0}, nil)
+		_, err := unpacker.Unpack(firstHdr.Raw, firstHdr, nil)
+		Expect(err).ToNot(HaveOccurred())
+		// the real packet number is 0x1338, but only the last byte is sent
+		secondHdr := &wire.ExtendedHeader{
+			PacketNumber:    0x38,
+			PacketNumberLen: 1,
+		}
+		// expect the call with the decoded packet number
+		aead.EXPECT().Open1RTT(gomock.Any(), gomock.Any(), protocol.PacketNumber(0x1338), gomock.Any()).Return([]byte{0}, nil)
+		_, err = unpacker.Unpack(secondHdr.Raw, secondHdr, nil)
+		Expect(err).ToNot(HaveOccurred())
+	})
+
 	It("unpacks the frames", func() {
 		buf := &bytes.Buffer{}
 		(&wire.PingFrame{}).Write(buf, protocol.VersionWhatever)

session.go

@@ -112,9 +112,8 @@ type session struct {
 	handshakeCompleteChan chan struct{} // is closed when the handshake completes
 	handshakeComplete     bool
 
-	receivedFirstPacket              bool // since packet numbers start at 0, we can't use largestRcvdPacketNumber != 0 for this
+	receivedFirstPacket              bool
 	receivedFirstForwardSecurePacket bool
-	largestRcvdPacketNumber          protocol.PacketNumber // used to calculate the next packet number
 
 	sessionCreationTime     time.Time
 	lastNetworkActivityTime time.Time
@@ -504,13 +503,6 @@ func (s *session) handlePacketImpl(p *receivedPacket) error {
 		// TODO(#1312): implement parsing of compound packets
 	}
 
-	// Calculate packet number
-	hdr.PacketNumber = protocol.DecodePacketNumber(
-		hdr.PacketNumberLen,
-		s.largestRcvdPacketNumber,
-		hdr.PacketNumber,
-	)
-
 	packet, err := s.unpacker.Unpack(hdr.Raw, hdr, data)
 	if s.logger.Debug() {
 		if err != nil {
@@ -545,9 +537,6 @@ func (s *session) handlePacketImpl(p *receivedPacket) error {
 		}
 	}
 
-	// Only do this after decrypting, so we are sure the packet is not attacker-controlled
-	s.largestRcvdPacketNumber = utils.MaxPacketNumber(s.largestRcvdPacketNumber, hdr.PacketNumber)
-
 	// If this is a Retry packet, there's no need to send an ACK.
 	// The session will be closed and recreated as soon as the crypto setup processed the HRR.
 	if hdr.Type != protocol.PacketTypeRetry {

session_test.go

@@ -461,18 +461,6 @@ var _ = Describe("Session", func() {
 			return buf.Bytes()
 		}
 
-		It("sets the largestRcvdPacketNumber", func() {
-			hdr := &wire.ExtendedHeader{
-				PacketNumber:    5,
-				PacketNumberLen: protocol.PacketNumberLen4,
-			}
-			hdrRaw := getData(hdr)
-			data := append(hdrRaw, []byte("foobar")...)
-			unpacker.EXPECT().Unpack(hdrRaw, gomock.Any(), []byte("foobar")).Return(&unpackedPacket{}, nil)
-			Expect(sess.handlePacketImpl(&receivedPacket{hdr: &hdr.Header, data: data})).To(Succeed())
-			Expect(sess.largestRcvdPacketNumber).To(Equal(protocol.PacketNumber(5)))
-		})
-
 		It("informs the ReceivedPacketHandler", func() {
 			hdr := &wire.ExtendedHeader{
 				Raw:             []byte("raw header"),