TLS messages are now sent in the right encryption level instead of pi…

…ping everything in Handshake packets

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 {

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]
 }

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)

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