From 8bcad17297024c160a994e154c2b28cbe5097f88 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 12 Nov 2016 16:42:03 +0700
Subject: [PATCH] create secureAEAD after receiving all necessary values

---
 handshake/crypto_setup_client.go      | 48 +++++++++++++++++++++++++--
 handshake/crypto_setup_client_test.go | 41 +++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/handshake/crypto_setup_client.go b/handshake/crypto_setup_client.go
index 853df8dd..1240d2e7 100644
--- a/handshake/crypto_setup_client.go
+++ b/handshake/crypto_setup_client.go
@@ -28,6 +28,9 @@ type cryptoSetupClient struct {
 	diversificationNonce []byte
 	lastSentCHLO         []byte
 	certManager          crypto.CertManager
+
+	keyDerivation KeyDerivationFunction
+	secureAEAD    crypto.AEAD
 }
 
 var _ crypto.AEAD = &cryptoSetupClient{}
@@ -50,12 +53,19 @@ func NewCryptoSetupClient(
 		version:      version,
 		cryptoStream: cryptoStream,
 		certManager:  crypto.NewCertManager(),
+
+		keyDerivation: crypto.DeriveKeysAESGCM,
 	}, nil
 }
 
 func (h *cryptoSetupClient) HandleCryptoStream() error {
 	for {
-		err := h.sendCHLO()
+		err := h.maybeUpgradeCrypto()
+		if err != nil {
+			return err
+		}
+
+		err = h.sendCHLO()
 		if err != nil {
 			return err
 		}
@@ -125,10 +135,20 @@ func (h *cryptoSetupClient) handleREJMessage(cryptoData map[Tag][]byte) error {
 }
 
 func (h *cryptoSetupClient) Open(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, error) {
+	if h.secureAEAD != nil {
+		data, err := h.secureAEAD.Open(dst, src, packetNumber, associatedData)
+		if err == nil {
+			return data, nil
+		}
+		return nil, err
+	}
 	return (&crypto.NullAEAD{}).Open(dst, src, packetNumber, associatedData)
 }
 
 func (h *cryptoSetupClient) Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) []byte {
+	if h.secureAEAD != nil {
+		return h.secureAEAD.Seal(dst, src, packetNumber, associatedData)
+	}
 	return (&crypto.NullAEAD{}).Seal(dst, src, packetNumber, associatedData)
 }
 
@@ -139,7 +159,7 @@ func (h *cryptoSetupClient) DiversificationNonce() []byte {
 func (h *cryptoSetupClient) SetDiversificationNonce(data []byte) error {
 	if len(h.diversificationNonce) == 0 {
 		h.diversificationNonce = data
-		return nil
+		return h.maybeUpgradeCrypto()
 	}
 	if !bytes.Equal(h.diversificationNonce, data) {
 		return errConflictingDiversificationNonces
@@ -206,6 +226,30 @@ func (h *cryptoSetupClient) getTags() map[Tag][]byte {
 	return tags
 }
 
+func (h *cryptoSetupClient) maybeUpgradeCrypto() error {
+	leafCert := h.certManager.GetLeafCert()
+
+	if h.secureAEAD == nil && (h.serverConfig != nil && len(h.serverConfig.sharedSecret) > 0 && len(h.nonc) > 0 && len(leafCert) > 0 && len(h.diversificationNonce) > 0 && len(h.lastSentCHLO) > 0) {
+		var err error
+		h.secureAEAD, err = h.keyDerivation(
+			false,
+			h.serverConfig.sharedSecret,
+			h.nonc,
+			h.connID,
+			h.lastSentCHLO,
+			h.serverConfig.Get(),
+			leafCert,
+			h.diversificationNonce,
+			protocol.PerspectiveClient,
+		)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (h *cryptoSetupClient) generateClientNonce() error {
 	if len(h.nonc) > 0 {
 		return errClientNonceAlreadyExists
diff --git a/handshake/crypto_setup_client_test.go b/handshake/crypto_setup_client_test.go
index 3562a1a7..20a2bb71 100644
--- a/handshake/crypto_setup_client_test.go
+++ b/handshake/crypto_setup_client_test.go
@@ -207,6 +207,47 @@ var _ = Describe("Crypto setup", func() {
 		})
 	})
 
+	Context("escalating crypto", func() {
+		// sets all values necessary for escalating to secureAEAD
+		BeforeEach(func() {
+			kex, err := crypto.NewCurve25519KEX()
+			Expect(err).ToNot(HaveOccurred())
+			cs.serverConfig = &serverConfigClient{
+				kex:          kex,
+				obit:         []byte("obit"),
+				sharedSecret: []byte("sharedSecret"),
+				raw:          []byte("rawserverconfig"),
+			}
+			cs.lastSentCHLO = []byte("lastSentCHLO")
+			cs.nonc = []byte("nonc")
+			cs.diversificationNonce = []byte("divnonce")
+			certManager.leafCert = []byte("leafCert")
+		})
+
+		It("creates a secureAEAD once it has all necessary values", func() {
+			err := cs.maybeUpgradeCrypto()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(cs.secureAEAD).ToNot(BeNil())
+		})
+
+		It("tries to escalate before reading a handshake message", func() {
+			Expect(cs.secureAEAD).To(BeNil())
+			err := cs.HandleCryptoStream()
+			// this will throw a qerr.HandshakeFailed due to an EOF in WriteHandshakeMessage
+			// this is because the mockStream doesn't block if there's no data to read
+			Expect(err).To(MatchError(qerr.HandshakeFailed))
+			Expect(cs.secureAEAD).ToNot(BeNil())
+		})
+
+		It("tries to escalate the crypto after receiving a diversification nonce", func() {
+			cs.diversificationNonce = nil
+			Expect(cs.secureAEAD).To(BeNil())
+			err := cs.SetDiversificationNonce([]byte("div"))
+			Expect(err).ToNot(HaveOccurred())
+			Expect(cs.secureAEAD).ToNot(BeNil())
+		})
+	})
+
 	Context("Diversification Nonces", func() {
 		It("sets a diversification nonce", func() {
 			nonce := []byte("foobar")