From af56ff2aca0f761773c4776429993f8655d99cfa Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Thu, 28 Jul 2016 19:07:57 +0200
Subject: [PATCH] cache the ephermal key for up to 1 min

fixes #136
---
 handshake/crypto_setup.go        |  9 ++----
 handshake/crypto_setup_test.go   |  2 +-
 handshake/ephermal_cache.go      | 55 ++++++++++++++++++++++++++++++++
 handshake/ephermal_cache_test.go | 30 +++++++++++++++++
 protocol/server_parameters.go    |  3 ++
 5 files changed, 92 insertions(+), 7 deletions(-)
 create mode 100644 handshake/ephermal_cache.go
 create mode 100644 handshake/ephermal_cache_test.go

diff --git a/handshake/crypto_setup.go b/handshake/crypto_setup.go
index 943060f8c..6eec9502f 100644
--- a/handshake/crypto_setup.go
+++ b/handshake/crypto_setup.go
@@ -17,7 +17,7 @@ import (
 type KeyDerivationFunction func(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (crypto.AEAD, error)
 
 // KeyExchangeFunction is used to make a new KEX
-type KeyExchangeFunction func() (crypto.KeyExchange, error)
+type KeyExchangeFunction func() crypto.KeyExchange
 
 // The CryptoSetup handles all things crypto for the Session
 type CryptoSetup struct {
@@ -61,7 +61,7 @@ func NewCryptoSetup(
 		version:                     version,
 		scfg:                        scfg,
 		keyDerivation:               crypto.DeriveKeysAESGCM,
-		keyExchange:                 crypto.NewCurve25519KEX,
+		keyExchange:                 getEphermalKEX,
 		cryptoStream:                cryptoStream,
 		connectionParametersManager: connectionParametersManager,
 		aeadChanged:                 aeadChanged,
@@ -263,10 +263,7 @@ func (h *CryptoSetup) handleCHLO(sni string, data []byte, cryptoData map[Tag][]b
 	var fsNonce bytes.Buffer
 	fsNonce.Write(cryptoData[TagNONC])
 	fsNonce.Write(nonce)
-	ephermalKex, err := h.keyExchange()
-	if err != nil {
-		return nil, err
-	}
+	ephermalKex := h.keyExchange()
 	ephermalSharedSecret, err := ephermalKex.CalculateSharedKey(cryptoData[TagPUBS])
 	if err != nil {
 		return nil, err
diff --git a/handshake/crypto_setup_test.go b/handshake/crypto_setup_test.go
index cbb9b6d3e..25f6d3162 100644
--- a/handshake/crypto_setup_test.go
+++ b/handshake/crypto_setup_test.go
@@ -164,7 +164,7 @@ var _ = Describe("Crypto setup", func() {
 		cs, err = NewCryptoSetup(protocol.ConnectionID(42), ip, v, scfg, stream, cpm, aeadChanged)
 		Expect(err).NotTo(HaveOccurred())
 		cs.keyDerivation = mockKeyDerivation
-		cs.keyExchange = func() (crypto.KeyExchange, error) { return &mockKEX{ephermal: true}, nil }
+		cs.keyExchange = func() crypto.KeyExchange { return &mockKEX{ephermal: true} }
 	})
 
 	Context("diversification nonce", func() {
diff --git a/handshake/ephermal_cache.go b/handshake/ephermal_cache.go
new file mode 100644
index 000000000..8dcd7fcc6
--- /dev/null
+++ b/handshake/ephermal_cache.go
@@ -0,0 +1,55 @@
+package handshake
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/lucas-clemente/quic-go/crypto"
+	"github.com/lucas-clemente/quic-go/protocol"
+	"github.com/lucas-clemente/quic-go/utils"
+)
+
+var (
+	kexLifetime = protocol.EphermalKeyLifetime
+	kexCurrent  crypto.KeyExchange
+	kexSetOnce  sync.Once
+	kexMutex    sync.RWMutex
+)
+
+// getEphermalKEX returns the currently active KEX, which changes every protocol.EphermalKeyLifetime
+// See the explanation from the QUIC crypto doc:
+//
+// A single connection is the usual scope for forward security, but the security
+// difference between an ephemeral key used for a single connection, and one
+// used for all connections for 60 seconds is negligible. Thus we can amortise
+// the Diffie-Hellman key generation at the server over all the connections in a
+// small time span.
+func getEphermalKEX() crypto.KeyExchange {
+	kexSetOnce.Do(func() { go setKexRoutine() })
+	kexMutex.RLock()
+	defer kexMutex.RUnlock()
+	return kexCurrent
+}
+
+func setKexRoutine() {
+	for {
+		time.Sleep(kexLifetime)
+		kex, err := crypto.NewCurve25519KEX()
+		if err != nil {
+			utils.Errorf("could not set KEX: %s", err.Error())
+			continue
+		}
+		kexMutex.Lock()
+		kexCurrent = kex
+		kexMutex.Unlock()
+	}
+}
+
+func init() {
+	kex, err := crypto.NewCurve25519KEX()
+	if err != nil {
+		panic(fmt.Sprintf("Could not set KEX: %s", err.Error()))
+	}
+	kexCurrent = kex
+}
diff --git a/handshake/ephermal_cache_test.go b/handshake/ephermal_cache_test.go
new file mode 100644
index 000000000..bb77b29bb
--- /dev/null
+++ b/handshake/ephermal_cache_test.go
@@ -0,0 +1,30 @@
+package handshake
+
+import (
+	"time"
+
+	"github.com/lucas-clemente/quic-go/crypto"
+	"github.com/lucas-clemente/quic-go/protocol"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Ephermal KEX", func() {
+	It("has a consistent KEX", func() {
+		kex1 := getEphermalKEX()
+		Expect(kex1).ToNot(BeNil())
+		kex2 := getEphermalKEX()
+		Expect(kex2).ToNot(BeNil())
+		Expect(kex1).To(Equal(kex2))
+	})
+
+	It("changes KEX", func() {
+		kexLifetime = time.Millisecond
+		defer func() {
+			kexLifetime = protocol.EphermalKeyLifetime
+		}()
+		kex := getEphermalKEX()
+		Expect(kex).ToNot(BeNil())
+		Eventually(func() crypto.KeyExchange { return getEphermalKEX() }).ShouldNot(Equal(kex))
+	})
+})
diff --git a/protocol/server_parameters.go b/protocol/server_parameters.go
index 20a78414a..959e1ce39 100644
--- a/protocol/server_parameters.go
+++ b/protocol/server_parameters.go
@@ -66,3 +66,6 @@ const CryptoMaxParams = 128
 
 // CryptoParameterMaxLength is the upper limit for the length of a parameter in a crypto message.
 const CryptoParameterMaxLength = 2000
+
+// EphermalKeyLifetime is the lifetime of the ephermal key during the handshake, see handshake.getEphermalKEX.
+const EphermalKeyLifetime = time.Minute