diff --git a/config/config.go b/config/config.go index 5a8c0e1..f3d6bab 100644 --- a/config/config.go +++ b/config/config.go @@ -18,6 +18,9 @@ type Config struct { OpenaiApiBaseUrl string `yaml:"openai_api_base_url"` OpenaiModelName string `yaml:"openai_model_name"` OpenaiPrompt string `yaml:"openai_prompt"` + WebServer struct { + Token string `yaml:"token"` + } `yaml:"web_server"` } func init() { diff --git a/handler/restart/restart.go b/handler/restart/restart.go index b80f9ab..62ee07a 100644 --- a/handler/restart/restart.go +++ b/handler/restart/restart.go @@ -8,6 +8,7 @@ import ( "os/exec" "time" + "git.lxtend.com/qqbot/config" "git.lxtend.com/qqbot/constants" "git.lxtend.com/qqbot/handler" "git.lxtend.com/qqbot/model" @@ -186,6 +187,11 @@ func BuildBotHandler(c *gin.Context) { } func RestartBotHandler(c *gin.Context) { + token := c.Request.Header.Get("Authorization") + if token != config.ConfigManager.GetConfig().WebServer.Token { + c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"}) + return + } err := restartProgram() log.Println("重启程序...") if err != nil {