From 1bd1a9e0e5cb2b1fb2e1f0084965c590539ec810 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 24 Feb 2017 14:08:21 +0700
Subject: [PATCH] store the encryption level that a packet was sent with

---
 packet_packer.go      | 16 +++++++++-------
 packet_packer_test.go | 15 ++++++++++++++-
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/packet_packer.go b/packet_packer.go
index a44490a2..43499379 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -11,9 +11,10 @@ import (
 )
 
 type packedPacket struct {
-	number protocol.PacketNumber
-	raw    []byte
-	frames []frames.Frame
+	number          protocol.PacketNumber
+	raw             []byte
+	frames          []frames.Frame
+	encryptionLevel protocol.EncryptionLevel
 }
 
 type packetPacker struct {
@@ -138,7 +139,7 @@ func (p *packetPacker) packPacket(stopWaitingFrame *frames.StopWaitingFrame, lea
 	}
 
 	raw = raw[0:buffer.Len()]
-	p.cryptoSetup.Seal(raw[payloadStartIndex:payloadStartIndex], raw[payloadStartIndex:], currentPacketNumber, raw[:payloadStartIndex])
+	_, encryptionLevel := p.cryptoSetup.Seal(raw[payloadStartIndex:payloadStartIndex], raw[payloadStartIndex:], currentPacketNumber, raw[:payloadStartIndex])
 	raw = raw[0 : buffer.Len()+12]
 
 	num := p.packetNumberGenerator.Pop()
@@ -147,9 +148,10 @@ func (p *packetPacker) packPacket(stopWaitingFrame *frames.StopWaitingFrame, lea
 	}
 
 	return &packedPacket{
-		number: currentPacketNumber,
-		raw:    raw,
-		frames: payloadFrames,
+		number:          currentPacketNumber,
+		raw:             raw,
+		frames:          payloadFrames,
+		encryptionLevel: encryptionLevel,
 	}, nil
 }
 
diff --git a/packet_packer_test.go b/packet_packer_test.go
index dbc57d94..f2c0f836 100644
--- a/packet_packer_test.go
+++ b/packet_packer_test.go
@@ -12,6 +12,7 @@ import (
 type mockCryptoSetup struct {
 	diversificationNonce []byte
 	handshakeComplete    bool
+	encLevelSeal         protocol.EncryptionLevel
 }
 
 func (m *mockCryptoSetup) HandleCryptoStream() error { return nil }
@@ -20,7 +21,7 @@ func (m *mockCryptoSetup) Open(dst, src []byte, packetNumber protocol.PacketNumb
 	return nil, protocol.EncryptionUnspecified, nil
 }
 func (m *mockCryptoSetup) Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, protocol.EncryptionLevel) {
-	return append(src, bytes.Repeat([]byte{0}, 12)...), protocol.EncryptionUnspecified
+	return append(src, bytes.Repeat([]byte{0}, 12)...), m.encLevelSeal
 }
 func (m *mockCryptoSetup) LockForSealing()                      {}
 func (m *mockCryptoSetup) UnlockForSealing()                    {}
@@ -76,6 +77,18 @@ var _ = Describe("Packet packer", func() {
 		Expect(p.raw).To(ContainSubstring(string(b.Bytes())))
 	})
 
+	It("stores the encryption level a packet was sealed with", func() {
+		packer.cryptoSetup.(*mockCryptoSetup).encLevelSeal = protocol.EncryptionSecure
+		f := &frames.StreamFrame{
+			StreamID: 5,
+			Data:     []byte("foobar"),
+		}
+		streamFramer.AddFrameForRetransmission(f)
+		p, err := packer.PackPacket(nil, []frames.Frame{}, 0)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(p.encryptionLevel).To(Equal(protocol.EncryptionSecure))
+	})
+
 	It("includes a diversification nonce, when acting as a server", func() {
 		nonce := bytes.Repeat([]byte{'e'}, 32)
 		packer.cryptoSetup.(*mockCryptoSetup).diversificationNonce = nonce