From bba617474594938bdc3046ef56c4d0bc40d94a9f Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Sun, 7 Jan 2024 21:27:28 +0800
Subject: [PATCH] fix: fix gemini panic

---
 middleware/recover.go  | 28 ++++++++++++++++++++++++++++
 model/usedata.go       |  6 ++++++
 router/relay-router.go |  2 +-
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 middleware/recover.go

diff --git a/middleware/recover.go b/middleware/recover.go
new file mode 100644
index 0000000..8338a51
--- /dev/null
+++ b/middleware/recover.go
@@ -0,0 +1,28 @@
+package middleware
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"net/http"
+	"one-api/common"
+	"runtime/debug"
+)
+
+func RelayPanicRecover() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		defer func() {
+			if err := recover(); err != nil {
+				common.SysError(fmt.Sprintf("panic detected: %v", err))
+				common.SysError(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack())))
+				c.JSON(http.StatusInternalServerError, gin.H{
+					"error": gin.H{
+						"message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/songquanpeng/one-api", err),
+						"type":    "one_api_panic",
+					},
+				})
+				c.Abort()
+			}
+		}()
+		c.Next()
+	}
+}
diff --git a/model/usedata.go b/model/usedata.go
index e00c3ce..9788512 100644
--- a/model/usedata.go
+++ b/model/usedata.go
@@ -18,6 +18,12 @@ type QuotaData struct {
 }
 
 func UpdateQuotaData() {
+	// recover
+	defer func() {
+		if r := recover(); r != nil {
+			common.SysLog(fmt.Sprintf("UpdateQuotaData panic: %s", r))
+		}
+	}()
 	for {
 		if common.DataExportEnabled {
 			common.SysLog("正在更新数据看板数据...")
diff --git a/router/relay-router.go b/router/relay-router.go
index fd80b30..f98e7ca 100644
--- a/router/relay-router.go
+++ b/router/relay-router.go
@@ -17,7 +17,7 @@ func SetRelayRouter(router *gin.Engine) {
 		modelsRouter.GET("/:model", controller.RetrieveModel)
 	}
 	relayV1Router := router.Group("/v1")
-	relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
+	relayV1Router.Use(middleware.RelayPanicRecover(), middleware.TokenAuth(), middleware.Distribute())
 	{
 		relayV1Router.POST("/completions", controller.Relay)
 		relayV1Router.POST("/chat/completions", controller.Relay)