From f88cbd801b2072b15eba3492c01d557b57902552 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 24 Feb 2017 13:16:30 +0700
Subject: [PATCH] store the encryption level in the unpacked packet

---
 packet_unpacker.go      |  5 +++--
 packet_unpacker_test.go | 17 +++++++++++++++--
 unpacked_packet.go      |  8 ++++++--
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/packet_unpacker.go b/packet_unpacker.go
index 44eee48a..a826fe7a 100644
--- a/packet_unpacker.go
+++ b/packet_unpacker.go
@@ -23,7 +23,7 @@ type packetUnpacker struct {
 func (u *packetUnpacker) Unpack(publicHeaderBinary []byte, hdr *PublicHeader, data []byte) (*unpackedPacket, error) {
 	buf := getPacketBuffer()
 	defer putPacketBuffer(buf)
-	decrypted, _, err := u.aead.Open(buf, data, hdr.PacketNumber, publicHeaderBinary)
+	decrypted, encryptionLevel, err := u.aead.Open(buf, data, hdr.PacketNumber, publicHeaderBinary)
 	if err != nil {
 		// Wrap err in quicError so that public reset is sent by session
 		return nil, qerr.Error(qerr.DecryptionFailure, err.Error())
@@ -104,6 +104,7 @@ func (u *packetUnpacker) Unpack(publicHeaderBinary []byte, hdr *PublicHeader, da
 	}
 
 	return &unpackedPacket{
-		frames: fs,
+		encryptionLevel: encryptionLevel,
+		frames:          fs,
 	}, nil
 }
diff --git a/packet_unpacker_test.go b/packet_unpacker_test.go
index 2c7609fa..799fcff9 100644
--- a/packet_unpacker_test.go
+++ b/packet_unpacker_test.go
@@ -12,11 +12,13 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-type mockAEAD struct{}
+type mockAEAD struct {
+	encLevelOpen protocol.EncryptionLevel
+}
 
 func (m *mockAEAD) Open(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, protocol.EncryptionLevel, error) {
 	res, err := (&crypto.NullAEAD{}).Open(dst, src, packetNumber, associatedData)
-	return res, protocol.EncryptionUnspecified, err
+	return res, m.encLevelOpen, err
 }
 func (m *mockAEAD) Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) []byte {
 	return (&crypto.NullAEAD{}).Seal(dst, src, packetNumber, associatedData)
@@ -58,6 +60,17 @@ var _ = Describe("Packet unpacker", func() {
 		Expect(packet.frames).To(Equal([]frames.Frame{f}))
 	})
 
+	It("saves the encryption level", func() {
+		f := &frames.ConnectionCloseFrame{ReasonPhrase: "foo"}
+		err := f.Write(buf, 0)
+		Expect(err).ToNot(HaveOccurred())
+		setData(buf.Bytes())
+		unpacker.aead.(*mockAEAD).encLevelOpen = protocol.EncryptionSecure
+		packet, err := unpacker.Unpack(hdrBin, hdr, data)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(packet.encryptionLevel).To(Equal(protocol.EncryptionSecure))
+	})
+
 	It("unpacks STREAM frames", func() {
 		f := &frames.StreamFrame{
 			StreamID: 1,
diff --git a/unpacked_packet.go b/unpacked_packet.go
index 80792047..0636b8f1 100644
--- a/unpacked_packet.go
+++ b/unpacked_packet.go
@@ -1,9 +1,13 @@
 package quic
 
-import "github.com/lucas-clemente/quic-go/frames"
+import (
+	"github.com/lucas-clemente/quic-go/frames"
+	"github.com/lucas-clemente/quic-go/protocol"
+)
 
 type unpackedPacket struct {
-	frames []frames.Frame
+	encryptionLevel protocol.EncryptionLevel
+	frames          []frames.Frame
 }
 
 func (u *unpackedPacket) IsRetransmittable() bool {