sqlreplay: report an error immediately when the encryption file key is not set (#751)

Author: djshow832

lib/go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/tiancaiamao/gp v0.0.0-20230126082955-4f9e4f1ed9b5
 	go.uber.org/atomic v1.9.0
 	go.uber.org/zap v1.23.0
+	golang.org/x/term v0.29.0
 )
 
 require (
@@ -17,7 +18,6 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/sys v0.30.0 // indirect
-	golang.org/x/term v0.29.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
@@ -26,7 +26,6 @@ require (
 	github.com/kr/pretty v0.3.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
-	golang.org/x/crypto v0.33.0
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 )

lib/go.sum

@@ -50,8 +50,6 @@ go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8=
 go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/zap v1.23.0 h1:OjGQ5KQDEUawVHxNwQgPpiypGHOxo2mNZsOqTak4fFY=
 go.uber.org/zap v1.23.0/go.mod h1:D+nX8jyLsMHMYrln8A0rJjFt/T/9/bGgIhAqxv5URuY=
-golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
-golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
 golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
 golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU=

pkg/sqlreplay/capture/capture.go

@@ -61,6 +61,7 @@ type CaptureConfig struct {
 	StartTime          time.Time
 	Duration           time.Duration
 	Compress           bool
+	encryptionKey      []byte
 	cmdLogger          io.WriteCloser
 	bufferCap          int
 	flushThreshold     int
@@ -91,6 +92,11 @@ func (cfg *CaptureConfig) Validate() (storage.ExternalStorage, error) {
 	} else if cfg.StartTime.Add(cfg.Duration).Before(now) {
 		return storage, errors.New("start time should not be in the past")
 	}
+	key, err := store.LoadEncryptionKey(cfg.EncryptionMethod, cfg.KeyFile)
+	if err != nil {
+		return storage, errors.Wrapf(err, "failed to load encryption key")
+	}
+	cfg.encryptionKey = key
 	if cfg.bufferCap == 0 {
 		cfg.bufferCap = bufferCap
 	}
@@ -244,10 +250,10 @@ func (c *capture) flushBuffer(bufCh <-chan *bytes.Buffer) {
 	if cmdLogger == nil {
 		var err error
 		cmdLogger, err = store.NewWriter(c.lg.Named("writer"), c.storage, store.WriterCfg{
-			Dir:           c.cfg.Output,
-			EncryptMethod: c.cfg.EncryptionMethod,
-			KeyFile:       c.cfg.KeyFile,
-			Compress:      c.cfg.Compress,
+			Dir:              c.cfg.Output,
+			EncryptionMethod: c.cfg.EncryptionMethod,
+			EncryptionKey:    c.cfg.encryptionKey,
+			Compress:         c.cfg.Compress,
 		})
 		if err != nil {
 			c.lg.Error("failed to create capture writer", zap.Error(err))

pkg/sqlreplay/capture/capture_test.go

@@ -160,6 +160,13 @@ func TestCaptureCfgError(t *testing.T) {
 			Output:    path,
 			StartTime: now,
 		},
+		{
+			Duration:         10 * time.Second,
+			Output:           dir,
+			StartTime:        now,
+			EncryptionMethod: store.EncryptAes,
+			KeyFile:          "",
+		},
 	}
 
 	for i, cfg := range cfgs {

pkg/sqlreplay/replay/replay.go

@@ -56,9 +56,10 @@ type ReplayConfig struct {
 	KeyFile  string
 	// It's specified when executing with the statement `TRAFFIC REPLAY` so that all TiProxy instances
 	// use the same start time and the time acts as the job ID.
-	StartTime time.Time
-	Speed     float64
-	ReadOnly  bool
+	StartTime     time.Time
+	Speed         float64
+	ReadOnly      bool
+	encryptionKey []byte
 	// the following fields are for testing
 	reader            cmd.LineReader
 	report            report.Report
@@ -154,6 +155,12 @@ func (r *replay) Start(cfg ReplayConfig, backendTLSConfig *tls.Config, hsHandler
 	r.replayStats.Reset()
 	r.exceptionCh = make(chan conn.Exception, maxPendingExceptions)
 	r.closeCh = make(chan uint64, maxPendingExceptions)
+	key, err := store.LoadEncryptionKey(r.meta.EncryptMethod, cfg.KeyFile)
+	if err != nil {
+		return errors.Wrapf(err, "failed to load encryption key")
+	}
+	cfg.encryptionKey = key
+
 	hsHandler = NewHandshakeHandler(hsHandler)
 	r.connCreator = cfg.connCreator
 	if r.connCreator == nil {
@@ -190,9 +197,9 @@ func (r *replay) readCommands(ctx context.Context) {
 	if reader == nil {
 		var err error
 		reader, err = store.NewReader(r.lg.Named("loader"), r.storage, store.ReaderCfg{
-			Dir:           r.cfg.Input,
-			KeyFile:       r.cfg.KeyFile,
-			EncryptMethod: r.meta.EncryptMethod,
+			Dir:              r.cfg.Input,
+			EncryptionKey:    r.cfg.encryptionKey,
+			EncryptionMethod: r.meta.EncryptMethod,
 		})
 		if err != nil {
 			r.stop(err)

pkg/sqlreplay/replay/replay_test.go

@@ -273,3 +273,22 @@ func TestPendingCmds(t *testing.T) {
 	require.Contains(t, logs, `"total_wait_time": "150ms"`)
 	require.Contains(t, logs, "too many pending commands")
 }
+
+func TestStartError(t *testing.T) {
+	replay := NewReplay(zap.NewNop(), id.NewIDManager())
+	dir := t.TempDir()
+	meta := store.NewMeta(10*time.Second, 20, 0, store.EncryptAes)
+	storage, err := store.NewStorage(dir)
+	require.NoError(t, err)
+	defer storage.Close()
+	require.NoError(t, meta.Write(storage))
+	now := time.Now()
+
+	cfg := ReplayConfig{
+		Input:     dir,
+		Username:  "u1",
+		StartTime: now,
+	}
+	err = replay.Start(cfg, nil, nil, &backend.BCConfig{})
+	require.ErrorContains(t, err, "encryption")
+}

pkg/sqlreplay/store/encrypt.go

@@ -27,20 +27,15 @@ type aesCTRWriter struct {
 	stream cipher.Stream
 }
 
-func newWriterWithEncryptOpts(writer io.WriteCloser, encryptMethod string, keyFile string) (io.WriteCloser, error) {
-	switch strings.ToLower(encryptMethod) {
+func newWriterWithEncryptOpts(writer io.WriteCloser, encryptionMethod string, encrytionKey []byte) (io.WriteCloser, error) {
+	switch strings.ToLower(encryptionMethod) {
 	case "", EncryptPlain:
 		return writer, nil
 	case EncryptAes:
+		return newAESCTRWriter(writer, encrytionKey)
 	default:
-		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptMethod)
+		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptionMethod)
 	}
-
-	key, err := readAesKey(keyFile)
-	if err != nil {
-		return nil, err
-	}
-	return newAESCTRWriter(writer, key)
 }
 
 func newAESCTRWriter(writer io.WriteCloser, key []byte) (*aesCTRWriter, error) {
@@ -72,20 +67,15 @@ type aesCTRReader struct {
 	stream cipher.Stream
 }
 
-func newReaderWithEncryptOpts(reader io.Reader, encryptMethod string, keyFile string) (io.Reader, error) {
-	switch strings.ToLower(encryptMethod) {
+func newReaderWithEncryptOpts(reader io.Reader, encryptionMethod string, encryptionKey []byte) (io.Reader, error) {
+	switch strings.ToLower(encryptionMethod) {
 	case "", EncryptPlain:
 		return reader, nil
 	case EncryptAes:
+		return newAESCTRReader(reader, encryptionKey)
 	default:
-		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptMethod)
-	}
-
-	key, err := readAesKey(keyFile)
-	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptionMethod)
 	}
-	return newAESCTRReader(reader, key)
 }
 
 func newAESCTRReader(reader io.Reader, key []byte) (*aesCTRReader, error) {
@@ -118,16 +108,28 @@ func (ctr *aesCTRReader) Read(data []byte) (int, error) {
 	return n, errors.WithStack(err)
 }
 
+func LoadEncryptionKey(encryptionMethod, keyFile string) ([]byte, error) {
+	switch strings.ToLower(encryptionMethod) {
+	case "", EncryptPlain:
+		return nil, nil
+	case EncryptAes:
+		return readAesKey(keyFile)
+	default:
+		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptionMethod)
+	}
+}
+
 func readAesKey(filename string) ([]byte, error) {
 	if len(filename) == 0 {
-		return nil, errors.New("encryption key file name is not set")
+		return nil, errors.New("security.encryption-key-file is not set")
 	}
 	key, err := os.ReadFile(filename)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
-	if len(key) != 32 {
+	if len(key) < 32 {
 		return nil, errors.Errorf("invalid aes-256 key length: %d, expecting 32", len(key))
 	}
-	return key, nil
+	// in case it's ended with a new line
+	return key[:32], nil
 }

pkg/sqlreplay/store/encrypt_test.go

@@ -50,22 +50,21 @@ func TestAes256(t *testing.T) {
 func TestEncryptOpts(t *testing.T) {
 	dir := t.TempDir()
 	path := filepath.Join(dir, "test")
-	keyFile := filepath.Join(dir, "key")
-	require.NoError(t, os.WriteFile(keyFile, genAesKey(), 0600))
+	aesKey := genAesKey()
 
 	tests := []struct {
-		method  string
-		keyFile string
+		method string
+		key    []byte
 	}{
-		{EncryptPlain, ""},
-		{"", ""},
-		{EncryptAes, keyFile},
+		{EncryptPlain, nil},
+		{"", nil},
+		{EncryptAes, aesKey},
 	}
 	for i, test := range tests {
 		// write
 		file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
 		require.NoError(t, err)
-		writer, err := newWriterWithEncryptOpts(file, test.method, test.keyFile)
+		writer, err := newWriterWithEncryptOpts(file, test.method, test.key)
 		require.NoError(t, err, "case %d", i)
 		n, err := writer.Write([]byte("test"))
 		require.NoError(t, err)
@@ -75,7 +74,7 @@ func TestEncryptOpts(t *testing.T) {
 		// read
 		file, err = os.OpenFile(path, os.O_RDONLY, 0600)
 		require.NoError(t, err)
-		reader, err := newReaderWithEncryptOpts(file, test.method, test.keyFile)
+		reader, err := newReaderWithEncryptOpts(file, test.method, test.key)
 		require.NoError(t, err)
 		data := make([]byte, 100)
 		n, err = io.ReadFull(reader, data)
@@ -85,32 +84,40 @@ func TestEncryptOpts(t *testing.T) {
 	}
 }
 
-func TestAes256Error(t *testing.T) {
+func TestLoadKey(t *testing.T) {
 	dir := t.TempDir()
-	path := filepath.Join(dir, "test")
-	file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
-	require.NoError(t, err)
-	defer file.Close()
+	key := genAesKey()
+
 	keyFile := filepath.Join(dir, "key")
-	require.NoError(t, os.WriteFile(keyFile, genAesKey(), 0600))
+	require.NoError(t, os.WriteFile(keyFile, key, 0600))
 	invalidKeyFile := filepath.Join(dir, "invalid")
 	require.NoError(t, os.WriteFile(invalidKeyFile, []byte("invalid"), 0600))
 	noKeyFile := filepath.Join(dir, "nonexist")
+	longKeyFile := filepath.Join(dir, "valid")
+	longKey := make([]byte, 33)
+	copy(longKey, key)
+	longKey[32] = '\n'
+	require.NoError(t, os.WriteFile(longKeyFile, longKey, 0600))
 
 	tests := []struct {
 		method  string
 		keyFile string
+		err     bool
 	}{
-		{"unknown", keyFile},
-		{EncryptAes, ""},
-		{EncryptAes, noKeyFile},
-		{EncryptAes, invalidKeyFile},
+		{"unknown", keyFile, true},
+		{EncryptAes, "", true},
+		{EncryptAes, noKeyFile, true},
+		{EncryptAes, invalidKeyFile, true},
+		{EncryptAes, keyFile, false},
+		{EncryptAes, longKeyFile, false},
 	}
 	for i, test := range tests {
-		_, err = newWriterWithEncryptOpts(file, test.method, test.keyFile)
-		require.Error(t, err, "case %d", i)
-		_, err = newReaderWithEncryptOpts(file, test.method, test.keyFile)
-		require.Error(t, err, "case %d", i)
+		actual, err := LoadEncryptionKey(test.method, test.keyFile)
+		if test.err {
+			require.Error(t, err, "case %d", i)
+		} else {
+			require.Equal(t, key, actual, "case %d", i)
+		}
 	}
 }
 

pkg/sqlreplay/store/line.go

@@ -14,11 +14,11 @@ import (
 )
 
 type WriterCfg struct {
-	Dir           string
-	EncryptMethod string
-	KeyFile       string
-	FileSize      int
-	Compress      bool
+	Dir              string
+	EncryptionMethod string
+	EncryptionKey    []byte
+	FileSize         int
+	Compress         bool
 }
 
 // NewWriter just wraps the rotate writer. It doesn't use a buffer because Capture writes data in a big batch.
@@ -28,9 +28,9 @@ func NewWriter(lg *zap.Logger, externalStorage storage.ExternalStorage, cfg Writ
 }
 
 type ReaderCfg struct {
-	Dir           string
-	EncryptMethod string
-	KeyFile       string
+	Dir              string
+	EncryptionMethod string
+	EncryptionKey    []byte
 }
 
 var _ cmd.LineReader = (*loader)(nil)

pkg/sqlreplay/store/rotate.go

@@ -73,7 +73,7 @@ func (w *rotateWriter) createFile() error {
 	if w.cfg.Compress {
 		w.writer = newCompressWriter(w.lg, w.writer)
 	}
-	if w.writer, err = newWriterWithEncryptOpts(w.writer, w.cfg.EncryptMethod, w.cfg.KeyFile); err != nil {
+	if w.writer, err = newWriterWithEncryptOpts(w.writer, w.cfg.EncryptionMethod, w.cfg.EncryptionKey); err != nil {
 		return err
 	}
 	return nil
@@ -207,7 +207,7 @@ func (r *rotateReader) nextReader() error {
 			return err
 		}
 	}
-	r.reader, err = newReaderWithEncryptOpts(r.reader, r.cfg.EncryptMethod, r.cfg.KeyFile)
+	r.reader, err = newReaderWithEncryptOpts(r.reader, r.cfg.EncryptionMethod, r.cfg.EncryptionKey)
 	if err != nil {
 		return err
 	}