add reset_root

opencat.go

@@ -1,13 +1,27 @@
 package main
 
 import (
+	"log"
 	"opencatd-open/router"
-	_ "opencatd-open/store"
+	"opencatd-open/store"
+	"os"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 )
 
 func main() {
+	args := os.Args[1:]
+	if len(args) > 0 && args[0] == "reset_root" {
+		log.Println("reset root token...")
+		ntoken := uuid.NewString()
+		if err := store.UpdateUser(uint(1), ntoken); err != nil {
+			log.Fatalln(err)
+			return
+		}
+		log.Println("new root token:", ntoken)
+		return
+	}
 
 	r := gin.Default()
 	group := r.Group("/1")

router/router.go

@@ -53,11 +53,27 @@ func AuthMiddleware() gin.HandlerFunc {
 			rootToken = u.Token
 		}
 		token := c.GetHeader("Authorization")
-		if token == "" || token[:7] != "Bearer " || token[7:] != rootToken {
+		if token == "" || token[:7] != "Bearer " {
 			c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
 			c.Abort()
 			return
 		}
+		if token[7:] != rootToken {
+			u, err := store.GetUserByID(uint(1))
+			if err != nil {
+				log.Println(err)
+				c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+				c.Abort()
+				return
+			}
+			if token[:7] != u.Token {
+				c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+				c.Abort()
+				return
+			}
+			rootToken = u.Token
+			store.LoadAuthCache()
+		}
 		// 可以在这里对 token 进行验证并检查权限
 
 		c.Next()
@@ -71,7 +87,7 @@ func Handleinit(c *gin.Context) {
 			u := store.User{Name: "root", Token: uuid.NewString()}
 			u.ID = 1
 			if err := store.CreateUser(&u); err != nil {
-				c.JSON(http.StatusOK, gin.H{
+				c.JSON(http.StatusForbidden, gin.H{
 					"error": err.Error(),
 				})
 				return
@@ -95,7 +111,7 @@ func Handleinit(c *gin.Context) {
 		return
 	}
 	if user.ID == uint(1) {
-		c.JSON(http.StatusOK, gin.H{
+		c.JSON(http.StatusForbidden, gin.H{
 			"error": "super user already exists, use cli to reset password",
 		})
 	}
@@ -214,14 +230,17 @@ func HandleResetUserToken(c *gin.Context) {
 	id := to.Int(c.Param("id"))
 
 	if err := store.UpdateUser(uint(id), uuid.NewString()); err != nil {
-		c.JSON(http.StatusOK, gin.H{"error": err.Error()})
+		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
 		return
 	}
 	u, err := store.GetUserByID(uint(id))
 	if err != nil {
-		c.JSON(http.StatusOK, gin.H{"error": err.Error()})
+		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
 		return
 	}
+	if u.ID == uint(1) {
+		rootToken = u.Token
+	}
 	c.JSON(http.StatusOK, u)
 }