From b5c7bcda187e572190db525ee33a3176739c952e Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Sun, 17 Apr 2016 00:31:30 +0200
Subject: [PATCH] extract SHLO into separate method and test it

---
 handshake/crypto_setup.go      | 63 ++++++++++++++++++----------------
 handshake/crypto_setup_test.go | 18 ++++++++--
 2 files changed, 48 insertions(+), 33 deletions(-)

diff --git a/handshake/crypto_setup.go b/handshake/crypto_setup.go
index be2c7537..da3dd10e 100644
--- a/handshake/crypto_setup.go
+++ b/handshake/crypto_setup.go
@@ -83,39 +83,11 @@ func (h *CryptoSetup) HandleCryptoMessage(data []byte) ([]byte, error) {
 	}
 
 	if scid, ok := cryptoData[TagSCID]; ok && bytes.Equal(h.scfg.ID, scid) {
-		// We have a CHLO matching our server config, we can continue with the 0-RTT handshake
-		var sharedSecret []byte
-		sharedSecret, err = h.scfg.kex.CalculateSharedKey(cryptoData[TagPUBS])
-		if err != nil {
-			return nil, err
-		}
-		var nonce bytes.Buffer
-		nonce.Write(cryptoData[TagNONC])
-		nonce.Write(h.nonce)
-
-		h.secureAEAD, err = crypto.DeriveKeysChacha20(false, sharedSecret, nonce.Bytes(), h.connID, data, h.scfg.Get(), h.scfg.signer.GetCertUncompressed())
-		if err != nil {
-			return nil, err
-		}
-		// TODO: Use new curve
-		h.forwardSecureAEAD, err = crypto.DeriveKeysChacha20(true, sharedSecret, nonce.Bytes(), h.connID, data, h.scfg.Get(), h.scfg.signer.GetCertUncompressed())
-		if err != nil {
-			return nil, err
-		}
-
-		var reply bytes.Buffer
-		WriteHandshakeMessage(&reply, TagSHLO, map[Tag][]byte{
-			TagPUBS: h.scfg.kex.PublicKey(),
-			TagSNO:  h.nonce,
-			TagVER:  protocol.SupportedVersionsAsTags,
-			TagICSL: []byte{0x1e, 0x00, 0x00, 0x00}, //30
-			TagMSPC: []byte{0x64, 0x00, 0x00, 0x00}, //100
-		})
-		return reply.Bytes(), nil
+		// We have a CHLO with a proper server config ID, do a 0-RTT handshake
+		return h.handleCHLO(data, cryptoData)
 	}
 
 	// We have an inacholate or non-matching CHLO, we now send a rejection
-
 	return h.handleInchoateCHLO(data)
 }
 
@@ -139,3 +111,34 @@ func (h *CryptoSetup) handleInchoateCHLO(data []byte) ([]byte, error) {
 	})
 	return serverReply.Bytes(), nil
 }
+
+func (h *CryptoSetup) handleCHLO(data []byte, cryptoData map[Tag][]byte) ([]byte, error) {
+	// We have a CHLO matching our server config, we can continue with the 0-RTT handshake
+	sharedSecret, err := h.scfg.kex.CalculateSharedKey(cryptoData[TagPUBS])
+	if err != nil {
+		return nil, err
+	}
+	var nonce bytes.Buffer
+	nonce.Write(cryptoData[TagNONC])
+	nonce.Write(h.nonce)
+
+	h.secureAEAD, err = crypto.DeriveKeysChacha20(false, sharedSecret, nonce.Bytes(), h.connID, data, h.scfg.Get(), h.scfg.signer.GetCertUncompressed())
+	if err != nil {
+		return nil, err
+	}
+	// TODO: Use new curve
+	h.forwardSecureAEAD, err = crypto.DeriveKeysChacha20(true, sharedSecret, nonce.Bytes(), h.connID, data, h.scfg.Get(), h.scfg.signer.GetCertUncompressed())
+	if err != nil {
+		return nil, err
+	}
+
+	var reply bytes.Buffer
+	WriteHandshakeMessage(&reply, TagSHLO, map[Tag][]byte{
+		TagPUBS: h.scfg.kex.PublicKey(),
+		TagSNO:  h.nonce,
+		TagVER:  protocol.SupportedVersionsAsTags,
+		TagICSL: []byte{0x1e, 0x00, 0x00, 0x00}, //30
+		TagMSPC: []byte{0x64, 0x00, 0x00, 0x00}, //100
+	})
+	return reply.Bytes(), nil
+}
diff --git a/handshake/crypto_setup_test.go b/handshake/crypto_setup_test.go
index 3093f793..6d0e6893 100644
--- a/handshake/crypto_setup_test.go
+++ b/handshake/crypto_setup_test.go
@@ -10,7 +10,7 @@ import (
 type mockKEX struct{}
 
 func (*mockKEX) PublicKey() []byte {
-	return []byte("publickey")
+	return []byte("pubs-s")
 }
 func (*mockKEX) CalculateSharedKey(otherPublic []byte) ([]byte, error) {
 	return []byte("shared key"), nil
@@ -55,11 +55,11 @@ var _ = Describe("Crypto setup", func() {
 	})
 
 	It("generates REJ messages", func() {
-		response, err := cs.handleInchoateCHLO(sampleCHLO)
+		response, err := cs.handleInchoateCHLO([]byte("chlo"))
 		Expect(err).ToNot(HaveOccurred())
 		Expect(response).To(HavePrefix("REJ"))
 		Expect(response).To(ContainSubstring("certcompressed"))
-		Expect(response).To(ContainSubstring("publickey"))
+		Expect(response).To(ContainSubstring("pubs-s"))
 		Expect(signer.gotCHLO).To(BeTrue())
 	})
 
@@ -69,4 +69,16 @@ var _ = Describe("Crypto setup", func() {
 		Expect(err).ToNot(HaveOccurred())
 		Expect(signer.gotCHLO).To(BeFalse())
 	})
+
+	It("generates SHLO messages", func() {
+		response, err := cs.handleCHLO([]byte("chlo-data"), map[Tag][]byte{
+			TagPUBS: []byte("pubs-c"),
+		})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(response).To(ContainSubstring("pubs-s")) // TODO: Should be new pubs
+		Expect(response).To(ContainSubstring(string(cs.nonce)))
+		Expect(response).To(ContainSubstring(string(protocol.SupportedVersionsAsTags)))
+		Expect(cs.secureAEAD).ToNot(BeNil())
+		Expect(cs.forwardSecureAEAD).ToNot(BeNil())
+	})
 })