From 927fa1b4bbcf5b0f8fe64c03d19f5d4ffb5313d5 Mon Sep 17 00:00:00 2001
From: Jonathan Rudenberg <jonathan@titanous.com>
Date: Tue, 20 Apr 2021 18:02:02 -0400
Subject: [PATCH] Add ErrMaxNonce before n overflows

---
 noise_test.go | 15 +++++++++++++++
 state.go      | 23 +++++++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/noise_test.go b/noise_test.go
index 783a9eb..ab7c0ca 100644
--- a/noise_test.go
+++ b/noise_test.go
@@ -2,6 +2,7 @@ package noise
 
 import (
 	"encoding/hex"
+	"math"
 	"testing"
 
 	. "gopkg.in/check.v1"
@@ -582,9 +583,14 @@ func (NoiseSuite) TestRekey(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(string(serverMessage), Equals, string(res))
 
+	preNonce := csR1.Nonce()
+
 	csR1.Rekey()
 	csI1.Rekey()
 
+	postNonce := csR1.Nonce()
+	c.Assert(preNonce, Equals, postNonce)
+
 	serverMessage = []byte("bye bye")
 	msg, err = csR1.Encrypt(nil, nil, serverMessage)
 	c.Assert(err, IsNil)
@@ -600,4 +606,13 @@ func (NoiseSuite) TestRekey(c *C) {
 	res, err = csI1.Decrypt(nil, nil, msg)
 	c.Assert(err, NotNil)
 	c.Assert(string(serverMessage), Not(Equals), string(res))
+
+	// check nonce overflow handling
+	csI1.n = math.MaxUint64
+	msg, err = csI1.Encrypt(nil, nil, nil)
+	c.Assert(err, Equals, ErrMaxNonce)
+	c.Assert(msg, IsNil)
+	msg, err = csI1.Decrypt(nil, nil, nil)
+	c.Assert(err, Equals, ErrMaxNonce)
+	c.Assert(msg, IsNil)
 }
diff --git a/state.go b/state.go
index 4153e68..c4c766f 100644
--- a/state.go
+++ b/state.go
@@ -25,16 +25,25 @@ type CipherState struct {
 	invalid bool
 }
 
+// MaxNonce is the maximum value of n that is allowed. ErrMaxNonce is returned
+// by Encrypt and Decrypt after this has been reached. 2^64-1 is reserved for rekeys.
+const MaxNonce = uint64(math.MaxUint64) - 1
+
+var ErrMaxNonce = errors.New("noise: cipherstate has reached maximum n, a new handshake must be performed")
 var ErrCipherSuiteCopied = errors.New("noise: CipherSuite has been copied, state is invalid")
 
 // Encrypt encrypts the plaintext and then appends the ciphertext and an
 // authentication tag across the ciphertext and optional authenticated data to
 // out. This method automatically increments the nonce after every call, so
-// messages must be decrypted in the same order.
+// messages must be decrypted in the same order. ErrMaxNonce is returned after
+// the maximum nonce of 2^64-2 is reached.
 func (s *CipherState) Encrypt(out, ad, plaintext []byte) ([]byte, error) {
 	if s.invalid {
 		return nil, ErrCipherSuiteCopied
 	}
+	if s.n > MaxNonce {
+		return nil, ErrMaxNonce
+	}
 	out = s.c.Encrypt(out, s.n, ad, plaintext)
 	s.n++
 	return out, nil
@@ -43,11 +52,15 @@ func (s *CipherState) Encrypt(out, ad, plaintext []byte) ([]byte, error) {
 // Decrypt checks the authenticity of the ciphertext and authenticated data and
 // then decrypts and appends the plaintext to out. This method automatically
 // increments the nonce after every call, messages must be provided in the same
-// order that they were encrypted with no missing messages.
+// order that they were encrypted with no missing messages. ErrMaxNonce is
+// returned after the maximum nonce of 2^64-2 is reached.
 func (s *CipherState) Decrypt(out, ad, ciphertext []byte) ([]byte, error) {
 	if s.invalid {
 		return nil, ErrCipherSuiteCopied
 	}
+	if s.n > MaxNonce {
+		return nil, ErrMaxNonce
+	}
 	out, err := s.c.Decrypt(out, s.n, ad, ciphertext)
 	s.n++
 	return out, err
@@ -64,6 +77,12 @@ func (s *CipherState) Cipher() Cipher {
 	return s.c
 }
 
+// Nonce returns the current value of n. This can be used to determine if a
+// new handshake should be performed due to approaching MaxNonce.
+func (s *CipherState) Nonce() uint64 {
+	return s.n
+}
+
 func (s *CipherState) Rekey() {
 	var zeros [32]byte
 	var out []byte