diff --git a/controller/gcwork.go b/controller/gcwork.go new file mode 100644 index 000000000000..0bc17c6b7c18 --- /dev/null +++ b/controller/gcwork.go @@ -0,0 +1,278 @@ +package controller + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "strconv" + "sync" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// GC Work control plane endpoints. They are machine-to-machine calls made with +// an administrator personal access token holding the user:write scope, so they +// require no interactive security proof. They only act on common users: the +// control plane never provisions, disables, or mints relay tokens for +// administrators. + +const gcworkExternalIDMaxLength = 128 + +// gcworkEnsureLock serializes provisioning on this node; across nodes the +// unique username index decides, because every request for one external ID +// tries the same username candidates in the same order. +var gcworkEnsureLock sync.Mutex + +type gcworkEnsureUserRequest struct { + ExternalID string `json:"external_id"` + Username string `json:"username"` + DisplayName string `json:"display_name"` + Email string `json:"email"` +} + +// GcworkEnsureUser returns the user bound to a Casdoor subject, creating it in +// the default group with the default quota on first use. The subject is stored +// as the OIDC id, so NewAPI's own OIDC sign-in against the same Casdoor issuer +// reaches the same account. +func GcworkEnsureUser(c *gin.Context) { + var req gcworkEnsureUserRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil || req.ExternalID == "" || len(req.ExternalID) > gcworkExternalIDMaxLength { + common.ApiErrorT(c, "Invalid parameters") + return + } + gcworkEnsureLock.Lock() + defer gcworkEnsureLock.Unlock() + + if user, err := gcworkUserByExternalID(req.ExternalID); err != nil || user != nil { + gcworkEnsureResult(c, user, err) + return + } + candidates := []string{} + if req.Username != "" && len(req.Username) <= model.UserNameMaxLength { + candidates = append(candidates, req.Username) + } + digest := sha256.Sum256([]byte(req.ExternalID)) + candidates = append(candidates, "gcw_"+hex.EncodeToString(digest[:8])) + displayName := []rune(req.DisplayName) + if len(displayName) == 0 { + displayName = []rune(req.Username) + } + email := model.NormalizeEmail(req.Email) + if email != "" && model.EnsureEmailAvailable(email, 0) != nil { + email = "" + } + + var insertErr error + for _, username := range candidates { + user := model.User{ + Username: username, + DisplayName: string(displayName[:min(len(displayName), 20)]), + Email: email, + OidcId: req.ExternalID, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + insertErr = model.DB.Transaction(func(tx *gorm.DB) error { return user.InsertWithTx(tx, 0) }) + if errors.Is(insertErr, model.ErrEmailAlreadyTaken) { + user.Email, email = "", "" + insertErr = model.DB.Transaction(func(tx *gorm.DB) error { return user.InsertWithTx(tx, 0) }) + } + if insertErr == nil { + user.FinalizeOAuthUserCreation(0) + recordManageAuditFor(c, user.Id, "gcwork.user.ensure", map[string]any{"id": user.Id, "username": user.Username}) + common.ApiSuccess(c, gin.H{"id": user.Id}) + return + } + // Another node may have created the user with this candidate meanwhile. + if existing, err := gcworkUserByExternalID(req.ExternalID); err != nil || existing != nil { + gcworkEnsureResult(c, existing, err) + return + } + } + common.ApiError(c, insertErr) +} + +func gcworkUserByExternalID(externalID string) (*model.User, error) { + var user model.User + err := model.DB.Where("oidc_id = ?", externalID).Order("id").First(&user).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &user, nil +} + +func gcworkEnsureResult(c *gin.Context, user *model.User, err error) { + if err != nil { + common.ApiError(c, err) + return + } + if user.Role >= common.RoleAdminUser { + common.ApiErrorT(c, "No permission to update users of same or higher permission level") + return + } + common.ApiSuccess(c, gin.H{"id": user.Id}) +} + +// gcworkCommonUser loads the common user named by the :id path parameter, or +// writes the error response and returns nil. +func gcworkCommonUser(c *gin.Context) *model.User { + id, err := strconv.Atoi(c.Param("id")) + if err != nil || id <= 0 { + common.ApiErrorT(c, "Invalid parameters") + return nil + } + user := model.User{Id: id} + if err := model.DB.Where(&user).First(&user).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + common.ApiErrorT(c, "User does not exist") + } else { + common.ApiError(c, err) + } + return nil + } + if user.Role >= common.RoleAdminUser { + common.ApiErrorT(c, "No permission to update users of same or higher permission level") + return nil + } + return &user +} + +// GcworkSetUserStatus enables or disables a common user. Disabling stops every +// relay token of the user at the next request. +func GcworkSetUserStatus(c *gin.Context) { + var req struct { + Enabled *bool `json:"enabled"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil || req.Enabled == nil { + common.ApiErrorT(c, "Invalid parameters") + return + } + user := gcworkCommonUser(c) + if user == nil { + return + } + user.Status = common.UserStatusDisabled + if *req.Enabled { + user.Status = common.UserStatusEnabled + } + if err := user.Update(false); err != nil { + common.ApiError(c, err) + return + } + if err := model.InvalidateUserTokensCache(user.Id); err != nil { + common.SysLog(common.LogText("failed to invalidate tokens cache for user %d: %s", user.Id, err.Error())) + } + recordManageAuditFor(c, user.Id, "gcwork.user.status", map[string]any{"id": user.Id, "username": user.Username, "enabled": *req.Enabled}) + common.ApiSuccess(c, gin.H{"status": user.Status}) +} + +type gcworkCreateTokenRequest struct { + Name string `json:"name"` + ExpiredTime int64 `json:"expired_time"` + RemainQuota int `json:"remain_quota"` + UnlimitedQuota bool `json:"unlimited_quota"` +} + +// GcworkCreateToken creates a relay token owned by a common user, applying the +// validation of the user-facing token API, and returns its full key once. +func GcworkCreateToken(c *gin.Context) { + var req gcworkCreateTokenRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiErrorT(c, "Invalid parameters") + return + } + if len(req.Name) > 50 { + common.ApiErrorT(c, "Token name is too long") + return + } + if req.ExpiredTime != -1 && req.ExpiredTime <= common.GetTimestamp() { + common.ApiErrorT(c, "Invalid parameters") + return + } + if !req.UnlimitedQuota { + if req.RemainQuota < 0 { + common.ApiErrorT(c, "Quota value cannot be negative") + return + } + if maxQuotaValue := maxTokenQuota(); req.RemainQuota > maxQuotaValue { + common.ApiErrorT(c, "Quota value exceeds valid range, maximum is {{max}}", map[string]any{"max": logger.FormatQuota(maxQuotaValue)}) + return + } + } + user := gcworkCommonUser(c) + if user == nil { + return + } + count, err := model.CountUserTokens(user.Id) + if err != nil { + common.ApiError(c, err) + return + } + if maxTokens := operation_setting.GetMaxUserTokens(); int(count) >= maxTokens { + common.ApiErrorT(c, "Maximum number of tokens reached ({{max}})", map[string]any{"max": maxTokens}) + return + } + key, err := common.GenerateKey() + if err != nil { + common.ApiErrorT(c, "Failed to generate token") + common.SysLog(common.LogText("failed to generate token key: %s", err.Error())) + return + } + now := common.GetTimestamp() + token := model.Token{ + UserId: user.Id, + Name: req.Name, + Key: key, + CreatedTime: now, + AccessedTime: now, + ExpiredTime: req.ExpiredTime, + RemainQuota: req.RemainQuota, + UnlimitedQuota: req.UnlimitedQuota, + AllowUserQuotaQuery: true, + } + if err := token.Insert(); err != nil { + common.ApiError(c, err) + return + } + recordManageAuditFor(c, user.Id, "gcwork.token.create", map[string]any{"id": token.Id, "user_id": user.Id, "name": token.Name}) + common.ApiSuccess(c, gin.H{"id": token.Id, "key": "sk-" + key}) +} + +// GcworkDeleteToken deletes a relay token owned by a common user. A token that +// no longer exists counts as deleted. +func GcworkDeleteToken(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil || id <= 0 { + common.ApiErrorT(c, "Invalid parameters") + return + } + token, err := model.GetTokenById(id) + if errors.Is(err, gorm.ErrRecordNotFound) { + common.ApiSuccess(c, nil) + return + } + if err != nil { + common.ApiError(c, err) + return + } + var owner model.User + if err := model.DB.Unscoped().Select("id", "role").First(&owner, token.UserId).Error; err == nil && owner.Role >= common.RoleAdminUser { + common.ApiErrorT(c, "No permission to update users of same or higher permission level") + return + } + if err := token.Delete(); err != nil { + common.ApiError(c, err) + return + } + recordManageAuditFor(c, token.UserId, "gcwork.token.delete", map[string]any{"id": token.Id, "user_id": token.UserId}) + common.ApiSuccess(c, nil) +} diff --git a/controller/gcwork_test.go b/controller/gcwork_test.go new file mode 100644 index 000000000000..c4b33ba4c080 --- /dev/null +++ b/controller/gcwork_test.go @@ -0,0 +1,158 @@ +package controller + +import ( + "fmt" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/middleware" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newGcworkTestRouter mirrors router.SetGcworkRouter, so the access token +// route rules keyed by gin full path apply unchanged. +func newGcworkTestRouter() *gin.Engine { + router := gin.New() + router.Use(middleware.RequestId(), middleware.AccessTokenAudit()) + group := router.Group("/api/gcwork", middleware.AdminAuth()) + for _, route := range []struct { + method, path string + handler gin.HandlerFunc + }{ + {http.MethodPost, "/users/ensure", GcworkEnsureUser}, + {http.MethodPost, "/users/:id/status", GcworkSetUserStatus}, + {http.MethodPost, "/users/:id/tokens", GcworkCreateToken}, + {http.MethodDelete, "/tokens/:id", GcworkDeleteToken}, + } { + middleware.DeclareAccessTokenScopeRoute(route.method, "/api/gcwork"+route.path, "user:write") + group.Handle(route.method, route.path, route.handler) + } + return router +} + +func TestGcworkControlPlaneEndpoints(t *testing.T) { + admin, _ := setupAccessTokenAudit(t) + require.NoError(t, model.DB.AutoMigrate(&model.Token{})) + sqlDB, err := model.DB.DB() + require.NoError(t, err) + // Every goroutine must share the one in-memory SQLite database. + sqlDB.SetMaxOpenConns(1) + router := newGcworkTestRouter() + expiresAt := time.Now().Add(time.Hour).Unix() + writer, _ := createScopedAccessToken(t, admin.Id, expiresAt, "user:write") + reader, _ := createScopedAccessToken(t, admin.Id, expiresAt, "user:read") + call := func(credential, method, path, body string) tokenAPIResponse { + response := accessTokenRequest(router, method, path, credential, "", body) + var payload tokenAPIResponse + if response.Code == http.StatusOK { + require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload)) + } + return payload + } + ensure := `{"external_id":"casdoor-sub-1","username":"alice","display_name":"Alice Zhang","email":"alice@example.com"}` + + t.Run("requires the user:write scope", func(t *testing.T) { + response := accessTokenRequest(router, http.MethodPost, "/api/gcwork/users/ensure", reader, "", ensure) + assert.Equal(t, http.StatusForbidden, response.Code) + }) + + var userID int + t.Run("ensure creates one user per subject, also under concurrency", func(t *testing.T) { + ids := make([]int, 8) + var wg sync.WaitGroup + for i := range ids { + wg.Go(func() { + var data struct { + ID int `json:"id"` + } + payload := call(writer, http.MethodPost, "/api/gcwork/users/ensure", ensure) + if payload.Success && common.Unmarshal(payload.Data, &data) == nil { + ids[i] = data.ID + } + }) + } + wg.Wait() + userID = ids[0] + require.NotZero(t, userID) + for _, id := range ids { + assert.Equal(t, userID, id) + } + var users []model.User + require.NoError(t, model.DB.Where("oidc_id = ?", "casdoor-sub-1").Find(&users).Error) + require.Len(t, users, 1) + assert.Equal(t, "alice", users[0].Username) + assert.Equal(t, "Alice Zhang", users[0].DisplayName) + assert.Equal(t, "alice@example.com", users[0].Email) + assert.Equal(t, common.RoleCommonUser, users[0].Role) + assert.Equal(t, common.QuotaForNewUser, users[0].Quota) + }) + + t.Run("ensure falls back to a derived username and drops a taken email", func(t *testing.T) { + payload := call(writer, http.MethodPost, "/api/gcwork/users/ensure", `{"external_id":"casdoor-sub-2","username":"alice","email":"alice@example.com"}`) + require.True(t, payload.Success, payload.Message) + var user model.User + require.NoError(t, model.DB.Where("oidc_id = ?", "casdoor-sub-2").First(&user).Error) + assert.True(t, strings.HasPrefix(user.Username, "gcw_")) + assert.Len(t, user.Username, model.UserNameMaxLength) + assert.Empty(t, user.Email) + }) + + t.Run("ensure refuses a subject bound to an administrator", func(t *testing.T) { + require.NoError(t, model.DB.Model(admin).Update("oidc_id", "casdoor-admin").Error) + payload := call(writer, http.MethodPost, "/api/gcwork/users/ensure", `{"external_id":"casdoor-admin","username":"root2"}`) + assert.False(t, payload.Success) + }) + + var tokenID int + t.Run("creates a relay token owned by the user and returns its key", func(t *testing.T) { + body := fmt.Sprintf(`{"name":"gc-work device","expired_time":%d,"remain_quota":0,"unlimited_quota":true}`, expiresAt) + payload := call(writer, http.MethodPost, fmt.Sprintf("/api/gcwork/users/%d/tokens", userID), body) + require.True(t, payload.Success, payload.Message) + var data struct { + ID int `json:"id"` + Key string `json:"key"` + } + require.NoError(t, common.Unmarshal(payload.Data, &data)) + require.True(t, strings.HasPrefix(data.Key, "sk-")) + token, err := model.GetTokenById(data.ID) + require.NoError(t, err) + assert.Equal(t, userID, token.UserId) + assert.Equal(t, "sk-"+token.Key, data.Key) + tokenID = data.ID + }) + + t.Run("refuses tokens for administrators and past expiry times", func(t *testing.T) { + valid := `{"name":"x","expired_time":-1,"unlimited_quota":true}` + assert.False(t, call(writer, http.MethodPost, fmt.Sprintf("/api/gcwork/users/%d/tokens", admin.Id), valid).Success) + expired := `{"name":"x","expired_time":1,"unlimited_quota":true}` + assert.False(t, call(writer, http.MethodPost, fmt.Sprintf("/api/gcwork/users/%d/tokens", userID), expired).Success) + }) + + t.Run("disables and enables the user", func(t *testing.T) { + path := fmt.Sprintf("/api/gcwork/users/%d/status", userID) + require.True(t, call(writer, http.MethodPost, path, `{"enabled":false}`).Success) + user, err := model.GetUserById(userID, false) + require.NoError(t, err) + assert.Equal(t, common.UserStatusDisabled, user.Status) + require.True(t, call(writer, http.MethodPost, path, `{"enabled":true}`).Success) + user, err = model.GetUserById(userID, false) + require.NoError(t, err) + assert.Equal(t, common.UserStatusEnabled, user.Status) + assert.False(t, call(writer, http.MethodPost, fmt.Sprintf("/api/gcwork/users/%d/status", admin.Id), `{"enabled":false}`).Success) + }) + + t.Run("deletes a token, and deleting a missing token succeeds", func(t *testing.T) { + path := fmt.Sprintf("/api/gcwork/tokens/%d", tokenID) + require.True(t, call(writer, http.MethodDelete, path, "").Success) + _, err := model.GetTokenById(tokenID) + require.Error(t, err) + assert.True(t, call(writer, http.MethodDelete, path, "").Success) + }) +} diff --git a/middleware/gcwork_access_token_routes.go b/middleware/gcwork_access_token_routes.go new file mode 100644 index 000000000000..86a18a880d19 --- /dev/null +++ b/middleware/gcwork_access_token_routes.go @@ -0,0 +1,20 @@ +package middleware + +import "fmt" + +// DeclareAccessTokenScopeRoute records the access token scope a downstream +// (GC Work) route requires. It runs while routes are registered, so the rule +// exists only when the route does, and panics on conflicting declarations. +func DeclareAccessTokenScopeRoute(method, fullPath, scope string) { + key := method + " " + fullPath + rule := accessTokenScopeRule(scope) + if _, ok := accessTokenRouteRules[key]; ok { + panic(fmt.Sprintf("access token route %s is declared twice", key)) + } + permissionRouteRulesMu.Lock() + defer permissionRouteRulesMu.Unlock() + if existing, ok := permissionRouteRules[key]; ok && existing != rule { + panic(fmt.Sprintf("access token route %s is declared with %s and %s", key, existing.scope, rule.scope)) + } + permissionRouteRules[key] = rule +} diff --git a/router/gcwork-router.go b/router/gcwork-router.go new file mode 100644 index 000000000000..99133454eebd --- /dev/null +++ b/router/gcwork-router.go @@ -0,0 +1,39 @@ +package router + +import ( + "github.com/QuantumNous/new-api/controller" + "github.com/QuantumNous/new-api/middleware" + + "github.com/gin-contrib/gzip" + "github.com/gin-gonic/gin" +) + +// gcworkAccessTokenScope is the personal access token scope the GC Work +// control plane needs; it is the existing administrator user-management scope. +const gcworkAccessTokenScope = "user:write" + +// SetGcworkRouter adds the GC Work control plane endpoints: provisioning +// users by Casdoor subject and issuing per-device relay tokens for them. +func SetGcworkRouter(router *gin.Engine) { + gcworkRouter := router.Group("/api/gcwork") + gcworkRouter.Use(middleware.RouteTag("api")) + gcworkRouter.Use(gzip.Gzip(gzip.DefaultCompression)) + gcworkRouter.Use(middleware.AccessTokenAudit()) + gcworkRouter.Use(middleware.GlobalAPIRateLimit()) + gcworkRouter.Use(middleware.DisableCache()) + gcworkRouter.Use(middleware.AdminAuth()) + routes := []struct { + method string + path string + handler gin.HandlerFunc + }{ + {"POST", "/users/ensure", controller.GcworkEnsureUser}, + {"POST", "/users/:id/status", controller.GcworkSetUserStatus}, + {"POST", "/users/:id/tokens", controller.GcworkCreateToken}, + {"DELETE", "/tokens/:id", controller.GcworkDeleteToken}, + } + for _, route := range routes { + middleware.DeclareAccessTokenScopeRoute(route.method, joinPaths(gcworkRouter.BasePath(), route.path), gcworkAccessTokenScope) + gcworkRouter.Handle(route.method, route.path, route.handler) + } +} diff --git a/router/main.go b/router/main.go index c74e32c38692..342074fa792e 100644 --- a/router/main.go +++ b/router/main.go @@ -15,6 +15,7 @@ import ( func SetRouter(router *gin.Engine, assets WebAssets) { SetApiRouter(router) SetChatlogRouter(router) + SetGcworkRouter(router) SetDashboardRouter(router) SetRelayRouter(router) SetTaskPluginProtocolRouter(router)