diff --git a/go.mod b/go.mod index bd48126..aa1fbbb 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index abcc27c..239927d 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/model/hosts.go b/model/hosts.go index ee8bce0..6237921 100644 --- a/model/hosts.go +++ b/model/hosts.go @@ -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，与类型取并集 diff --git a/utils/encodehelper.go b/utils/encodehelper.go index 4bd4db5..3dfd05e 100644 --- a/utils/encodehelper.go +++ b/utils/encodehelper.go @@ -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) +} diff --git a/wafenginecore/content_encoding.go b/wafenginecore/content_encoding.go index 756e027..b60fd93 100644 --- a/wafenginecore/content_encoding.go +++ b/wafenginecore/content_encoding.go @@ -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 diff --git a/wafenginecore/response_compress.go b/wafenginecore/response_compress.go index 1b64b97..3d64b0c 100644 --- a/wafenginecore/response_compress.go +++ b/wafenginecore/response_compress.go @@ -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 "" } }