store: compress traffic files after encryption (#731)

Author: djshow832

pkg/sqlreplay/store/rotate.go

@@ -62,20 +62,20 @@ func (w *rotateWriter) createFile() error {
 	}
 	w.fileIdx++
 	fileName := fmt.Sprintf("%s%d%s%s", fileNamePrefix, w.fileIdx, fileNameSuffix, ext)
-	// compressWriter -> encryptWriter -> file
+	// rotateWriter -> encryptWriter -> compressWriter -> file
 	ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
 	fileWriter, err := w.storage.Create(ctx, fileName, &storage.WriterOption{})
 	cancel()
-	writer := NewStorageWriter(fileWriter)
+	w.writer = NewStorageWriter(fileWriter)
 	if err != nil {
 		return errors.WithStack(err)
 	}
-	if w.writer, err = newWriterWithEncryptOpts(writer, w.cfg.EncryptMethod, w.cfg.KeyFile); err != nil {
-		return err
-	}
 	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 {
+		return err
+	}
 	return nil
 }
 
@@ -199,18 +199,19 @@ func (r *rotateReader) nextReader() error {
 		return errors.WithStack(err)
 	}
 	r.externalFile = fileReader
+	r.reader = fileReader
 	r.curFileIdx = minFileIdx
 	r.curFileName = minFileName
-	// compressReader -> encryptReader -> file
-	r.reader, err = newReaderWithEncryptOpts(fileReader, r.cfg.EncryptMethod, r.cfg.KeyFile)
-	if err != nil {
-		return err
-	}
+	// rotateReader -> encryptReader -> compressReader -> file
 	if strings.HasSuffix(minFileName, fileCompressFormat) {
 		if r.reader, err = newCompressReader(r.reader); err != nil {
 			return err
 		}
 	}
+	r.reader, err = newReaderWithEncryptOpts(r.reader, r.cfg.EncryptMethod, r.cfg.KeyFile)
+	if err != nil {
+		return err
+	}
 	r.lg.Info("reading next file", zap.String("file", minFileName))
 	return nil
 }

pkg/sqlreplay/store/rotate_test.go

@@ -5,6 +5,7 @@ package store
 
 import (
 	"compress/gzip"
+	"context"
 	"errors"
 	"io"
 	"os"
@@ -257,3 +258,49 @@ func TestReadGZip(t *testing.T) {
 		l.Close()
 	}
 }
+
+func TestCompressAndEncrypt(t *testing.T) {
+	tmpDir := t.TempDir()
+	storage, err := NewStorage(tmpDir)
+	require.NoError(t, err)
+	defer storage.Close()
+	require.NoError(t, storage.WriteFile(context.Background(), "key", genAesKey()))
+	keyFile := filepath.Join(tmpDir, "key")
+
+	// write with compression and encryption
+	writer, err := newRotateWriter(zap.NewNop(), storage, WriterCfg{
+		Dir:           tmpDir,
+		FileSize:      1000,
+		Compress:      true,
+		EncryptMethod: EncryptAes,
+		KeyFile:       keyFile,
+	})
+	require.NoError(t, err)
+	_, err = writer.Write([]byte("test"))
+	require.NoError(t, err)
+	require.NoError(t, writer.Close())
+
+	// make sure data is compressed after encryption
+	file, err := os.Open(filepath.Join(tmpDir, "traffic-1.log.gz"))
+	require.NoError(t, err)
+	greader, err := gzip.NewReader(file)
+	require.NoError(t, err)
+	data := make([]byte, 1000)
+	n, err := io.ReadFull(greader, data)
+	require.ErrorContains(t, err, "EOF")
+	require.Equal(t, 20, n)
+	require.NoError(t, file.Close())
+
+	// rotateReader is able to read the file
+	reader, err := newRotateReader(zap.NewNop(), storage, ReaderCfg{
+		Dir:           tmpDir,
+		EncryptMethod: EncryptAes,
+		KeyFile:       keyFile,
+	})
+	require.NoError(t, err)
+	n, err = io.ReadFull(reader, data)
+	require.ErrorContains(t, err, "EOF")
+	require.Equal(t, 4, n)
+	require.Equal(t, []byte("test"), data[:4])
+	require.NoError(t, reader.Close())
+}