ear_test.go

// 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 (
	"bytes"
	"context"
	"crypto/rand"
	"fmt"
	"path/filepath"
	"sort"
	"strings"
	"testing"

	"github.com/cockroachdb/cockroach/pkg/ccl/baseccl"
	// The following import is required for the hook that populates
	// NewEncryptedEnvFunc in `pkg/storage`.
	_ "github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl"
	"github.com/cockroachdb/cockroach/pkg/cli"
	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
	"github.com/cockroachdb/cockroach/pkg/storage"
	"github.com/cockroachdb/cockroach/pkg/storage/fs"
	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
	"github.com/cockroachdb/cockroach/pkg/util/envutil"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/randutil"
	"github.com/cockroachdb/datadriven"
	"github.com/cockroachdb/pebble/vfs"
	"github.com/spf13/cobra"
	"github.com/stretchr/testify/require"
)

func TestDecrypt(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	ctx := context.Background()
	dir := t.TempDir()

	// Generate a new encryption key to use.
	keyPath := filepath.Join(dir, "aes.key")
	err := genEncryptionKeyCmd.RunE(nil, []string{keyPath})
	require.NoError(t, err)

	// Spin up a new encrypted store.
	encSpecStr := fmt.Sprintf("path=%s,key=%s,old-key=plain", dir, keyPath)
	encSpec, err := baseccl.NewStoreEncryptionSpec(encSpecStr)
	require.NoError(t, err)
	encOpts, err := encSpec.ToEncryptionOptions()
	require.NoError(t, err)

	env, err := fs.InitEnv(ctx, vfs.Default, dir, fs.EnvConfig{EncryptionOptions: encOpts})
	require.NoError(t, err)
	p, err := storage.Open(ctx, env, cluster.MakeClusterSettings())
	require.NoError(t, err)

	// Find a manifest file to check.
	files, err := p.List(dir)
	require.NoError(t, err)
	sort.Strings(files)
	var manifestPath string
	for _, basename := range files {
		if strings.HasPrefix(basename, "MANIFEST-") {
			manifestPath = filepath.Join(dir, basename)
			break
		}
	}
	// Should have found a manifest file.
	require.NotEmpty(t, manifestPath)

	// Close the DB.
	p.Close()

	// Pluck the `pebble manifest dump` command out of the debug command.
	dumpCmd := getTool(cli.DebugPebbleCmd, []string{"pebble", "manifest", "dump"})
	require.NotNil(t, dumpCmd)

	dumpManifest := func(cmd *cobra.Command, path string) string {
		var b bytes.Buffer
		dumpCmd.SetOut(&b)
		dumpCmd.SetErr(&b)
		dumpCmd.Run(cmd, []string{path})
		return b.String()
	}
	out := dumpManifest(dumpCmd, manifestPath)
	// Check for the presence of the comparator line in the manifest dump, as a
	// litmus test for the manifest file being readable. This line should only
	// appear once the file has been decrypted.
	const checkStr = "comparer:     cockroach_comparator"
	require.NotContains(t, out, checkStr)

	// Decrypt the manifest file.
	outPath := filepath.Join(dir, "manifest.plain")
	decryptCmd := getTool(cli.DebugCmd, []string{"debug", "encryption-decrypt"})
	require.NotNil(t, decryptCmd)
	err = decryptCmd.Flags().Set("enterprise-encryption", encSpecStr)
	require.NoError(t, err)
	err = decryptCmd.RunE(decryptCmd, []string{dir, manifestPath, outPath})
	require.NoError(t, err)

	// Check that the decrypted manifest file can now be read.
	out = dumpManifest(dumpCmd, outPath)
	require.Contains(t, out, checkStr)
}

func TestList(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	// Pin the random generator to use a fixed seed. This also requires overriding
	// the PRNG for the duration of the test. This ensures that all the keys
	// generated by the key manager are deterministic.
	reset := envutil.TestSetEnv(t, "COCKROACH_RANDOM_SEED", "1665612120123601000")
	defer reset()
	randBefore := rand.Reader
	randOverride, _ := randutil.NewPseudoRand()
	rand.Reader = randOverride
	defer func() { rand.Reader = randBefore }()

	ctx := context.Background()
	dir := t.TempDir()

	// Generate a new encryption key to use.
	keyPath := filepath.Join(dir, "aes.key")
	err := genEncryptionKeyCmd.RunE(nil, []string{keyPath})
	require.NoError(t, err)

	// Spin up a new encrypted store.
	encSpecStr := fmt.Sprintf("path=%s,key=%s,old-key=plain", dir, keyPath)
	encSpec, err := baseccl.NewStoreEncryptionSpec(encSpecStr)
	require.NoError(t, err)
	encOpts, err := encSpec.ToEncryptionOptions()
	require.NoError(t, err)
	env, err := fs.InitEnv(ctx, vfs.Default, dir, fs.EnvConfig{EncryptionOptions: encOpts})
	require.NoError(t, err)
	p, err := storage.Open(ctx, env, cluster.MakeClusterSettings())
	require.NoError(t, err)

	// Write a key and flush, to create a table in the store.
	err = p.PutUnversioned([]byte("foo"), nil)
	require.NoError(t, err)
	err = p.Flush()
	require.NoError(t, err)
	p.Close()

	datadriven.RunTest(t, datapathutils.TestDataPath(t, "ear-list"), func(t *testing.T, d *datadriven.TestData) string {
		if d.Cmd != "list" {
			d.Fatalf(t, "invalid command %q", d.Cmd)
		}
		// List the files in the registry.
		cmd := getTool(cli.DebugCmd, []string{"debug", "encryption-registry-list"})
		require.NotNil(t, cmd)
		var b bytes.Buffer
		cmd.SetOut(&b)
		cmd.SetErr(&b)
		require.NoError(t, cmd.Flags().Set("enterprise-encryption", encSpecStr))
		require.NoError(t, runList(cmd, []string{dir}))
		return b.String()
	})
}

// getTool traverses the given cobra.Command recursively, searching for a tool
// matching the given command.
func getTool(cmd *cobra.Command, want []string) *cobra.Command {
	// Base cases.
	if cmd.Name() != want[0] {
		return nil
	}
	if len(want) == 1 {
		return cmd
	}
	// Recursive case.
	for _, subCmd := range cmd.Commands() {
		found := getTool(subCmd, want[1:])
		if found != nil {
			return found
		}
	}
	return nil
}