feat: support zstd request decompression (#6545)
This commit is contained in:
@@ -31,6 +31,7 @@ require (
|
||||
github.com/jfreymuth/oggvorbis v1.0.5
|
||||
github.com/jinzhu/copier v0.4.0
|
||||
github.com/joho/godotenv v1.5.1
|
||||
github.com/klauspost/compress v1.18.0
|
||||
github.com/mewkiz/flac v1.0.13
|
||||
github.com/nicksnyder/go-i18n/v2 v2.6.1
|
||||
github.com/pkg/errors v0.9.1
|
||||
@@ -126,7 +127,6 @@ require (
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/klauspost/compress v1.18.0 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
|
||||
github.com/leodido/go-urn v1.4.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
|
||||
+23
-2
@@ -8,6 +8,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/andybalholm/brotli"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
type readCloser struct {
|
||||
@@ -38,6 +39,7 @@ func DecompressRequestMiddleware() gin.HandlerFunc {
|
||||
wrapMaxBytes := func(body io.ReadCloser) io.ReadCloser {
|
||||
return http.MaxBytesReader(c.Writer, body, maxBytes)
|
||||
}
|
||||
decompressed := false
|
||||
|
||||
switch c.GetHeader("Content-Encoding") {
|
||||
case "gzip":
|
||||
@@ -55,7 +57,7 @@ func DecompressRequestMiddleware() gin.HandlerFunc {
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
c.Request.Header.Del("Content-Encoding")
|
||||
decompressed = true
|
||||
case "br":
|
||||
reader := brotli.NewReader(origBody)
|
||||
c.Request.Body = wrapMaxBytes(&readCloser{
|
||||
@@ -64,12 +66,31 @@ func DecompressRequestMiddleware() gin.HandlerFunc {
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
c.Request.Header.Del("Content-Encoding")
|
||||
decompressed = true
|
||||
case "zstd":
|
||||
reader, err := zstd.NewReader(origBody)
|
||||
if err != nil {
|
||||
_ = origBody.Close()
|
||||
c.AbortWithStatus(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
c.Request.Body = wrapMaxBytes(&readCloser{
|
||||
Reader: reader,
|
||||
closeFn: func() error {
|
||||
reader.Close()
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
decompressed = true
|
||||
default:
|
||||
// Even for uncompressed bodies, enforce a max size to avoid huge request allocations.
|
||||
c.Request.Body = wrapMaxBytes(origBody)
|
||||
}
|
||||
|
||||
if decompressed {
|
||||
c.Request.Header.Del("Content-Encoding")
|
||||
}
|
||||
|
||||
// Continue processing the request
|
||||
c.Next()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user