From 4c809277aaf8ca202cf4aad80ae58518c75ef347 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Wed, 4 Dec 2024 23:53:02 +0800
Subject: [PATCH] feat: support br

---
 go.mod                 |  1 +
 go.sum                 |  4 ++++
 middleware/gzip.go     | 15 +++++++++++++--
 router/relay-router.go |  2 +-
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 4d278f78..0b78a62e 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index f4465e23..15bce874 100644
--- a/go.sum
+++ b/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=
diff --git a/middleware/gzip.go b/middleware/gzip.go
index 4d4ce0c2..5b9d566a 100644
--- a/middleware/gzip.go
+++ b/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
diff --git a/router/relay-router.go b/router/relay-router.go
index ccbe3025..63f5c36d 100644
--- a/router/relay-router.go
+++ b/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())