github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/storageccl/engineccl/encrypted_fs_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package engineccl
    10  
    11  import (
    12  	"bytes"
    13  	"context"
    14  	"io"
    15  	"os"
    16  	"strconv"
    17  	"strings"
    18  	"testing"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/base"
    21  	"github.com/cockroachdb/cockroach/pkg/ccl/baseccl"
    22  	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl/enginepbccl"
    23  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    24  	"github.com/cockroachdb/cockroach/pkg/storage"
    25  	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
    26  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    27  	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
    28  	"github.com/cockroachdb/pebble"
    29  	"github.com/cockroachdb/pebble/vfs"
    30  	"github.com/stretchr/testify/require"
    31  )
    32  
    33  func TestEncryptedFS(t *testing.T) {
    34  	defer leaktest.AfterTest(t)()
    35  
    36  	memFS := vfs.NewMem()
    37  
    38  	fileRegistry := &storage.PebbleFileRegistry{FS: memFS, DBDir: "/bar"}
    39  	require.NoError(t, fileRegistry.Load())
    40  
    41  	// Using a StoreKeyManager for the test since it is easy to create. Write a key for the
    42  	// StoreKeyManager.
    43  	var b []byte
    44  	for i := 0; i < keyIDLength+16; i++ {
    45  		b = append(b, 'a')
    46  	}
    47  	f, err := memFS.Create("keyfile")
    48  	require.NoError(t, err)
    49  	bReader := bytes.NewReader(b)
    50  	_, err = io.Copy(f, bReader)
    51  	require.NoError(t, err)
    52  	require.NoError(t, f.Close())
    53  	keyManager := &StoreKeyManager{fs: memFS, activeKeyFilename: "keyfile", oldKeyFilename: "plain"}
    54  	require.NoError(t, keyManager.Load(context.Background()))
    55  
    56  	streamCreator := &FileCipherStreamCreator{keyManager: keyManager, envType: enginepb.EnvType_Store}
    57  
    58  	fs := &encryptedFS{FS: memFS, fileRegistry: fileRegistry, streamCreator: streamCreator}
    59  
    60  	// Style (and most code) is from Pebble's mem_fs_test.go. We are mainly testing the integration of
    61  	// encryptedFS with FileRegistry and FileCipherStreamCreator. This uses real encryption but the
    62  	// strings here are not very long since we've tested that in lower-level unit tests.
    63  	testCases := []string{
    64  		// Make the /bar/baz directory; create a third-level file.
    65  		"1a: mkdirall /bar/baz",
    66  		"1b: f = create /bar/baz/y",
    67  		"1c: f.stat.name == y",
    68  		// Write more than a block of data; read it back.
    69  		"2a: f.write abcdefghijklmnopqrstuvwxyz",
    70  		"2b: f.close",
    71  		"2c: f = open /bar/baz/y",
    72  		"2d: f.read 5 == abcde",
    73  		"2e: f.readat 2 1 == bc",
    74  		"2f: f.readat 5 20 == uvwxy",
    75  		"2g: f.close",
    76  		// Link /bar/baz/y to /bar/z. We should be able to read from both files
    77  		// and remove them independently.
    78  		"3a: link /bar/baz/y /bar/z",
    79  		"3b: f = open /bar/z",
    80  		"3c: f.read 5 == abcde",
    81  		"3d: f.close",
    82  		"3e: remove /bar/baz/y",
    83  		"3f: f = open /bar/z",
    84  		"3g: f.read 5 == abcde",
    85  		"3h: f.close",
    86  		// Rename /bar/z to /foo
    87  		"4a: rename /bar/z /foo",
    88  		"4b: f = open /foo",
    89  		"4c: f.readat 5 20 == uvwxy",
    90  		"4d: f.close",
    91  		"4e: open /bar/z fails",
    92  		// ReuseForWrite /foo /baz
    93  		"5a: f = reuseForWrite /foo /baz",
    94  		"5b: f.write abc",
    95  		"5c: f.close",
    96  		"5d: f = open /baz",
    97  		"5e: f.read 3 == abc",
    98  	}
    99  
   100  	for _, tc := range testCases {
   101  		s := strings.Split(tc, " ")[1:]
   102  
   103  		saveF := s[0] == "f" && s[1] == "="
   104  		if saveF {
   105  			s = s[2:]
   106  		}
   107  
   108  		fails := s[len(s)-1] == "fails"
   109  		if fails {
   110  			s = s[:len(s)-1]
   111  		}
   112  
   113  		var (
   114  			fi  os.FileInfo
   115  			g   vfs.File
   116  			err error
   117  		)
   118  		switch s[0] {
   119  		case "create":
   120  			g, err = fs.Create(s[1])
   121  		case "link":
   122  			err = fs.Link(s[1], s[2])
   123  		case "open":
   124  			g, err = fs.Open(s[1])
   125  		case "mkdirall":
   126  			err = fs.MkdirAll(s[1], 0755)
   127  		case "remove":
   128  			err = fs.Remove(s[1])
   129  		case "rename":
   130  			err = fs.Rename(s[1], s[2])
   131  		case "reuseForWrite":
   132  			g, err = fs.ReuseForWrite(s[1], s[2])
   133  		case "f.write":
   134  			_, err = f.Write([]byte(s[1]))
   135  		case "f.read":
   136  			n, _ := strconv.Atoi(s[1])
   137  			buf := make([]byte, n)
   138  			_, err = io.ReadFull(f, buf)
   139  			if err != nil {
   140  				break
   141  			}
   142  			if got, want := string(buf), s[3]; got != want {
   143  				t.Fatalf("%q: got %q, want %q", tc, got, want)
   144  			}
   145  		case "f.readat":
   146  			n, _ := strconv.Atoi(s[1])
   147  			off, _ := strconv.Atoi(s[2])
   148  			buf := make([]byte, n)
   149  			_, err = f.ReadAt(buf, int64(off))
   150  			if err != nil {
   151  				break
   152  			}
   153  			if got, want := string(buf), s[4]; got != want {
   154  				t.Fatalf("%q: got %q, want %q", tc, got, want)
   155  			}
   156  		case "f.close":
   157  			f, err = nil, f.Close()
   158  		case "f.stat.name":
   159  			fi, err = f.Stat()
   160  			if err != nil {
   161  				break
   162  			}
   163  			if got, want := fi.Name(), s[2]; got != want {
   164  				t.Fatalf("%q: got %q, want %q", tc, got, want)
   165  			}
   166  		default:
   167  			t.Fatalf("bad test case: %q", tc)
   168  		}
   169  
   170  		if saveF {
   171  			f, g = g, nil
   172  		} else if g != nil {
   173  			g.Close()
   174  		}
   175  
   176  		if fails {
   177  			if err == nil {
   178  				t.Fatalf("%q: got nil error, want non-nil", tc)
   179  			}
   180  		} else {
   181  			if err != nil {
   182  				t.Fatalf("%q: %v", tc, err)
   183  			}
   184  		}
   185  	}
   186  }
   187  
   188  // Minimal test that creates an encrypted Pebble that exercises creation and reading of encrypted
   189  // files, rereading data after reopening the engine, and stats code.
   190  func TestPebbleEncryption(t *testing.T) {
   191  	defer leaktest.AfterTest(t)()
   192  
   193  	opts := storage.DefaultPebbleOptions()
   194  	opts.Cache = pebble.NewCache(1 << 20)
   195  	defer opts.Cache.Unref()
   196  
   197  	memFS := vfs.NewMem()
   198  	opts.FS = memFS
   199  	keyFile128 := "111111111111111111111111111111111234567890123456"
   200  	writeToFile(t, opts.FS, "16.key", []byte(keyFile128))
   201  	var encOptions baseccl.EncryptionOptions
   202  	encOptions.KeySource = baseccl.EncryptionKeySource_KeyFiles
   203  	encOptions.KeyFiles = &baseccl.EncryptionKeyFiles{
   204  		CurrentKey: "16.key",
   205  		OldKey:     "plain",
   206  	}
   207  	encOptions.DataKeyRotationPeriod = 1000 // arbitrary seconds
   208  	encOptionsBytes, err := protoutil.Marshal(&encOptions)
   209  	require.NoError(t, err)
   210  	db, err := storage.NewPebble(
   211  		context.Background(),
   212  		storage.PebbleConfig{
   213  			StorageConfig: base.StorageConfig{
   214  				Attrs:           roachpb.Attributes{},
   215  				MaxSize:         512 << 20,
   216  				UseFileRegistry: true,
   217  				ExtraOptions:    encOptionsBytes,
   218  			},
   219  			Opts: opts,
   220  		})
   221  	require.NoError(t, err)
   222  	// TODO(sbhola): Ensure that we are not returning the secret data keys by mistake.
   223  	r, err := db.GetEncryptionRegistries()
   224  	require.NoError(t, err)
   225  
   226  	var fileRegistry enginepb.FileRegistry
   227  	require.NoError(t, protoutil.Unmarshal(r.FileRegistry, &fileRegistry))
   228  	var keyRegistry enginepbccl.DataKeysRegistry
   229  	require.NoError(t, protoutil.Unmarshal(r.KeyRegistry, &keyRegistry))
   230  
   231  	stats, err := db.GetEnvStats()
   232  	require.NoError(t, err)
   233  	// Opening the DB should've created OPTIONS, CURRENT, MANIFEST and the
   234  	// WAL, all under the active key.
   235  	require.Equal(t, uint64(4), stats.TotalFiles)
   236  	require.Equal(t, uint64(4), stats.ActiveKeyFiles)
   237  	var s enginepbccl.EncryptionStatus
   238  	require.NoError(t, protoutil.Unmarshal(stats.EncryptionStatus, &s))
   239  	require.Equal(t, "16.key", s.ActiveStoreKey.Source)
   240  	require.Equal(t, int32(enginepbccl.EncryptionType_AES128_CTR), stats.EncryptionType)
   241  	t.Logf("EnvStats:\n%+v\n\n", *stats)
   242  
   243  	batch := db.NewWriteOnlyBatch()
   244  	require.NoError(t, batch.Put(storage.MVCCKey{Key: roachpb.Key("a")}, []byte("a")))
   245  	require.NoError(t, batch.Commit(true))
   246  	require.NoError(t, db.Flush())
   247  	val, err := db.Get(storage.MVCCKey{Key: roachpb.Key("a")})
   248  	require.NoError(t, err)
   249  	require.Equal(t, "a", string(val))
   250  	db.Close()
   251  
   252  	opts2 := storage.DefaultPebbleOptions()
   253  	opts2.Cache = pebble.NewCache(1 << 20)
   254  	defer opts2.Cache.Unref()
   255  
   256  	opts2.FS = memFS
   257  	db, err = storage.NewPebble(
   258  		context.Background(),
   259  		storage.PebbleConfig{
   260  			StorageConfig: base.StorageConfig{
   261  				Attrs:           roachpb.Attributes{},
   262  				MaxSize:         512 << 20,
   263  				UseFileRegistry: true,
   264  				ExtraOptions:    encOptionsBytes,
   265  			},
   266  			Opts: opts2,
   267  		})
   268  	require.NoError(t, err)
   269  	val, err = db.Get(storage.MVCCKey{Key: roachpb.Key("a")})
   270  	require.NoError(t, err)
   271  	require.Equal(t, "a", string(val))
   272  
   273  	// Flushing should've created a new sstable under the active key.
   274  	stats, err = db.GetEnvStats()
   275  	require.NoError(t, err)
   276  	require.Equal(t, uint64(5), stats.TotalFiles)
   277  	require.Equal(t, uint64(5), stats.ActiveKeyFiles)
   278  	require.Equal(t, stats.TotalBytes, stats.ActiveKeyBytes)
   279  	t.Logf("EnvStats:\n%+v\n\n", *stats)
   280  
   281  	db.Close()
   282  }