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