diff --git a/opencat.go b/opencat.go index 7d5806c..6a091cd 100644 --- a/opencat.go +++ b/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") diff --git a/router/router.go b/router/router.go index c4bc2df..a18ccd9 100644 --- a/router/router.go +++ b/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) }