diff --git a/go.mod b/go.mod index 86e1156..4b76395 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module SamWaf go 1.23.0 require ( + github.com/andybalholm/brotli v1.2.0 github.com/bytedance/godlp v1.2.15 github.com/corazawaf/coraza/v3 v3.2.2 github.com/corazawaf/libinjection-go v0.2.2 diff --git a/go.sum b/go.sum index 4849924..8812674 100644 --- a/go.sum +++ b/go.sum @@ -15,6 +15,8 @@ github.com/ProtonMail/go-crypto v0.0.0-20230828082145-3c4c8a2d2371/go.mod h1:EjA github.com/ajstarks/svgo v0.0.0-20180226025133-644b8db467af/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw= github.com/aliyun/alibaba-cloud-sdk-go v1.63.100 h1:yUkCbrSM1cWtgBfRVKMQtdt22KhDvKY7g4V+92eG9wA= github.com/aliyun/alibaba-cloud-sdk-go v1.63.100/go.mod h1:SOSDHfe1kX91v3W5QiBsWSLqeLxImobbMX1mxrFHsVQ= +github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= +github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 h1:HblK3eJHq54yET63qPCTJnks3loDse5xRmmqHgHzwoI= github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6/go.mod h1:pbiaLIeYLUbgMY1kwEAdwO6UKD5ZNwdPGQlwokS9fe8= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= @@ -341,6 +343,8 @@ github.com/xuri/excelize/v2 v2.9.0 h1:1tgOaEq92IOEumR1/JfYS/eR0KHOCsRv/rYXXh6YJQ github.com/xuri/excelize/v2 v2.9.0/go.mod h1:uqey4QBZ9gdMeWApPLdhm9x+9o2lq4iVmjiLfBS5hdE= github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7 h1:hPVCafDV85blFTabnqKgNhDCkJX25eik94Si9cTER4A= github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= diff --git a/utils/deflatehelper.go b/utils/deflatehelper.go deleted file mode 100644 index 9deafda..0000000 --- a/utils/deflatehelper.go +++ /dev/null @@ -1,41 +0,0 @@ -package utils - -import ( - "bytes" - "compress/flate" -) - -func DeflateEncode(input []byte) ([]byte, error) { - // 创建一个新的 byte 输出流 - var buf bytes.Buffer - // 创建一个新的 gzip 输出流 - deflateWriter, err := flate.NewWriter(&buf, flate.DefaultCompression) - - // 将 input byte 数组写入到此输出流中 - _, err = deflateWriter.Write(input) - if err != nil { - _ = deflateWriter.Close() - return nil, err - } - if err := deflateWriter.Close(); err != nil { - return nil, err - } - // 返回压缩后的 bytes 数组 - return buf.Bytes(), nil -} -func DeflateDecode(input []byte) ([]byte, error) { - // 创建一个新的 gzip.Reader - bytesReader := bytes.NewReader(input) - flateReader := flate.NewReader(bytesReader) - - defer func() { - // defer 中关闭 gzipReader - _ = flateReader.Close() - }() - buf := new(bytes.Buffer) - // 从 Reader 中读取出数据 - if _, err := buf.ReadFrom(flateReader); err != nil { - return nil, err - } - return buf.Bytes(), nil -} diff --git a/utils/encodehelper.go b/utils/encodehelper.go new file mode 100644 index 0000000..4bd4db5 --- /dev/null +++ b/utils/encodehelper.go @@ -0,0 +1,110 @@ +package utils + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "github.com/andybalholm/brotli" +) + +func BrotliEncode(input []byte) ([]byte, error) { + // 创建一个新的 byte 输出流 + var buf bytes.Buffer + // 创建一个新的 brotli 输出流 + brotliWriter := brotli.NewWriter(&buf) + // 将 input byte 数组写入到此输出流中 + _, err := brotliWriter.Write(input) + if err != nil { + _ = brotliWriter.Close() + return nil, err + } + if err := brotliWriter.Close(); err != nil { + return nil, err + } + // 返回压缩后的 bytes 数组 + return buf.Bytes(), nil +} + +func BrotliDecode(input []byte) ([]byte, error) { + // 创建一个新的 brotli.Reader + bytesReader := bytes.NewReader(input) + brotliReader := brotli.NewReader(bytesReader) + + buf := new(bytes.Buffer) + // 从 Reader 中读取出数据 + if _, err := buf.ReadFrom(brotliReader); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func DeflateEncode(input []byte) ([]byte, error) { + // 创建一个新的 byte 输出流 + var buf bytes.Buffer + // 创建一个新的 gzip 输出流 + deflateWriter, err := flate.NewWriter(&buf, flate.DefaultCompression) + + // 将 input byte 数组写入到此输出流中 + _, err = deflateWriter.Write(input) + if err != nil { + _ = deflateWriter.Close() + return nil, err + } + if err := deflateWriter.Close(); err != nil { + return nil, err + } + // 返回压缩后的 bytes 数组 + return buf.Bytes(), nil +} +func DeflateDecode(input []byte) ([]byte, error) { + // 创建一个新的 gzip.Reader + bytesReader := bytes.NewReader(input) + flateReader := flate.NewReader(bytesReader) + + defer func() { + // defer 中关闭 gzipReader + _ = flateReader.Close() + }() + buf := new(bytes.Buffer) + // 从 Reader 中读取出数据 + if _, err := buf.ReadFrom(flateReader); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func GZipEncode(input []byte) ([]byte, error) { + // 创建一个新的 byte 输出流 + var buf bytes.Buffer + // 创建一个新的 gzip 输出流 + gzipWriter := gzip.NewWriter(&buf) + // 将 input byte 数组写入到此输出流中 + _, err := gzipWriter.Write(input) + if err != nil { + _ = gzipWriter.Close() + return nil, err + } + if err := gzipWriter.Close(); err != nil { + return nil, err + } + // 返回压缩后的 bytes 数组 + return buf.Bytes(), nil +} +func GZipDecode(input []byte) ([]byte, error) { + // 创建一个新的 gzip.Reader + bytesReader := bytes.NewReader(input) + gzipReader, err := gzip.NewReader(bytesReader) + if err != nil { + return nil, err + } + defer func() { + // defer 中关闭 gzipReader + _ = gzipReader.Close() + }() + buf := new(bytes.Buffer) + // 从 Reader 中读取出数据 + if _, err := buf.ReadFrom(gzipReader); err != nil { + return nil, err + } + return buf.Bytes(), nil +} diff --git a/utils/gziphelper.go b/utils/gziphelper.go deleted file mode 100644 index 81342b0..0000000 --- a/utils/gziphelper.go +++ /dev/null @@ -1,42 +0,0 @@ -package utils - -import ( - "bytes" - "compress/gzip" -) - -func GZipEncode(input []byte) ([]byte, error) { - // 创建一个新的 byte 输出流 - var buf bytes.Buffer - // 创建一个新的 gzip 输出流 - gzipWriter := gzip.NewWriter(&buf) - // 将 input byte 数组写入到此输出流中 - _, err := gzipWriter.Write(input) - if err != nil { - _ = gzipWriter.Close() - return nil, err - } - if err := gzipWriter.Close(); err != nil { - return nil, err - } - // 返回压缩后的 bytes 数组 - return buf.Bytes(), nil -} -func GZipDecode(input []byte) ([]byte, error) { - // 创建一个新的 gzip.Reader - bytesReader := bytes.NewReader(input) - gzipReader, err := gzip.NewReader(bytesReader) - if err != nil { - return nil, err - } - defer func() { - // defer 中关闭 gzipReader - _ = gzipReader.Close() - }() - buf := new(bytes.Buffer) - // 从 Reader 中读取出数据 - if _, err := buf.ReadFrom(gzipReader); err != nil { - return nil, err - } - return buf.Bytes(), nil -} diff --git a/wafenginecore/content_encoding.go b/wafenginecore/content_encoding.go index a475db2..fe25b28 100644 --- a/wafenginecore/content_encoding.go +++ b/wafenginecore/content_encoding.go @@ -9,6 +9,7 @@ import ( "compress/gzip" "errors" "fmt" + "github.com/andybalholm/brotli" "golang.org/x/net/html/charset" "golang.org/x/text/encoding" "golang.org/x/text/encoding/simplifiedchinese" @@ -69,6 +70,8 @@ func (waf *WafEngine) compressContent(res *http.Response, isStaticAssist bool, i respBytes, err = utils.GZipEncode(encodedBytes) case "deflate": respBytes, err = utils.DeflateEncode(encodedBytes) + case "br": + respBytes, err = utils.BrotliEncode(encodedBytes) default: respBytes = encodedBytes } @@ -106,6 +109,9 @@ func (waf *WafEngine) getOrgContent(resp *http.Response, isStaticAssist bool, de deflateReader := flate.NewReader(resp.Body) bodyReader = deflateReader defer deflateReader.Close() + case "br": + brotliReader := brotli.NewReader(resp.Body) + bodyReader = brotliReader default: bodyReader = resp.Body }