From fc37cdc5c5d894e1fdb49800de21b5a4a57ccd72 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 4 Aug 2019 13:49:36 +0700
Subject: [PATCH] implement unpacking of 0-RTT packets

---
 mock_sealing_manager_test.go | 15 +++++++++++++++
 packet_unpacker.go           | 10 ++++++++++
 packet_unpacker_test.go      | 23 +++++++++++++++++++++++
 server.go                    |  1 +
 session.go                   |  4 ++--
 5 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/mock_sealing_manager_test.go b/mock_sealing_manager_test.go
index f8c661e8..238fff32 100644
--- a/mock_sealing_manager_test.go
+++ b/mock_sealing_manager_test.go
@@ -34,6 +34,21 @@ func (m *MockSealingManager) EXPECT() *MockSealingManagerMockRecorder {
 	return m.recorder
 }
 
+// Get0RTTSealer mocks base method
+func (m *MockSealingManager) Get0RTTSealer() (handshake.LongHeaderSealer, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Get0RTTSealer")
+	ret0, _ := ret[0].(handshake.LongHeaderSealer)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Get0RTTSealer indicates an expected call of Get0RTTSealer
+func (mr *MockSealingManagerMockRecorder) Get0RTTSealer() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get0RTTSealer", reflect.TypeOf((*MockSealingManager)(nil).Get0RTTSealer))
+}
+
 // Get1RTTSealer mocks base method
 func (m *MockSealingManager) Get1RTTSealer() (handshake.ShortHeaderSealer, error) {
 	m.ctrl.T.Helper()
diff --git a/packet_unpacker.go b/packet_unpacker.go
index 3cf943e5..27eb52ec 100644
--- a/packet_unpacker.go
+++ b/packet_unpacker.go
@@ -65,6 +65,16 @@ func (u *packetUnpacker) Unpack(hdr *wire.Header, rcvTime time.Time, data []byte
 		if err != nil {
 			return nil, err
 		}
+	case protocol.PacketType0RTT:
+		encLevel = protocol.Encryption0RTT
+		opener, err := u.cs.Get0RTTOpener()
+		if err != nil {
+			return nil, err
+		}
+		extHdr, decrypted, err = u.unpackLongHeaderPacket(opener, hdr, data)
+		if err != nil {
+			return nil, err
+		}
 	default:
 		if hdr.IsLongHeader {
 			return nil, fmt.Errorf("unknown packet type: %s", hdr.Type)
diff --git a/packet_unpacker_test.go b/packet_unpacker_test.go
index 9b1eb260..1f7fe241 100644
--- a/packet_unpacker_test.go
+++ b/packet_unpacker_test.go
@@ -78,6 +78,29 @@ var _ = Describe("Packet Unpacker", func() {
 		Expect(packet.data).To(Equal([]byte("decrypted")))
 	})
 
+	It("opens 0-RTT packets", func() {
+		extHdr := &wire.ExtendedHeader{
+			Header: wire.Header{
+				IsLongHeader:     true,
+				Type:             protocol.PacketType0RTT,
+				Length:           3 + 6, // packet number len + payload
+				DestConnectionID: connID,
+				Version:          version,
+			},
+			PacketNumber:    2,
+			PacketNumberLen: 3,
+		}
+		hdr, hdrRaw := getHeader(extHdr)
+		opener := mocks.NewMockLongHeaderOpener(mockCtrl)
+		cs.EXPECT().Get0RTTOpener().Return(opener, nil)
+		opener.EXPECT().DecryptHeader(gomock.Any(), gomock.Any(), gomock.Any())
+		opener.EXPECT().Open(gomock.Any(), payload, extHdr.PacketNumber, hdrRaw).Return([]byte("decrypted"), nil)
+		packet, err := unpacker.Unpack(hdr, time.Now(), append(hdrRaw, payload...))
+		Expect(err).ToNot(HaveOccurred())
+		Expect(packet.encryptionLevel).To(Equal(protocol.Encryption0RTT))
+		Expect(packet.data).To(Equal([]byte("decrypted")))
+	})
+
 	It("returns the error when getting the sealer fails", func() {
 		extHdr := &wire.ExtendedHeader{
 			Header:          wire.Header{DestConnectionID: connID},
diff --git a/server.go b/server.go
index 9dea3694..7638d037 100644
--- a/server.go
+++ b/server.go
@@ -361,6 +361,7 @@ func (s *baseServer) handlePacketImpl(p *receivedPacket) bool /* was the packet
 		// Drop long header packets.
 		// There's litte point in sending a Stateless Reset, since the client
 		// might not have received the token yet.
+		s.logger.Debugf("Dropping long header packet of type %s (%d bytes)", hdr.Type, len(p.data))
 		return false
 	}
 
diff --git a/session.go b/session.go
index fd3cf3d5..d65bf9e0 100644
--- a/session.go
+++ b/session.go
@@ -683,8 +683,8 @@ func (s *session) handleSinglePacket(p *receivedPacket, hdr *wire.Header) bool /
 		s.logger.Debugf("Dropping %s packet with unexpected source connection ID: %s (expected %s)", hdr.PacketType(), hdr.SrcConnectionID, s.handshakeDestConnID)
 		return false
 	}
-	// drop 0-RTT packets
-	if hdr.Type == protocol.PacketType0RTT {
+	// drop 0-RTT packets, if we are a client
+	if s.perspective == protocol.PerspectiveClient && hdr.Type == protocol.PacketType0RTT {
 		return false
 	}