feat: Secure document encryption key exchange

cli/collection_create.go

@@ -17,8 +17,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/sourcenetwork/defradb/client"
-	"github.com/sourcenetwork/defradb/datastore"
-	"github.com/sourcenetwork/defradb/internal/db"
 	"github.com/sourcenetwork/defradb/internal/encryption"
 )
 
@@ -89,8 +87,7 @@ Example: create from stdin:
 				return cmd.Usage()
 			}
 
-			txn, _ := db.TryGetContextTxn(cmd.Context())
-			setContextDocEncryption(cmd, shouldEncryptDoc, encryptedFields, txn)
+			setContextDocEncryption(cmd, shouldEncryptDoc, encryptedFields)
 
 			if client.IsJSONArray(docData) {
 				docs, err := client.NewDocsFromJSON(docData, col.Definition())
@@ -116,14 +113,11 @@ Example: create from stdin:
 }
 
 // setContextDocEncryption sets doc encryption for the current command context.
-func setContextDocEncryption(cmd *cobra.Command, shouldEncryptDoc bool, encryptFields []string, txn datastore.Txn) {
+func setContextDocEncryption(cmd *cobra.Command, shouldEncryptDoc bool, encryptFields []string) {
 	if !shouldEncryptDoc && len(encryptFields) == 0 {
 		return
 	}
 	ctx := cmd.Context()
-	if txn != nil {
-		ctx = encryption.ContextWithStore(ctx, txn)
-	}
 	ctx = encryption.SetContextConfigFromParams(ctx, shouldEncryptDoc, encryptFields)
 	cmd.SetContext(ctx)
 }

client/db.go

@@ -52,6 +52,11 @@ type DB interface {
 	// It sits within the rootstore returned by [Root].
 	Blockstore() datastore.Blockstore
 
+	// Encstore returns the store, that contains all known encryption keys for documents and their fields.
+	//
+	// It sits within the rootstore returned by [Root].
+	Encstore() datastore.Blockstore
+
 	// Peerstore returns the peerstore where known host information is stored.
 	//
 	// It sits within the rootstore returned by [Root].

crypto/aes.go

@@ -0,0 +1,94 @@
+// Copyright 2024 Democratized Data Foundation
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package crypto
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+)
+
+// EncryptAES encrypts data using AES-GCM with a provided key and additional data.
+// It generates a nonce internally and optionally prepends it to the cipherText.
+//
+// Parameters:
+//   - plainText: The data to be encrypted
+//   - key: The AES encryption key
+//   - additionalData: Additional authenticated data (AAD) to be used in the encryption
+//   - prependNonce: If true, the nonce is prepended to the returned cipherText
+//
+// Returns:
+//   - cipherText: The encrypted data, with the nonce prepended if prependNonce is true
+//   - nonce: The generated nonce
+//   - error: Any error encountered during the encryption process
+func EncryptAES(plainText, key, additionalData []byte, prependNonce bool) ([]byte, []byte, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	nonce, err := generateNonceFunc()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	aesGCM, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	var cipherText []byte
+	if prependNonce {
+		cipherText = aesGCM.Seal(nonce, nonce, plainText, additionalData)
+	} else {
+		cipherText = aesGCM.Seal(nil, nonce, plainText, additionalData)
+	}
+
+	return cipherText, nonce, nil
+}
+
+// DecryptAES decrypts AES-GCM encrypted data with a provided key and additional data.
+// If no separate nonce is provided, it assumes the nonce is prepended to the cipherText.
+//
+// Parameters:
+//   - nonce: The nonce used for decryption. If empty, it's assumed to be prepended to cipherText
+//   - cipherText: The data to be decrypted
+//   - key: The AES decryption key
+//   - additionalData: Additional authenticated data (AAD) used during encryption
+//
+// Returns:
+//   - plainText: The decrypted data
+//   - error: Any error encountered during the decryption process, including authentication failures
+func DecryptAES(nonce, cipherText, key, additionalData []byte) ([]byte, error) {
+	if len(nonce) == 0 {
+		if len(cipherText) < AESNonceSize {
+			return nil, ErrCipherTextTooShort
+		}
+		nonce = cipherText[:AESNonceSize]
+		cipherText = cipherText[AESNonceSize:]
+	}
+
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	aesGCM, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+
+	plainText, err := aesGCM.Open(nil, nonce, cipherText, additionalData)
+	if err != nil {
+		return nil, err
+	}
+
+	return plainText, nil
+}

crypto/aes_test.go

@@ -0,0 +1,175 @@
+// Copyright 2024 Democratized Data Foundation
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package crypto
+
+import (
+	"bytes"
+	"crypto/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncryptAES(t *testing.T) {
+	validKey := make([]byte, 32) // AES-256
+	_, err := rand.Read(validKey)
+	require.NoError(t, err)
+	validPlaintext := []byte("Hello, World!")
+	validAAD := []byte("Additional Authenticated Data")
+
+	tests := []struct {
+		name           string
+		plainText      []byte
+		key            []byte
+		additionalData []byte
+		prependNonce   bool
+		expectError    bool
+		errorContains  string
+	}{
+		{
+			name:           "Valid encryption with prepended nonce",
+			plainText:      validPlaintext,
+			key:            validKey,
+			additionalData: validAAD,
+			prependNonce:   true,
+			expectError:    false,
+		},
+		{
+			name:           "Valid encryption without prepended nonce",
+			plainText:      validPlaintext,
+			key:            validKey,
+			additionalData: validAAD,
+			prependNonce:   false,
+			expectError:    false,
+		},
+		{
+			name:           "Invalid key size",
+			plainText:      validPlaintext,
+			key:            make([]byte, 31), // Invalid key size
+			additionalData: validAAD,
+			prependNonce:   true,
+			expectError:    true,
+			errorContains:  "invalid key size",
+		},
+		{
+			name:           "Nil plaintext",
+			plainText:      nil,
+			key:            validKey,
+			additionalData: validAAD,
+			prependNonce:   true,
+			expectError:    false, // AES-GCM can encrypt nil/empty plaintext
+		},
+		{
+			name:           "Nil additional data",
+			plainText:      validPlaintext,
+			key:            validKey,
+			additionalData: nil,
+			prependNonce:   true,
+			expectError:    false, // Nil AAD is valid
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cipherText, nonce, err := EncryptAES(tt.plainText, tt.key, tt.additionalData, tt.prependNonce)
+
+			if tt.expectError {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorContains)
+			} else {
+				require.NoError(t, err)
+				if tt.prependNonce {
+					require.Greater(t, len(cipherText), len(nonce), "Ciphertext length not greater than nonce length")
+				} else {
+					require.Equal(t, AESNonceSize, len(nonce), "Nonce length != AESNonceSize")
+				}
+			}
+		})
+	}
+}
+
+func TestDecryptAES(t *testing.T) {
+	validKey := make([]byte, 32) // AES-256
+	_, err := rand.Read(validKey)
+	require.NoError(t, err)
+	validPlaintext := []byte("Hello, World!")
+	validAAD := []byte("Additional Authenticated Data")
+	validCiphertext, validNonce, _ := EncryptAES(validPlaintext, validKey, validAAD, true)
+
+	tests := []struct {
+		name           string
+		nonce          []byte
+		cipherText     []byte
+		key            []byte
+		additionalData []byte
+		expectError    bool
+		errorContains  string
+	}{
+		{
+			name:           "Valid decryption",
+			nonce:          nil, // Should be extracted from cipherText
+			cipherText:     validCiphertext,
+			key:            validKey,
+			additionalData: validAAD,
+			expectError:    false,
+		},
+		{
+			name:           "Invalid key size",
+			nonce:          validNonce,
+			cipherText:     validCiphertext[AESNonceSize:],
+			key:            make([]byte, 31), // Invalid key size
+			additionalData: validAAD,
+			expectError:    true,
+			errorContains:  "invalid key size",
+		},
+		{
+			name:           "Ciphertext too short",
+			nonce:          nil,
+			cipherText:     make([]byte, AESNonceSize-1), // Too short to contain nonce
+			key:            validKey,
+			additionalData: validAAD,
+			expectError:    true,
+			errorContains:  errCipherTextTooShort,
+		},
+		{
+			name:           "Invalid additional data",
+			nonce:          validNonce,
+			cipherText:     validCiphertext[AESNonceSize:],
+			key:            validKey,
+			additionalData: []byte("Wrong AAD"),
+			expectError:    true,
+			errorContains:  "message authentication failed",
+		},
+		{
+			name:           "Tampered ciphertext",
+			nonce:          validNonce,
+			cipherText:     append([]byte{0}, validCiphertext[AESNonceSize+1:]...),
+			key:            validKey,
+			additionalData: validAAD,
+			expectError:    true,
+			errorContains:  "message authentication failed",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			plainText, err := DecryptAES(tt.nonce, tt.cipherText, tt.key, tt.additionalData)
+
+			if tt.expectError {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorContains)
+			} else {
+				require.NoError(t, err)
+				require.True(t, bytes.Equal(plainText, validPlaintext), "Decrypted plaintext does not match original")
+			}
+		})
+	}
+}