From 55d74be4a3517c1a796bac47228bec58d53e1800 Mon Sep 17 00:00:00 2001
From: Maxime Piraux <maxime.piraux@uclouvain.be>
Date: Mon, 17 Sep 2018 14:03:51 +0200
Subject: [PATCH] TLS messages are now sent in the right encryption level
 instead of piping everything in Handshake packets

---
 agents/tls_agent.go | 13 ++++---------
 common.go           |  7 +++++++
 connection.go       |  3 ++-
 crypto.go           |  7 +++++++
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/agents/tls_agent.go b/agents/tls_agent.go
index 0fdad00..7fb61e1 100644
--- a/agents/tls_agent.go
+++ b/agents/tls_agent.go
@@ -55,7 +55,6 @@ func (a *TLSAgent) Run(conn *Connection) {
 				if _, ok := packet.(Framer); !ok {
 					break
 				}
-				cryptoStream := conn.CryptoStreams.Get(packet.PNSpace())
 				cryptoChan := cryptoChans[packet.PNSpace()]
 
 				var handshakeData []byte
@@ -75,7 +74,7 @@ func (a *TLSAgent) Run(conn *Connection) {
 				switch packet.(type) {
 				case Framer:
 					if len(handshakeData) > 0 {
-						responseData, notCompleted, err := conn.Tls.HandleMessage(handshakeData, PNSpaceToEpoch[packet.PNSpace()])
+						tlsOutput, notCompleted, err := conn.Tls.HandleMessage(handshakeData, PNSpaceToEpoch[packet.PNSpace()])
 
 						if err != nil {
 							a.Logger.Printf("TLS error occured: %s\n", err.Error())
@@ -97,14 +96,10 @@ func (a *TLSAgent) Run(conn *Connection) {
 							}
 						}
 
-						if len(responseData) > 0 && !a.DisableFrameSending {
-							var responseEncryptionLevel EncryptionLevel
-							if packet.EncryptionLevel() == EncryptionLevelInitial {
-								responseEncryptionLevel = EncryptionLevelHandshake
-							} else {
-								responseEncryptionLevel = packet.EncryptionLevel()
+						if len(tlsOutput) > 0 && !a.DisableFrameSending {
+							for _, m := range tlsOutput {
+								conn.FrameQueue.Submit(QueuedFrame{NewCryptoFrame(conn.CryptoStreams.Get(EpochToPNSpace[m.Epoch]), m.Data), EpochToEncryptionLevel[m.Epoch]})
 							}
-							conn.FrameQueue.Submit(QueuedFrame{NewCryptoFrame(cryptoStream, responseData), responseEncryptionLevel})
 						}
 
 						if !notCompleted && conn.CryptoStates[EncryptionLevel1RTT] == nil {
diff --git a/common.go b/common.go
index 5d687a0..fca3ac8 100644
--- a/common.go
+++ b/common.go
@@ -179,6 +179,13 @@ var PNSpaceToEpoch = map[PNSpace]pigotls.Epoch{
 	PNSpaceAppData: pigotls.Epoch1RTT,
 }
 
+var EpochToPNSpace = map[pigotls.Epoch]PNSpace {
+	pigotls.EpochInitial: PNSpaceInitial,
+	pigotls.EpochHandshake: PNSpaceHandshake,
+	pigotls.Epoch0RTT: PNSpaceAppData,
+	pigotls.Epoch1RTT: PNSpaceAppData,
+}
+
 func (pns PNSpace) String() string {
 	return PNSpaceToString[pns]
 }
diff --git a/connection.go b/connection.go
index 5897fdb..5a6befc 100644
--- a/connection.go
+++ b/connection.go
@@ -102,11 +102,12 @@ func (c *Connection) GetInitialPacket() *InitialPacket {
 	}
 	c.Tls.SetQUICTransportParameters(extensionData)
 
-	clientHello, notComplete, err := c.Tls.HandleMessage(nil, pigotls.EpochInitial)
+	tlsOutput, notComplete, err := c.Tls.HandleMessage(nil, pigotls.EpochInitial)
 	if err != nil || !notComplete {
 		println(err.Error())
 		return nil
 	}
+	clientHello := tlsOutput[0].Data
 	c.ClientRandom = make([]byte, 32, 32)
 	copy(c.ClientRandom, clientHello[11:11+32])
 	cryptoFrame := NewCryptoFrame(c.CryptoStreams.Get(PNSpaceInitial), clientHello)
diff --git a/crypto.go b/crypto.go
index 533e76f..5fd379a 100644
--- a/crypto.go
+++ b/crypto.go
@@ -49,6 +49,13 @@ var packetTypeToEncryptionLevel = map[PacketType]EncryptionLevel{
 	ShortHeaderPacket: EncryptionLevel1RTT,
 }
 
+var EpochToEncryptionLevel = map[pigotls.Epoch]EncryptionLevel {
+	pigotls.EpochInitial: EncryptionLevelInitial,
+	pigotls.Epoch0RTT: EncryptionLevel0RTT,
+	pigotls.EpochHandshake: EncryptionLevelHandshake,
+	pigotls.Epoch1RTT: EncryptionLevel1RTT,
+}
+
 type DirectionalEncryptionLevel struct {
 	EncryptionLevel
 	Read bool