engineccl: Support JWKS for encryption-at-rest keys and hook up v2

pkg/ccl/cliccl/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "demo.go",
         "ear.go",
         "flags.go",
+        "gen.go",
         "mt.go",
         "mt_proxy.go",
         "mt_test_directory.go",
@@ -44,6 +45,7 @@ go_library(
         "@com_github_cockroachdb_errors//oserror",
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",
+        "@com_github_lestrrat_go_jwx//jwk",
         "@com_github_olekukonko_tablewriter//:tablewriter",
         "@com_github_spf13_cobra//:cobra",
         "@com_github_spf13_pflag//:pflag",
@@ -55,6 +57,7 @@ go_test(
     size = "medium",
     srcs = [
         "ear_test.go",
+        "gen_test.go",
         "main_test.go",
     ],
     data = glob(["testdata/**"]),
@@ -80,7 +83,9 @@ go_test(
         "//pkg/util/log",
         "//pkg/util/randutil",
         "@com_github_cockroachdb_datadriven//:datadriven",
+        "@com_github_cockroachdb_pebble//vfs",
         "@com_github_spf13_cobra//:cobra",
+        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/ccl/cliccl/ear_test.go

@@ -44,7 +44,7 @@ func TestDecrypt(t *testing.T) {
 
 	// Generate a new encryption key to use.
 	keyPath := filepath.Join(dir, "aes.key")
-	err := cli.GenEncryptionKeyCmd.RunE(nil, []string{keyPath})
+	err := genEncryptionKeyCmd.RunE(nil, []string{keyPath})
 	require.NoError(t, err)
 
 	// Spin up a new encrypted store.
@@ -124,7 +124,7 @@ func TestList(t *testing.T) {
 
 	// Generate a new encryption key to use.
 	keyPath := filepath.Join(dir, "aes.key")
-	err := cli.GenEncryptionKeyCmd.RunE(nil, []string{keyPath})
+	err := genEncryptionKeyCmd.RunE(nil, []string{keyPath})
 	require.NoError(t, err)
 
 	// Spin up a new encrypted store.

pkg/ccl/cliccl/gen.go

@@ -0,0 +1,142 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package cliccl
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+
+	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl/enginepbccl"
+	"github.com/cockroachdb/cockroach/pkg/cli"
+	"github.com/cockroachdb/errors"
+	"github.com/lestrrat-go/jwx/jwk"
+	"github.com/spf13/cobra"
+)
+
+var aesSizeFlag int
+var overwriteKeyFlag bool
+var keyVersionFlag int
+
+func genEncryptionKey(
+	encryptionKeyPath string, aesSize int, overwriteKey bool, keyVersion int,
+) error {
+	// Check encryptionKeySize is suitable for the encryption algorithm.
+	if aesSize != 128 && aesSize != 192 && aesSize != 256 {
+		return fmt.Errorf("store key size should be 128, 192, or 256 bits, got %d", aesSize)
+	}
+
+	keyID := make([]byte, 32)
+	if _, err := rand.Read(keyID); err != nil {
+		return fmt.Errorf("failed to create random key ID")
+	}
+	key := make([]byte, aesSize/8)
+	if _, err := rand.Read(key); err != nil {
+		return fmt.Errorf("failed to create key with size %d bytes", aesSize/8)
+	}
+
+	var b []byte
+	switch keyVersion {
+	case 1:
+		b = append(b, keyID...)
+		b = append(b, key...)
+	case 2:
+		var et enginepbccl.EncryptionType
+		switch aesSize {
+		case 128:
+			et = enginepbccl.EncryptionType_AES_128_CTR_V2
+		case 192:
+			et = enginepbccl.EncryptionType_AES_192_CTR_V2
+		case 256:
+			et = enginepbccl.EncryptionType_AES_256_CTR_V2
+		default:
+			// Redundant since we checked this at the start of the function too.
+			return fmt.Errorf("store key size should be 128, 192, or 256 bits, got %d", aesSize)
+		}
+
+		symKey := jwk.NewSymmetricKey()
+		if err := symKey.FromRaw(key); err != nil {
+			return errors.Wrap(err, "error setting key bytes")
+		}
+		if err := symKey.Set(jwk.KeyIDKey, hex.EncodeToString(keyID)); err != nil {
+			return errors.Wrap(err, "error setting key id")
+		}
+		alg, err := et.JWKAlgorithm()
+		if err != nil {
+			return err
+		}
+		if err := symKey.Set(jwk.AlgorithmKey, alg); err != nil {
+			return errors.Wrap(err, "error setting algorithm")
+		}
+
+		keySet := jwk.NewSet()
+		keySet.Add(symKey)
+
+		b, err = json.Marshal(keySet)
+		if err != nil {
+			return errors.Wrap(err, "error writing key to json: %s")
+		}
+	default:
+		return fmt.Errorf("unsupported version %d", keyVersion)
+	}
+
+	// Write key to the file with owner read/write permission.
+	openMode := os.O_WRONLY | os.O_CREATE | os.O_TRUNC
+	if !overwriteKey {
+		openMode |= os.O_EXCL
+	}
+
+	f, err := os.OpenFile(encryptionKeyPath, openMode, 0600)
+	if err != nil {
+		return err
+	}
+	_, err = f.Write(b)
+	if err1 := f.Close(); err == nil {
+		err = err1
+	}
+	return err
+}
+
+// genEncryptionKeyCmd is a command to generate a store key for Encryption At
+// Rest.
+var genEncryptionKeyCmd = &cobra.Command{
+	Use:   "encryption-key <key-file>",
+	Short: "generate store key for encryption at rest",
+	Long: `Generate store key for encryption at rest.
+
+Generates a key suitable for use as a store key for Encryption At Rest.
+The resulting key file will be 32 bytes (random key ID) + key_size in bytes.
+`,
+	Args: cobra.ExactArgs(1),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		encryptionKeyPath := args[0]
+
+		err := genEncryptionKey(encryptionKeyPath, aesSizeFlag, overwriteKeyFlag, keyVersionFlag)
+
+		if err != nil {
+			return err
+		}
+
+		fmt.Printf("successfully created AES-%d key: %s\n", aesSizeFlag, encryptionKeyPath)
+		return nil
+	},
+}
+
+func init() {
+	cli.GenCmd.AddCommand(genEncryptionKeyCmd)
+
+	genEncryptionKeyCmd.PersistentFlags().IntVarP(&aesSizeFlag, "size", "s", 128,
+		"AES key size for encryption at rest (one of: 128, 192, 256)")
+	genEncryptionKeyCmd.PersistentFlags().BoolVar(&overwriteKeyFlag, "overwrite", false,
+		"Overwrite key if it exists")
+	genEncryptionKeyCmd.PersistentFlags().IntVar(&keyVersionFlag, "version", 1,
+		"Encryption format version (1 or 2)")
+}

pkg/ccl/cliccl/gen_test.go

@@ -0,0 +1,61 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package cliccl
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/pebble/vfs"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGenEncryptionKey(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	dir := t.TempDir()
+
+	for _, keyVersion := range []int{1, 2} {
+		for _, keySize := range []int{128, 192, 256} {
+			t.Run(fmt.Sprintf("version=%d/size=%d", keyVersion, keySize), func(t *testing.T) {
+				keyName := fmt.Sprintf("aes-%d-v%d.key", keySize, keyVersion)
+				keyPath := filepath.Join(dir, keyName)
+
+				err := genEncryptionKey(keyPath, keySize, false, keyVersion)
+				require.NoError(t, err)
+
+				if keyVersion == 1 {
+					info, err := os.Stat(keyPath)
+					require.NoError(t, err)
+					// 32-byte id plus the key.
+					assert.EqualValues(t, 32+(keySize/8), info.Size())
+				}
+
+				key, err := engineccl.LoadKeyFromFile(vfs.Default, keyPath)
+				require.NoError(t, err)
+				assert.EqualValues(t, keySize/8, len(key.Key))
+				// Key ID is hex encoded on load so it's 64 bytes here but 32 in the file size.
+				assert.EqualValues(t, 64, len(key.Info.KeyId))
+
+				err = genEncryptionKey(keyPath, keySize, false, keyVersion)
+				require.ErrorContains(t, err, fmt.Sprintf("%s: file exists", keyName))
+
+				err = genEncryptionKey(keyPath, keySize, true /* overwrite */, keyVersion)
+				require.NoError(t, err)
+			})
+		}
+	}
+}

pkg/ccl/storageccl/engineccl/BUILD.bazel

@@ -31,6 +31,8 @@ go_library(
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_pebble//vfs/atomicfs",
         "@com_github_gogo_protobuf//proto",
+        "@com_github_lestrrat_go_jwx//jwa",
+        "@com_github_lestrrat_go_jwx//jwk",
     ],
 )
 

pkg/ccl/storageccl/engineccl/ctr_stream.go

@@ -51,27 +51,54 @@ func (c *FileCipherStreamCreator) CreateNew(
 	settings := &enginepbccl.EncryptionSettings{}
 	if key == nil || key.Info.EncryptionType == enginepbccl.EncryptionType_Plaintext {
 		settings.EncryptionType = enginepbccl.EncryptionType_Plaintext
-		stream := &filePlainStream{}
-		return settings, stream, nil
+	} else {
+		settings.EncryptionType = key.Info.EncryptionType
+		settings.KeyId = key.Info.KeyId
+		settings.Nonce = make([]byte, ctrNonceSize)
+		_, err = rand.Read(settings.Nonce)
+		if err != nil {
+			return nil, nil, err
+		}
+		counterBytes := make([]byte, 4)
+		if _, err = rand.Read(counterBytes); err != nil {
+			return nil, nil, err
+		}
+		// Does not matter how we convert 4 random bytes into uint32
+		settings.Counter = binary.LittleEndian.Uint32(counterBytes)
 	}
-	settings.EncryptionType = key.Info.EncryptionType
-	settings.KeyId = key.Info.KeyId
-	settings.Nonce = make([]byte, ctrNonceSize)
-	_, err = rand.Read(settings.Nonce)
+
+	fcs, err := createFileCipherStream(settings, key)
 	if err != nil {
 		return nil, nil, err
 	}
-	counterBytes := make([]byte, 4)
-	if _, err = rand.Read(counterBytes); err != nil {
-		return nil, nil, err
-	}
-	// Does not matter how we convert 4 random bytes into uint32
-	settings.Counter = binary.LittleEndian.Uint32(counterBytes)
-	ctrCS, err := newCTRBlockCipherStream(key, settings.Nonce, settings.Counter)
-	if err != nil {
-		return nil, nil, err
+	return settings, fcs, nil
+}
+
+func createFileCipherStream(
+	settings *enginepbccl.EncryptionSettings, key *enginepbccl.SecretKey,
+) (FileStream, error) {
+	switch settings.EncryptionType {
+	case enginepbccl.EncryptionType_Plaintext:
+		return &filePlainStream{}, nil
+
+	case enginepbccl.EncryptionType_AES128_CTR, enginepbccl.EncryptionType_AES192_CTR, enginepbccl.EncryptionType_AES256_CTR:
+		ctrCS, err := newCTRBlockCipherStream(key, settings.Nonce, settings.Counter)
+		if err != nil {
+			return nil, err
+		}
+		return &fileCipherStream{bcs: ctrCS}, nil
+
+	case enginepbccl.EncryptionType_AES_128_CTR_V2, enginepbccl.EncryptionType_AES_192_CTR_V2, enginepbccl.EncryptionType_AES_256_CTR_V2:
+		var iv [ctrBlockSize]byte
+		copy(iv[:ctrNonceSize], settings.Nonce)
+		binary.BigEndian.PutUint32(iv[ctrNonceSize:ctrNonceSize+4], settings.Counter)
+		fcs, err := newFileCipherStreamV2(key.Key, iv[:])
+		if err != nil {
+			return nil, err
+		}
+		return fcs, nil
 	}
-	return settings, &fileCipherStream{bcs: ctrCS}, nil
+	return nil, fmt.Errorf("unknown encryption type %s", settings.EncryptionType)
 }
 
 // CreateExisting creates a FileStream for an existing file by looking up the key described by
@@ -86,11 +113,7 @@ func (c *FileCipherStreamCreator) CreateExisting(
 	if err != nil {
 		return nil, err
 	}
-	ctrCS, err := newCTRBlockCipherStream(key, settings.Nonce, settings.Counter)
-	if err != nil {
-		return nil, err
-	}
-	return &fileCipherStream{bcs: ctrCS}, nil
+	return createFileCipherStream(settings, key)
 }
 
 // FileStream encrypts/decrypts byte slices at arbitrary file offsets.

pkg/ccl/storageccl/engineccl/ctr_stream_test.go

@@ -155,8 +155,8 @@ func TestCTRStreamDataDriven(t *testing.T) {
 						d.MaybeScanArgs(t, "key", &keyName)
 						ivName := "default"
 						d.MaybeScanArgs(t, "iv", &ivName)
-						skipV1 := false
-						d.MaybeScanArgs(t, "skip-v1", &skipV1)
+						var onlyVersion string
+						d.MaybeScanArgs(t, "only-version", &onlyVersion)
 						iv := ivs[ivName]
 						var fcs FileStream
 						if impl == "v1" {
@@ -178,7 +178,7 @@ func TestCTRStreamDataDriven(t *testing.T) {
 						}
 
 						outputString := string(output)
-						if skipV1 && impl == "v1" {
+						if onlyVersion != "" && impl != onlyVersion {
 							return d.Expected
 						}
 						_, isDuplicate := seenCiphertexts[outputString]
@@ -320,9 +320,9 @@ func TestFileCipherStreamCreator(t *testing.T) {
 	// Make the active key = nil, so encryption/decryption is a noop.
 	km.activeID = "bar"
 	encSettings, fs5, err := fcs.CreateNew(context.Background())
+	require.NoError(t, err)
 	require.Equal(t, "", encSettings.KeyId)
 	require.Equal(t, enginepbccl.EncryptionType_Plaintext, encSettings.EncryptionType)
-	require.NoError(t, err)
 	fs5.Encrypt(5, data)
 	if diff := pretty.Diff(data, testData); diff != nil {
 		t.Fatalf("%s\n%s", strings.Join(diff, "\n"), data)

pkg/ccl/storageccl/engineccl/enginepbccl/BUILD.bazel

@@ -22,6 +22,7 @@ go_proto_library(
 
 go_library(
     name = "enginepbccl",
+    srcs = ["key_registry.go"],
     embed = [":enginepbccl_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl/enginepbccl",
     visibility = ["//visibility:public"],