Răsfoiți Sursa

Merge pull request #593 from Calcium-Ion/gzip

feat: support br
Calcium-Ion 1 an în urmă
părinte
comite
6b2f675308
4 a modificat fișierele cu 19 adăugiri și 3 ștergeri
  1. 1 0
      go.mod
  2. 4 0
      go.sum
  3. 13 2
      middleware/gzip.go
  4. 1 1
      router/relay-router.go

+ 1 - 0
go.mod

@@ -7,6 +7,7 @@ toolchain go1.22.4
 
 require (
 	github.com/Calcium-Ion/go-epay v0.0.2
+	github.com/andybalholm/brotli v1.1.1
 	github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0
 	github.com/aws/aws-sdk-go-v2 v1.26.1
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.11

+ 4 - 0
go.sum

@@ -1,5 +1,7 @@
 github.com/Calcium-Ion/go-epay v0.0.2 h1:3knFBuaBFpHzsGeGQU/QxUqZSHh5s0+jGo0P62pJzWc=
 github.com/Calcium-Ion/go-epay v0.0.2/go.mod h1:cxo/ZOg8ClvE3VAnCmEzbuyAZINSq7kFEN9oHj5WQ2U=
+github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
+github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA=
 github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0 h1:onfun1RA+KcxaMk1lfrRnwCd1UUuOjJM/lri5eM1qMs=
 github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0/go.mod h1:4yg+jNTYlDEzBjhGS96v+zjyA3lfXlFd5CiTLIkPBLI=
 github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 h1:HblK3eJHq54yET63qPCTJnks3loDse5xRmmqHgHzwoI=
@@ -198,6 +200,8 @@ github.com/ugorji/go/codec v1.1.7/go.mod h1:Ax+UKWsSmolVDwsd+7N3ZtXu+yMGCf907BLY
 github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY=
 github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
 github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
+github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
 github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw=
 github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
 golang.org/x/arch v0.12.0 h1:UsYJhbzPYGsT0HbEdmYcqtCv8UNGvnaL561NnIUvaKg=

+ 13 - 2
middleware/gzip.go

@@ -2,14 +2,20 @@ package middleware
 
 import (
 	"compress/gzip"
+	"github.com/andybalholm/brotli"
 	"github.com/gin-gonic/gin"
 	"io"
 	"net/http"
 )
 
-func GzipDecodeMiddleware() gin.HandlerFunc {
+func DecompressRequestMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		if c.GetHeader("Content-Encoding") == "gzip" {
+		if c.Request.Body == nil || c.Request.Method == http.MethodGet {
+			c.Next()
+			return
+		}
+		switch c.GetHeader("Content-Encoding") {
+		case "gzip":
 			gzipReader, err := gzip.NewReader(c.Request.Body)
 			if err != nil {
 				c.AbortWithStatus(http.StatusBadRequest)
@@ -19,6 +25,11 @@ func GzipDecodeMiddleware() gin.HandlerFunc {
 
 			// Replace the request body with the decompressed data
 			c.Request.Body = io.NopCloser(gzipReader)
+			c.Request.Header.Del("Content-Encoding")
+		case "br":
+			reader := brotli.NewReader(c.Request.Body)
+			c.Request.Body = io.NopCloser(reader)
+			c.Request.Header.Del("Content-Encoding")
 		}
 
 		// Continue processing the request

+ 1 - 1
router/relay-router.go

@@ -9,7 +9,7 @@ import (
 
 func SetRelayRouter(router *gin.Engine) {
 	router.Use(middleware.CORS())
-	router.Use(middleware.GzipDecodeMiddleware())
+	router.Use(middleware.DecompressRequestMiddleware())
 	// https://platform.openai.com/docs/api-reference/introduction
 	modelsRouter := router.Group("/v1/models")
 	modelsRouter.Use(middleware.TokenAuth())