feat: Secure document encryption key exchange

cli/collection_create.go

@@ -122,7 +122,7 @@ func setContextDocEncryption(cmd *cobra.Command, shouldEncryptDoc bool, encryptF
 	}
 	ctx := cmd.Context()
 	if txn != nil {
-		ctx = encryption.ContextWithStore(ctx, txn)
+		ctx = encryption.ContextWithStore(ctx, txn.Encstore())
 	}
 	ctx = encryption.SetContextConfigFromParams(ctx, shouldEncryptDoc, encryptFields)
 	cmd.SetContext(ctx)

client/db.go

@@ -50,6 +50,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.DSReaderWriter
+
 	// Peerstore returns the peerstore where known host information is stored.
 	//
 	// It sits within the rootstore returned by [Root].

internal/encryption/aes.go → crypto/aes.go

@@ -8,69 +8,65 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
-package encryption
+package crypto
 
 import (
 	"crypto/aes"
 	"crypto/cipher"
-	"encoding/base64"
 	"fmt"
 )
 
 // EncryptAES encrypts data using AES-GCM with a provided key.
-func EncryptAES(plainText, key []byte) ([]byte, error) {
+// The nonce is prepended to the cipherText.
+func EncryptAES(plainText, key, additionalData []byte, prependNonce bool) ([]byte, []byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	nonce, err := generateNonceFunc()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	aesGCM, err := cipher.NewGCM(block)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	cipherText := aesGCM.Seal(nonce, nonce, plainText, nil)
-
-	buf := make([]byte, base64.StdEncoding.EncodedLen(len(cipherText)))
-	base64.StdEncoding.Encode(buf, cipherText)
+	var cipherText []byte
+	if prependNonce {
+		cipherText = aesGCM.Seal(nonce, nonce, plainText, additionalData)
+	} else {
+		cipherText = aesGCM.Seal(nil, nonce, plainText, additionalData)
+	}
 
-	return buf, nil
+	return cipherText, nonce, nil
 }
 
 // DecryptAES decrypts AES-GCM encrypted data with a provided key.
-func DecryptAES(cipherTextBase64, key []byte) ([]byte, error) {
-	cipherText := make([]byte, base64.StdEncoding.DecodedLen(len(cipherTextBase64)))
-	n, err := base64.StdEncoding.Decode(cipherText, []byte(cipherTextBase64))
-
-	if err != nil {
-		return nil, err
+// The nonce is expected to be prepended to the cipherText.
+func DecryptAES(nonce, cipherText, key, additionalData []byte) ([]byte, error) {
+	if len(nonce) == 0 {
+		if len(cipherText) < AESNonceSize {
+			// TODO return typed error
+			return nil, fmt.Errorf("cipherText too short")
+		}
+		nonce = cipherText[:AESNonceSize]
+		cipherText = cipherText[AESNonceSize:]
 	}
 
-	cipherText = cipherText[:n]
-
 	block, err := aes.NewCipher(key)
 	if err != nil {
 		return nil, err
 	}
 
-	if len(cipherText) < nonceLength {
-		return nil, fmt.Errorf("cipherText too short")
-	}
-
-	nonce := cipherText[:nonceLength]
-	cipherText = cipherText[nonceLength:]
-
 	aesGCM, err := cipher.NewGCM(block)
 	if err != nil {
 		return nil, err
 	}
 
-	plainText, err := aesGCM.Open(nil, nonce, cipherText, nil)
+	plainText, err := aesGCM.Open(nil, nonce, cipherText, additionalData)
 	if err != nil {
 		return nil, err
 	}

crypto/ecies.go

@@ -0,0 +1,164 @@
+// Current Implementation Analysis:
+// 1. Key Generation: Correctly uses X25519 for key generation.
+// 2. ECDH: Properly performs the ECDH operation.
+// 3. Key Derivation: Uses SHA-256 on the shared secret, which is simplistic.
+// 4. Encryption: Uses AES (implementation not shown).
+// 5. MAC: Not implemented.
+
+// Improvements Needed:
+// 1. Use a proper Key Derivation Function (KDF)
+// 2. Implement HMAC for message authentication
+// 3. Use authenticated encryption (e.g., AES-GCM) instead of AES
+// 4. Standardize the output format
+
+// Here's an improved version of the EncryptECDH and DecryptECDH functions:
+
+package crypto
+
+import (
+	"crypto/ecdh"
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"fmt"
+
+	"golang.org/x/crypto/hkdf"
+)
+
+const X25519PublicKeySize = 32
+const HMACSize = 32
+const AESKeySize = 32
+
+const minCipherTextSize = 16
+
+// GenerateX25519 generates a new X25519 private key.
+func GenerateX25519() (*ecdh.PrivateKey, error) {
+	return ecdh.X25519().GenerateKey(rand.Reader)
+}
+
+// X25519PublicKeyFromBytes creates a new X25519 public key from the given bytes.
+func X25519PublicKeyFromBytes(publicKeyBytes []byte) (*ecdh.PublicKey, error) {
+	return ecdh.X25519().NewPublicKey(publicKeyBytes)
+}
+
+// EncryptECIES encrypts plaintext using a custom Elliptic Curve Integrated Encryption Scheme (ECIES)
+// with X25519 for key agreement, HKDF for key derivation, AES for encryption, and HMAC for authentication.
+//
+// The function:
+// - Generates an ephemeral X25519 key pair
+// - Performs ECDH with the provided public key
+// - Derives encryption and HMAC keys using HKDF
+// - Encrypts the plaintext using a custom AES encryption function
+// - Computes an HMAC over the ciphertext
+//
+// The output format is: [ephemeral public key | encrypted data (including nonce) | HMAC]
+//
+// Parameters:
+//   - plainText: The message to encrypt
+//   - publicKey: The recipient's X25519 public key
+//   - associatedData: Optional associated data for additional authentication
+//
+// Returns:
+//   - Byte slice containing the encrypted message and necessary metadata for decryption
+//   - Error if any step of the encryption process fails
+func EncryptECIES(plainText []byte, publicKey *ecdh.PublicKey, associatedData []byte) ([]byte, error) {
+	ephemeralPrivate, err := GenerateX25519()
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate ephemeral key: %w", err)
+	}
+	ephemeralPublic := ephemeralPrivate.PublicKey()
+
+	sharedSecret, err := ephemeralPrivate.ECDH(publicKey)
+	if err != nil {
+		return nil, fmt.Errorf("ECDH failed: %w", err)
+	}
+
+	kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
+	aesKey := make([]byte, AESKeySize)
+	hmacKey := make([]byte, HMACSize)
+	if _, err := kdf.Read(aesKey); err != nil {
+		return nil, fmt.Errorf("KDF failed for AES key: %w", err)
+	}
+	if _, err := kdf.Read(hmacKey); err != nil {
+		return nil, fmt.Errorf("KDF failed for HMAC key: %w", err)
+	}
+
+	fullAssociatedData := append(ephemeralPublic.Bytes(), associatedData...)
+	cipherText, _, err := EncryptAES(plainText, aesKey, fullAssociatedData, true)
+	if err != nil {
+		return nil, fmt.Errorf("failed to encrypt: %w", err)
+	}
+
+	mac := hmac.New(sha256.New, hmacKey)
+	mac.Write(cipherText)
+	macSum := mac.Sum(nil)
+
+	result := append(ephemeralPublic.Bytes(), cipherText...)
+	result = append(result, macSum...)
+
+	return result, nil
+}
+
+// DecryptECIES decrypts ciphertext encrypted with EncryptECIES using the provided private key.
+//
+// The function:
+// - Extracts the ephemeral public key from the ciphertext
+// - Performs ECDH with the provided private key
+// - Derives decryption and HMAC keys using HKDF
+// - Verifies the HMAC
+// - Decrypts the message using a custom AES decryption function
+//
+// The expected input format is: [ephemeral public key | encrypted data (including nonce) | HMAC]
+//
+// Parameters:
+//   - cipherText: The encrypted message, including all necessary metadata
+//   - privateKey: The recipient's X25519 private key
+//   - associatedData: Optional associated data used during encryption for additional authentication
+//
+// Returns:
+//   - Byte slice containing the decrypted plaintext
+//   - Error if any step of the decryption process fails, including authentication failure
+func DecryptECIES(cipherText []byte, privateKey *ecdh.PrivateKey, associatedData []byte) ([]byte, error) {
+	if len(cipherText) < X25519PublicKeySize+AESNonceSize+HMACSize+minCipherTextSize {
+		return nil, fmt.Errorf("ciphertext too short")
+	}
+
+	ephemeralPublicBytes := cipherText[:X25519PublicKeySize]
+	ephemeralPublic, err := ecdh.X25519().NewPublicKey(ephemeralPublicBytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse ephemeral public key: %w", err)
+	}
+
+	sharedSecret, err := privateKey.ECDH(ephemeralPublic)
+	if err != nil {
+		return nil, fmt.Errorf("ECDH failed: %w", err)
+	}
+
+	kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
+	aesKey := make([]byte, AESKeySize)
+	hmacKey := make([]byte, HMACSize)
+	if _, err := kdf.Read(aesKey); err != nil {
+		return nil, fmt.Errorf("KDF failed for AES key: %w", err)
+	}
+	if _, err := kdf.Read(hmacKey); err != nil {
+		return nil, fmt.Errorf("KDF failed for HMAC key: %w", err)
+	}
+
+	macSum := cipherText[len(cipherText)-HMACSize:]
+	cipherTextWithNonce := cipherText[X25519PublicKeySize : len(cipherText)-HMACSize]
+
+	mac := hmac.New(sha256.New, hmacKey)
+	mac.Write(cipherTextWithNonce)
+	expectedMAC := mac.Sum(nil)
+	if !hmac.Equal(macSum, expectedMAC) {
+		return nil, fmt.Errorf("HMAC verification failed")
+	}
+
+	fullAssociatedData := append(ephemeralPublicBytes, associatedData...)
+	plainText, err := DecryptAES(nil, cipherTextWithNonce, aesKey, fullAssociatedData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt: %w", err)
+	}
+
+	return plainText, nil
+}

crypto/ephemeral.go

@@ -0,0 +1,47 @@
+package crypto
+
+import (
+	"crypto/ecdh"
+	"crypto/rand"
+	"crypto/sha256"
+	"fmt"
+)
+
+const (
+	// EphemeralKeyLength is the size of the ECDH ephemeral key in bytes.
+	EphemeralKeyLength = 65
+)
+
+// EncryptWithEphemeralKey encrypts a key using a randomly generated ephemeral ECDH key and a provided public key.
+// It returns the encrypted key prepended with the ephemeral public key.
+func EncryptWithEphemeralKey(plainText, publicKeyBytes []byte) ([]byte, error) {
+	ephemeralPriv, err := ecdh.P256().GenerateKey(rand.Reader)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate ephemeral key: %w", err)
+	}
+
+	ephPubKeyBytes := ephemeralPriv.PublicKey().Bytes()
+	sharedSecret := sha256.Sum256(append(ephPubKeyBytes, publicKeyBytes...))
+
+	return append(ephPubKeyBytes, xorBytes(plainText, sharedSecret[:])...), nil
+}
+
+func xorBytes(data, xor []byte) []byte {
+	result := make([]byte, len(data))
+	for i := range data {
+		result[i] = data[i] ^ xor[i%len(xor)]
+	}
+	return result
+}
+
+// DecryptWithEphemeralKey decrypts data that was encrypted using EncryptWithEphemeralKey.
+// It expects the input to be the ephemeral public key followed by the encrypted data.
+func DecryptWithEphemeralKey(encryptedData, publicKeyBytes []byte) ([]byte, error) {
+	ephPubKeyBytes := encryptedData[:EphemeralKeyLength]
+	cipherText := make([]byte, len(encryptedData)-EphemeralKeyLength)
+	copy(cipherText, encryptedData[EphemeralKeyLength:])
+
+	sharedSecret := sha256.Sum256(append(ephPubKeyBytes, publicKeyBytes...))
+
+	return xorBytes(cipherText, sharedSecret[:]), nil
+}

crypto/nonce.go

@@ -0,0 +1,52 @@
+// 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/rand"
+	"errors"
+	"io"
+	"os"
+	"strings"
+)
+
+const AESNonceSize = 12
+
+var generateNonceFunc = generateNonce
+
+func generateNonce() ([]byte, error) {
+	nonce := make([]byte, AESNonceSize)
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+
+	return nonce, nil
+}
+
+// generateTestNonce generates a deterministic nonce for testing.
+func generateTestNonce() ([]byte, error) {
+	nonce := []byte("deterministic nonce for testing")
+
+	if len(nonce) < AESNonceSize {
+		return nil, errors.New("nonce length is longer than available deterministic nonce")
+	}
+
+	return nonce[:AESNonceSize], nil
+}
+
+func init() {
+	arg := os.Args[0]
+	// If the binary is a test binary, use a deterministic nonce.
+	// TODO: We should try to find a better way to detect this https://github.com/sourcenetwork/defradb/issues/2801
+	if strings.HasSuffix(arg, ".test") || strings.Contains(arg, "/defradb/tests/") {
+		generateNonceFunc = generateTestNonce
+	}
+}

datastore/multi.go

@@ -33,8 +33,7 @@ type multistore struct {
 	head   DSReaderWriter
 	peer   DSBatching
 	system DSReaderWriter
-	// block DSReaderWriter
-	dag Blockstore
+	dag    Blockstore
 }
 
 var _ MultiStore = (*multistore)(nil)