From c03f15dfdd1c193148740a88a444bca1bd25b4df Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 9 Nov 2016 17:39:39 +0700
Subject: [PATCH] add a method to set the diversification nonce in the
 CryptoSetup

---
 handshake/crypto_setup_client.go      | 15 ++++++++++++++-
 handshake/crypto_setup_client_test.go | 27 +++++++++++++++++++++++++++
 handshake/crypto_setup_interface.go   |  4 +++-
 handshake/crypto_setup_server.go      |  4 ++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/handshake/crypto_setup_client.go b/handshake/crypto_setup_client.go
index 9a28a4d93..628c16236 100644
--- a/handshake/crypto_setup_client.go
+++ b/handshake/crypto_setup_client.go
@@ -21,13 +21,15 @@ type cryptoSetupClient struct {
 	cryptoStream utils.Stream
 
 	serverConfig *serverConfigClient
+	diversificationNonce []byte
 }
 
 var _ crypto.AEAD = &cryptoSetupClient{}
 var _ CryptoSetup = &cryptoSetupClient{}
 
 var (
-	errNoObitForClientNonce = errors.New("No OBIT for client nonce available")
+	errNoObitForClientNonce             = errors.New("No OBIT for client nonce available")
+	errConflictingDiversificationNonces = errors.New("Received two different diversification nonces")
 )
 
 // NewCryptoSetupClient creates a new CryptoSetup instance for a client
@@ -81,6 +83,17 @@ func (h *cryptoSetupClient) Seal(dst, src []byte, packetNumber protocol.PacketNu
 }
 
 func (h *cryptoSetupClient) DiversificationNonce() []byte {
+	panic("not needed for cryptoSetupClient")
+}
+
+func (h *cryptoSetupClient) SetDiversificationNonce(data []byte) error {
+	if len(h.diversificationNonce) == 0 {
+		h.diversificationNonce = data
+		return nil
+	}
+	if !bytes.Equal(h.diversificationNonce, data) {
+		return errConflictingDiversificationNonces
+	}
 	return nil
 }
 
diff --git a/handshake/crypto_setup_client_test.go b/handshake/crypto_setup_client_test.go
index 7cb4ee4b6..4a3545dd5 100644
--- a/handshake/crypto_setup_client_test.go
+++ b/handshake/crypto_setup_client_test.go
@@ -36,6 +36,33 @@ var _ = Describe("Crypto setup", func() {
 		})
 	})
 
+	Context("Diversification Nonces", func() {
+		It("sets a diversification nonce", func() {
+			nonce := []byte("foobar")
+			err := cs.SetDiversificationNonce(nonce)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(cs.diversificationNonce).To(Equal(nonce))
+		})
+
+		It("doesn't do anything when called multiple times with the same nonce", func() {
+			nonce := []byte("foobar")
+			err := cs.SetDiversificationNonce(nonce)
+			Expect(err).ToNot(HaveOccurred())
+			err = cs.SetDiversificationNonce(nonce)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(cs.diversificationNonce).To(Equal(nonce))
+		})
+
+		It("rejects a different diversification nonce", func() {
+			nonce1 := []byte("foobar")
+			nonce2 := []byte("raboof")
+			err := cs.SetDiversificationNonce(nonce1)
+			Expect(err).ToNot(HaveOccurred())
+			err = cs.SetDiversificationNonce(nonce2)
+			Expect(err).To(MatchError(errConflictingDiversificationNonces))
+		})
+	})
+
 	Context("Client Nonce generation", func() {
 		BeforeEach(func() {
 			cs.serverConfig.obit = []byte{0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8}
diff --git a/handshake/crypto_setup_interface.go b/handshake/crypto_setup_interface.go
index 1452b5af2..4822cfb8f 100644
--- a/handshake/crypto_setup_interface.go
+++ b/handshake/crypto_setup_interface.go
@@ -7,8 +7,10 @@ type CryptoSetup interface {
 	HandleCryptoStream() error
 	Open(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, error)
 	Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) []byte
-	DiversificationNonce() []byte
 	LockForSealing()
 	UnlockForSealing()
 	HandshakeComplete() bool
+	// TODO: clean up this interface
+	DiversificationNonce() []byte         // only needed for cryptoSetupServer
+	SetDiversificationNonce([]byte) error // only needed for cryptoSetupClient
 }
diff --git a/handshake/crypto_setup_server.go b/handshake/crypto_setup_server.go
index 2c154c6c6..b8a8d0bbe 100644
--- a/handshake/crypto_setup_server.go
+++ b/handshake/crypto_setup_server.go
@@ -361,6 +361,10 @@ func (h *cryptoSetupServer) DiversificationNonce() []byte {
 	return h.diversificationNonce
 }
 
+func (h *cryptoSetupServer) SetDiversificationNonce(data []byte) error {
+	panic("not needed for cryptoSetupServer")
+}
+
 // LockForSealing should be called before Seal(). It is needed so that diversification nonces can be obtained before packets are sealed, and the AEADs are not changed in the meantime.
 func (h *cryptoSetupServer) LockForSealing() {
 	h.mutex.RLock()