diff --git a/go.mod b/go.mod index 8575613..468f14d 100644 --- a/go.mod +++ b/go.mod @@ -91,6 +91,7 @@ require ( github.com/jinzhu/now v1.1.5 // indirect github.com/joho/godotenv v1.5.1 // indirect github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/compress v1.19.0 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e // indirect diff --git a/go.sum b/go.sum index 61d800c..63d5655 100644 --- a/go.sum +++ b/go.sum @@ -286,6 +286,8 @@ github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= +github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= diff --git a/internal/utils/compress.go b/internal/utils/compress.go index 88e1917..832a50a 100644 --- a/internal/utils/compress.go +++ b/internal/utils/compress.go @@ -5,9 +5,14 @@ import ( "compress/zlib" "encoding/base64" "io" + "strings" "sync" + + "github.com/klauspost/compress/zstd" ) +const zstdPrefix = "zstd:" + var zlibWriterPool = sync.Pool{ New: func() interface{} { return zlib.NewWriter(io.Discard) @@ -26,30 +31,59 @@ func PutZlibWriter(zw *zlib.Writer) { zlibWriterPool.Put(zw) } -// CompressToBase64 compresses data using zlib and encodes to base64 +var ( + zstdEncoder *zstd.Encoder + zstdDecoder *zstd.Decoder + zstdOnce sync.Once +) + +func initZstd() { + zstdOnce.Do(func() { + var err error + // 默认级别适合常规压缩 + zstdEncoder, err = zstd.NewWriter(nil) + if err != nil { + panic(err) + } + zstdDecoder, err = zstd.NewReader(nil) + if err != nil { + panic(err) + } + }) +} + +// CompressToBase64 compresses data using zstd and encodes to base64 with a prefix func CompressToBase64(data string) (string, error) { if data == "" { return "", nil } - var buf bytes.Buffer - zw := GetZlibWriter(&buf) - defer PutZlibWriter(zw) + initZstd() - if _, err := zw.Write([]byte(data)); err != nil { - return "", err - } - if err := zw.Close(); err != nil { - return "", err - } - - return base64.StdEncoding.EncodeToString(buf.Bytes()), nil + compressed := zstdEncoder.EncodeAll([]byte(data), nil) + return zstdPrefix + base64.StdEncoding.EncodeToString(compressed), nil } -// DecompressFromBase64 decodes base64 and decompresses zlib data +// DecompressFromBase64 decodes base64 and decompresses data (supports zstd prefix and falls back to zlib) func DecompressFromBase64(data string) (string, error) { if data == "" { return "", nil } + + if strings.HasPrefix(data, zstdPrefix) { + initZstd() + encoded := data[len(zstdPrefix):] + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", err + } + decompressed, err := zstdDecoder.DecodeAll(decoded, nil) + if err != nil { + return "", err + } + return string(decompressed), nil + } + + // Fallback to legacy zlib decoded, err := base64.StdEncoding.DecodeString(data) if err != nil { return "", err @@ -65,3 +99,4 @@ func DecompressFromBase64(data string) (string, error) { } return string(result), nil } + diff --git a/internal/utils/compress_test.go b/internal/utils/compress_test.go new file mode 100644 index 0000000..8c73746 --- /dev/null +++ b/internal/utils/compress_test.go @@ -0,0 +1,93 @@ +package utils + +import ( + "bytes" + "compress/zlib" + "encoding/base64" + "strings" + "testing" +) + +// legacyZlibCompress 以前的 zlib 压缩逻辑,用于构造测试样本 +func legacyZlibCompress(data string) (string, error) { + if data == "" { + return "", nil + } + var buf bytes.Buffer + zw := zlib.NewWriter(&buf) + if _, err := zw.Write([]byte(data)); err != nil { + return "", err + } + if err := zw.Close(); err != nil { + return "", err + } + return base64.StdEncoding.EncodeToString(buf.Bytes()), nil +} + +func TestCompressAndDecompressZstd(t *testing.T) { + originalText := "Hello, this is a test log message for ZSTD compression in Baihu Panel! Repeat: Hello, this is a test log message for ZSTD compression in Baihu Panel!" + + // 1. 测试 ZSTD 压缩 + compressed, err := CompressToBase64(originalText) + if err != nil { + t.Fatalf("CompressToBase64 failed: %v", err) + } + + // 验证前缀是否正确 + if !strings.HasPrefix(compressed, "zstd:") { + t.Errorf("Expected compressed output to have 'zstd:' prefix, got: %s", compressed) + } + + // 2. 测试 ZSTD 解密解压 + decompressed, err := DecompressFromBase64(compressed) + if err != nil { + t.Fatalf("DecompressFromBase64 failed: %v", err) + } + + if decompressed != originalText { + t.Errorf("Decompressed text mismatch.\nExpected: %s\nGot: %s", originalText, decompressed) + } +} + +func TestDecompressLegacyZlibCompatibility(t *testing.T) { + originalText := "This is a legacy log message compressed using zlib. It should be decompressed successfully." + + // 1. 用老逻辑压缩生成旧数据 + legacyCompressed, err := legacyZlibCompress(originalText) + if err != nil { + t.Fatalf("legacyZlibCompress failed: %v", err) + } + + // 验证没有 zstd 前缀 + if strings.HasPrefix(legacyCompressed, "zstd:") { + t.Fatalf("Legacy compressed string shouldn't have 'zstd:' prefix") + } + + // 2. 使用新版的 DecompressFromBase64 解压,验证其对旧格式的兼容性 + decompressed, err := DecompressFromBase64(legacyCompressed) + if err != nil { + t.Fatalf("DecompressFromBase64 failed to decompress legacy data: %v", err) + } + + if decompressed != originalText { + t.Errorf("Decompressed legacy text mismatch.\nExpected: %s\nGot: %s", originalText, decompressed) + } +} + +func TestEmptyString(t *testing.T) { + compressed, err := CompressToBase64("") + if err != nil { + t.Fatalf("CompressToBase64 for empty string failed: %v", err) + } + if compressed != "" { + t.Errorf("Expected empty string for empty input, got: %q", compressed) + } + + decompressed, err := DecompressFromBase64("") + if err != nil { + t.Fatalf("DecompressFromBase64 for empty string failed: %v", err) + } + if decompressed != "" { + t.Errorf("Expected empty string for empty input, got: %q", decompressed) + } +}