Merge pull request #736 from samwafgo/feat_zstd

feat:add zstd compress
2026-11-02 19:21:44 +08:00 · 2026-03-31 13:51:43 +08:00
parent da2ed94411 d67700d17e
commit 8ad22b7119
6 changed files with 65 additions and 4 deletions
@@ -113,7 +113,7 @@ require (
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.13-0.20220915233716-71ac16282d12 // indirect
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
-	github.com/klauspost/compress v1.17.8 // indirect
+	github.com/klauspost/compress v1.18.5 // 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-20211012122336-39d0f177ccd0 // indirect
@@ -270,6 +270,8 @@ github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
 github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
+github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE=
+github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
 github.com/klauspost/cpuid v1.2.0/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
 github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
@@ -47,7 +47,7 @@ type Hosts struct {
 	LogOnlyMode               int    `json:"log_only_mode"`                //仅记录模式 1开启 0关闭
 	CustomHeadersJSON         string `json:"custom_headers_json"`          //自定义头信息配置 json
 	CustomResponseHeadersJSON string `json:"custom_response_headers_json"` //自定义响应头信息配置 json
-	ResponseCompressJSON      string `json:"response_compress_json"`       //响应压缩配置 json（Gzip/Brotli）
+	ResponseCompressJSON      string `json:"response_compress_json"`       //响应压缩配置 json（Gzip/Brotli/Zstd）
 	IPMode                    string `json:"ip_mode"`                      //IP提取模式: "nic" 网卡模式 或 "proxy" 代理模式
 }

@@ -172,10 +172,10 @@ type StaticSiteConfig struct {
 	SecurityHeaders []StaticSecurityHeader `json:"security_headers"`
 }

-// ResponseCompressConfig 反代响应压缩（类似 nginx gzip / brotli）
+// ResponseCompressConfig 反代响应压缩（类似 nginx gzip / brotli / zstd）
 type ResponseCompressConfig struct {
 	IsEnable                 int    `json:"is_enable"`                   // 1 开启 0 关闭
-	Prefer                   string `json:"prefer"`                      // br_first | gzip_only | br_only
+	Prefer                   string `json:"prefer"`                      // br_first | gzip_only | br_only | zstd_only
 	MinLength                int    `json:"min_length"`                  // 最小压缩字节数，0 表示用默认
 	IncludeTypes             string `json:"include_types"`               // 分号分隔 MIME 前缀或完整类型，空则用内置默认
 	IncludeExtensions        string `json:"include_extensions"`          // 分号分隔，如 .js;.css，与类型取并集
@@ -5,6 +5,8 @@ import (
 	"compress/flate"
 	"compress/gzip"
 	"github.com/andybalholm/brotli"
+	"github.com/klauspost/compress/zstd"
+	"io"
 )

 func BrotliEncode(input []byte) ([]byte, error) {
@@ -108,3 +110,29 @@ func GZipDecode(input []byte) ([]byte, error) {
 	}
 	return buf.Bytes(), nil
 }
+
+func ZstdEncode(input []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	zstdWriter, err := zstd.NewWriter(&buf)
+	if err != nil {
+		return nil, err
+	}
+	_, err = zstdWriter.Write(input)
+	if err != nil {
+		_ = zstdWriter.Close()
+		return nil, err
+	}
+	if err := zstdWriter.Close(); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func ZstdDecode(input []byte) ([]byte, error) {
+	reader, err := zstd.NewReader(bytes.NewReader(input))
+	if err != nil {
+		return nil, err
+	}
+	defer reader.Close()
+	return io.ReadAll(reader)
+}
@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/andybalholm/brotli"
+	"github.com/klauspost/compress/zstd"
 	"golang.org/x/net/html/charset"
 	"golang.org/x/text/encoding"
 	"golang.org/x/text/encoding/simplifiedchinese"
@@ -72,6 +73,8 @@ func (waf *WafEngine) compressContent(res *http.Response, isStaticAssist bool, i
 		respBytes, err = utils.DeflateEncode(encodedBytes)
 	case "br":
 		respBytes, err = utils.BrotliEncode(encodedBytes)
+	case "zstd":
+		respBytes, err = utils.ZstdEncode(encodedBytes)
 	default:
 		respBytes = encodedBytes
 	}
@@ -112,6 +115,14 @@ func (waf *WafEngine) getOrgContent(resp *http.Response, isStaticAssist bool, de
 	case "br":
 		brotliReader := brotli.NewReader(resp.Body)
 		bodyReader = brotliReader
+	case "zstd":
+		zstdReader, zstdErr := zstd.NewReader(resp.Body)
+		if zstdErr != nil {
+			zlog.Warn("zstd解压失败: %v", zstdErr)
+			return bodyBytes, "", fmt.Errorf("zstd解压失败: %v", zstdErr)
+		}
+		bodyReader = zstdReader
+		defer zstdReader.Close()
 	default:
 		bodyReader = resp.Body
 	}
@@ -271,6 +282,14 @@ func (waf *WafEngine) decompressRequestContent(req *http.Request) ([]byte, error
 	case "br":
 		brotliReader := brotli.NewReader(bytes.NewReader(bodyBytes))
 		bodyReader = brotliReader
+	case "zstd":
+		zstdReader, zstdErr := zstd.NewReader(bytes.NewReader(bodyBytes))
+		if zstdErr != nil {
+			zlog.Warn("请求zstd解压失败: %v", zstdErr)
+			return bodyBytes, fmt.Errorf("请求zstd解压失败: %v", zstdErr)
+		}
+		bodyReader = zstdReader
+		defer zstdReader.Close()
 	default:
 		// 没有压缩或不支持的压缩格式，直接返回原始内容
 		return bodyBytes, nil
@@ -56,6 +56,8 @@ func (waf *WafEngine) maybeApplyResponseCompress(req *http.Request, resp *http.R
 		out, err = utils.BrotliEncode(body)
 	case "gzip":
 		out, err = utils.GZipEncode(body)
+	case "zstd":
+		out, err = utils.ZstdEncode(body)
 	default:
 		return body
 	}
@@ -139,6 +141,8 @@ func (waf *WafEngine) maybeCompressStaticAssistResponse(req *http.Request, resp
 		out, err = utils.BrotliEncode(raw)
 	case "gzip":
 		out, err = utils.GZipEncode(raw)
+	case "zstd":
+		out, err = utils.ZstdEncode(raw)
 	default:
 		resp.Body = io.NopCloser(bytes.NewReader(raw))
 		resp.ContentLength = int64(len(raw))
@@ -184,6 +188,11 @@ func chooseCompressEncoding(acceptEncoding, prefer string) string {
 			return "br"
 		}
 		return ""
+	case "zstd_only":
+		if strings.Contains(ae, "zstd") {
+			return "zstd"
+		}
+		return ""
 	default: // br_first
 		if strings.Contains(ae, "br") {
 			return "br"
@@ -191,6 +200,9 @@ func chooseCompressEncoding(acceptEncoding, prefer string) string {
 		if acceptEncodingSupportsGzip(ae) {
 			return "gzip"
 		}
+		if strings.Contains(ae, "zstd") {
+			return "zstd"
+		}
 		return ""
 	}
 }