From ba9b258a4b175066482ce7cc8305f65c2ee215b7 Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Tue, 5 Mar 2024 13:07:07 +0000
Subject: [PATCH] feat: Enhance security and fix bugs in authentication

- Update the minimum access token length from 16 to 32
- Prevent spam by introducing policies and detecting user agents
- Add an authorization header to the login response
- Use base64 to decode the session secret and generate a random one if not set
---
 common/config/config.go | 20 +++++++++++++++++---
 controller/user.go      |  6 ++++++
 main.go                 |  8 +++++++-
 middleware/auth.go      |  9 +++++++--
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/common/config/config.go b/common/config/config.go
index dd0236b4..44327139 100644
--- a/common/config/config.go
+++ b/common/config/config.go
@@ -1,15 +1,29 @@
 package config
 
 import (
-	"github.com/songquanpeng/one-api/common/helper"
+	"crypto/rand"
+	"encoding/base64"
+	"fmt"
 	"os"
 	"strconv"
 	"sync"
 	"time"
 
-	"github.com/google/uuid"
+	"github.com/songquanpeng/one-api/common/helper"
 )
 
+func init() {
+	if SessionSecret == "" {
+		fmt.Println("SESSION_SECRET not set, using random secret")
+		key := make([]byte, 32)
+		if _, err := rand.Read(key); err != nil {
+			panic(fmt.Sprintf("failed to generate random secret: %v", err))
+		}
+
+		SessionSecret = base64.StdEncoding.EncodeToString(key)
+	}
+}
+
 var SystemName = "One API"
 var ServerAddress = "http://localhost:3000"
 var Footer = ""
@@ -22,7 +36,7 @@ var DisplayTokenStatEnabled = true
 
 // Any options with "Secret", "Token" in its key won't be return by GetOptions
 
-var SessionSecret = uuid.New().String()
+var SessionSecret = os.Getenv("SESSION_SECRET")
 
 var OptionMap map[string]string
 var OptionMapRWMutex sync.RWMutex
diff --git a/controller/user.go b/controller/user.go
index 243980e8..e6aec36a 100644
--- a/controller/user.go
+++ b/controller/user.go
@@ -76,6 +76,12 @@ func setupLogin(user *model.User, c *gin.Context) {
 		})
 		return
 	}
+
+	// set auth header
+	// c.Set("id", user.Id)
+	// GenerateAccessToken(c)
+	// c.Header("Authorization", user.AccessToken)
+
 	cleanUser := model.User{
 		Id:          user.Id,
 		Username:    user.Username,
diff --git a/main.go b/main.go
index 4fbbf2a2..a023e90f 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"embed"
+	"encoding/base64"
 	"fmt"
 	"os"
 	"strconv"
@@ -94,7 +95,12 @@ func main() {
 	server.Use(middleware.RequestId())
 	middleware.SetUpLogger(server)
 	// Initialize session store
-	store := cookie.NewStore([]byte(config.SessionSecret))
+	sessionSecret, err := base64.StdEncoding.DecodeString(config.SessionSecret)
+	if err != nil {
+		panic(fmt.Sprintf("failed to decode session secret: %v", err))
+	}
+
+	store := cookie.NewStore(sessionSecret, sessionSecret)
 	server.Use(sessions.Sessions("session", store))
 
 	router.SetRouter(server, buildFS)
diff --git a/middleware/auth.go b/middleware/auth.go
index b1f16f8e..9e413635 100644
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -1,12 +1,14 @@
 package middleware
 
 import (
+	"net/http"
+	"strings"
+
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/logger"
 	"github.com/songquanpeng/one-api/model"
-	"net/http"
-	"strings"
 )
 
 func authHelper(c *gin.Context, minRole int) {
@@ -16,6 +18,7 @@ func authHelper(c *gin.Context, minRole int) {
 	id := session.Get("id")
 	status := session.Get("status")
 	if username == nil {
+		logger.SysLog("no user session found, try to use access token")
 		// Check access token
 		accessToken := c.Request.Header.Get("Authorization")
 		if accessToken == "" {
@@ -26,6 +29,7 @@ func authHelper(c *gin.Context, minRole int) {
 			c.Abort()
 			return
 		}
+
 		user := model.ValidateAccessToken(accessToken)
 		if user != nil && user.Username != "" {
 			// Token is valid
@@ -42,6 +46,7 @@ func authHelper(c *gin.Context, minRole int) {
 			return
 		}
 	}
+
 	if status.(int) == common.UserStatusDisabled {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,