From cee0ea235faccbb43b342b9ff45bbf7ca8eb1710 Mon Sep 17 00:00:00 2001
From: Seweryn Chlewicki <schlewicki@palantir.com>
Date: Tue, 5 Sep 2023 15:42:16 +0100
Subject: [PATCH] Add the nonce to the endpoint

---
 endpoint/crypto.go | 25 +++++++++++++++----------
 endpoint/labels.go | 10 +++++++++-
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/endpoint/crypto.go b/endpoint/crypto.go
index 1d6ebd1dd7..253cb227a5 100644
--- a/endpoint/crypto.go
+++ b/endpoint/crypto.go
@@ -29,6 +29,17 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const standardGcmNonceSize = 12
+
+// GenerateNonce creates a random nonce of a fixed size
+func GenerateNonce() ([]byte, error) {
+	nonce := make([]byte, standardGcmNonceSize)
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+	return []byte(base64.StdEncoding.EncodeToString(nonce)), nil
+}
+
 // EncryptText gzip input data and encrypts it using the supplied AES key
 func EncryptText(text string, aesKey []byte, nonceEncoded []byte) (string, error) {
 	block, err := aes.NewCipher(aesKey)
@@ -36,20 +47,14 @@ func EncryptText(text string, aesKey []byte, nonceEncoded []byte) (string, error
 		return "", err
 	}
 
-	gcm, err := cipher.NewGCM(block)
+	gcm, err := cipher.NewGCMWithNonceSize(block, standardGcmNonceSize)
 	if err != nil {
 		return "", err
 	}
 
-	nonce := make([]byte, gcm.NonceSize())
-	if nonceEncoded == nil {
-		if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
-			return "", err
-		}
-	} else {
-		if _, err = base64.StdEncoding.Decode(nonce, nonceEncoded); err != nil {
-			return "", err
-		}
+	nonce := make([]byte, standardGcmNonceSize)
+	if _, err = base64.StdEncoding.Decode(nonce, nonceEncoded); err != nil {
+		return "", err
 	}
 
 	data, err := compressData([]byte(text))
diff --git a/endpoint/labels.go b/endpoint/labels.go
index 54afaf8c82..571852e3b3 100644
--- a/endpoint/labels.go
+++ b/endpoint/labels.go
@@ -136,9 +136,15 @@ func (l Labels) Serialize(withQuotes bool, txtEncryptEnabled bool, aesKey []byte
 		return l.SerializePlain(withQuotes)
 	}
 
-	var encryptionNonce []byte = nil
+	var encryptionNonce []byte
 	if extractedNonce, nonceExists := l[txtEncryptionNonce]; nonceExists {
 		encryptionNonce = []byte(extractedNonce)
+	} else {
+		var err error
+		encryptionNonce, err = GenerateNonce()
+		if err != nil {
+			log.Fatalf("Failed to generate cryptographic nonce %#v.", err)
+		}
 	}
 
 	text := l.SerializePlain(false)
@@ -150,6 +156,8 @@ func (l Labels) Serialize(withQuotes bool, txtEncryptEnabled bool, aesKey []byte
 		log.Fatalf("Failed to encrypt the text %#v using the encryption key %#v. Got error %#v.", text, aesKey, err)
 	}
 
+	l[txtEncryptionNonce] = string(encryptionNonce)
+
 	if withQuotes {
 		text = fmt.Sprintf("\"%s\"", text)
 	}