fix: add authentication mechanism to management interface

PBK-B · PBK-B · commit 67e4885d28ac · 2023-03-10T21:52:45.000+08:00
diff --git a/helper/config.go b/helper/config.go
@@ -56,11 +56,13 @@ func init() {
 
 // &#33719;&#21462;&#37197;&#32622;&#30446;&#24405;
 func getConfigPath() string {
-	appPath, err := os.Executable()
-	if err == nil {
-		appPath = filepath.Dir(appPath)
-		if appPath != "" {
-			return appPath + "/app.conf"
+	if IsRelease() {
+		appPath, err := os.Executable()
+		if err == nil {
+			appPath = filepath.Dir(appPath)
+			if appPath != "" {
+				return appPath + "/app.conf"
+			}
 		}
 	}
 	return "./app.conf"
@@ -109,7 +111,10 @@ func ReadConfig() (*DefaultConfig, error) {
 // &#37325;&#26032;&#21152;&#36733;&#37197;&#32622;&#25991;&#20214;
 func LoadConfig(configStr string) (*DefaultConfig, error) {
 	var config = &DefaultConfig{}
-	yaml.Unmarshal([]byte(configStr), config)
+	err := yaml.Unmarshal([]byte(configStr), config)
+	if err != nil {
+		return nil, err
+	}
 	return config, nil
 }
 
diff --git a/helper/default.go b/helper/default.go
@@ -9,10 +9,11 @@ import (
 	"math/rand"
 	"os"
 	"reflect"
+	"strings"
 	"time"
 )
 
-var AppName = "server"
+var AppName = "gpt-zmide-server"
 
 func init() {
 	type obj struct{}
@@ -21,7 +22,8 @@ func init() {
 }
 
 func IsRelease() bool {
-	return os.Getenv("DEBUG") == ""
+	arg1 := strings.ToLower(os.Args[0])
+	return (!strings.Contains(arg1, "go-build") && os.Getenv("DEBUG") == "")
 }
 
 // &#29983;&#25104;&#38543;&#26426;&#23383;&#31526;&#20018;
diff --git a/middleware/admin.go b/middleware/admin.go
@@ -0,0 +1,54 @@
+/*
+ * @Author: Bin
+ * @Date: 2023-03-10
+ * @FilePath: /gpt-zmide-server/middleware/admin.go
+ */
+package middleware
+
+import (
+	"crypto/md5"
+	"encoding/base64"
+	"fmt"
+	"gpt-zmide-server/controllers/apis"
+	"gpt-zmide-server/helper"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+)
+
+func adminCredential(token string) bool {
+	if token == "" {
+		return false
+	}
+
+	if decoded, err := base64.StdEncoding.DecodeString(token); err == nil && decoded != nil {
+		if input := strings.Split(string(decoded), ":"); len(input) == 2 {
+			userObj := helper.Config.AdminUser
+			user, password := input[0], fmt.Sprintf("%x", md5.Sum([]byte(input[1])))
+			if user == userObj.User && password == userObj.Password {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+func BasicAuthAdmin() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// Search user in the slice of allowed credentials
+		auth := strings.Replace(c.Request.Header.Get("Authorization"), "Basic ", "", -1)
+		if auth == "" {
+			auth = c.Query("token")
+		}
+
+		if !adminCredential(auth) {
+			// Credentials doesn't match, we return 401 and abort handlers chain.
+			apis.APIDefaultController.Fail(c, "&#35831;&#30331;&#24405;&#31649;&#29702;&#21592;&#36134;&#21495;")
+			c.AbortWithStatus(http.StatusUnauthorized)
+			return
+		}
+
+	}
+}
diff --git a/routers/routers.go b/routers/routers.go
@@ -41,7 +41,7 @@ func BuildRouter(r *gin.Engine) *gin.Engine {
 		openApis.POST("/", apisCtlOpen.Index)
 		openApis.POST("/query", apisCtlOpen.Query)
 
-		adminApis := api.Group("/admin")
+		adminApis := api.Group("/admin", middleware.BasicAuthAdmin())
 
 		// &#21518;&#21488;&#31649;&#29702;&#24212;&#29992;&#25509;&#21475;
 		adminApp := adminApis.Group("/application")

Original file line number	Diff line number	Diff line change
`@@ -9,10 +9,11 @@ import (`
`9`	`9`	`"math/rand"`
`10`	`10`	`"os"`
`11`	`11`	`"reflect"`
	`12`	`+ "strings"`
`12`	`13`	`"time"`
`13`	`14`	`)`
`14`	`15`
`15`		`-var AppName = "server"`
	`16`	`+var AppName = "gpt-zmide-server"`
`16`	`17`
`17`	`18`	`func init() {`
`18`	`19`	`type obj struct{}`
`@@ -21,7 +22,8 @@ func init() {`
`21`	`22`	`}`
`22`	`23`
`23`	`24`	`func IsRelease() bool {`
`24`		`- return os.Getenv("DEBUG") == ""`
	`25`	`+ arg1 := strings.ToLower(os.Args[0])`
	`26`	`+ return (!strings.Contains(arg1, "go-build") && os.Getenv("DEBUG") == "")`
`25`	`27`	`}`
`26`	`28`
`27`	`29`	`// 生成随机字符串`