From 2bbe98e6942e0ab60ed4906d042704ac609cd6a4 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Thu, 4 Jan 2024 19:03:58 +0800
Subject: [PATCH] fix duplicated cors headers

---
 cors.go    | 15 ++++++++++++---
 main.go    |  6 +++++-
 process.go | 12 ++++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/cors.go b/cors.go
index e3655e6..5a68e02 100644
--- a/cors.go
+++ b/cors.go
@@ -4,10 +4,19 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+// this function is aborded
 func corsMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		c.Header("Access-Control-Allow-Origin", "*")
-		c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
-		c.Header("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+		// set cors header
+		header := c.Request.Header
+		if header.Get("Access-Control-Allow-Origin") == "" {
+			c.Header("Access-Control-Allow-Origin", "*")
+		}
+		if header.Get("Access-Control-Allow-Methods") == "" {
+			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+		}
+		if header.Get("Access-Control-Allow-Headers") == "" {
+			c.Header("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+		}
 	}
 }
diff --git a/main.go b/main.go
index e9358a5..4ff2da6 100644
--- a/main.go
+++ b/main.go
@@ -74,7 +74,7 @@ func main() {
 	m.Use(engine)
 
 	// CORS middleware
-	engine.Use(corsMiddleware())
+	// engine.Use(corsMiddleware())
 
 	// error handle middleware
 	engine.Use(func(c *gin.Context) {
@@ -90,6 +90,10 @@ func main() {
 
 	// CORS handler
 	engine.OPTIONS("/v1/*any", func(ctx *gin.Context) {
+		// set cros header
+		ctx.Header("Access-Control-Allow-Origin", "*")
+		ctx.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+		ctx.Header("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
 		ctx.AbortWithStatus(200)
 	})
 
diff --git a/process.go b/process.go
index 90bbbc5..1dc81aa 100644
--- a/process.go
+++ b/process.go
@@ -110,6 +110,18 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s
 		haveResponse = true
 		record.ResponseTime = time.Now().Sub(record.CreatedAt)
 		record.Status = r.StatusCode
+
+		// handle reverse proxy cors header if upstream do not set that
+		if r.Header.Get("Access-Control-Allow-Origin") == "" {
+			c.Header("Access-Control-Allow-Origin", "*")
+		}
+		if r.Header.Get("Access-Control-Allow-Methods") == "" {
+			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+		}
+		if r.Header.Get("Access-Control-Allow-Headers") == "" {
+			c.Header("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+		}
+
 		if !shouldResponse && r.StatusCode != 200 {
 			log.Println("upstream return not 200 and should not response", r.StatusCode)
 			return errors.New("upstream return not 200 and should not response")