diff --git a/middleware/auth.go b/middleware/auth.go index 10ec6f3..c026d2e 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -175,6 +175,20 @@ func RootAuth() func(c *gin.Context) { } } +func MasterNodeOnly() func(c *gin.Context) { + return func(c *gin.Context) { + if !common.IsMasterNode { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "only master node can access this endpoint", + }) + c.Abort() + return + } + c.Next() + } +} + func WssAuth(c *gin.Context) { } diff --git a/model/region_migration_user_test.go b/model/region_migration_user_test.go index f962e69..83707d1 100644 --- a/model/region_migration_user_test.go +++ b/model/region_migration_user_test.go @@ -74,6 +74,42 @@ func TestInsertImportedUserWithHashedPassword_RejectsExplicitId(t *testing.T) { require.Contains(t, err.Error(), "Id must not be set") } +func TestInsertImportedUserWithHashedPassword_RejectsOverlongFields(t *testing.T) { + db := setupRegionMigrationUserDB(t) + orig := DB + DB = db + defer func() { + DB = orig + }() + + user := &User{ + Username: "abcdefghijklmnopqrstu", + Password: "already-hashed-value", + Email: "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx@example.com", + DisplayName: "DDDDDDDDDDDDDDDDDDDDDDDDD", + Status: common.UserStatusEnabled, + Role: common.RoleCommonUser, + Group: "default", + Quota: 123, + AffCode: "A002", + CreatedAt: 1710000000, + } + + err := InsertImportedUserWithHashedPassword(user) + require.Error(t, err) + require.Contains(t, err.Error(), "username exceeds max length") + + user.Username = "valid-user" + err = InsertImportedUserWithHashedPassword(user) + require.Error(t, err) + require.Contains(t, err.Error(), "display_name exceeds max length") + + user.DisplayName = "Valid Name" + err = InsertImportedUserWithHashedPassword(user) + require.Error(t, err) + require.Contains(t, err.Error(), "email exceeds max length") +} + func TestCopyOAuthBindingsIfMissing_Idempotent(t *testing.T) { db := setupRegionMigrationUserDB(t) orig := DB diff --git a/model/user.go b/model/user.go index 13c9706..afe1b57 100644 --- a/model/user.go +++ b/model/user.go @@ -7,6 +7,7 @@ import ( "strconv" "strings" "time" + "unicode/utf8" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" @@ -579,6 +580,25 @@ func InsertImportedUserWithHashedPassword(user *User) error { return InsertImportedUserWithHashedPasswordTx(DB, user) } +func validateImportedUserFields(user *User) error { + if user == nil { + return errors.New("user is nil") + } + if strings.TrimSpace(user.Username) == "" { + return errors.New("username is required") + } + if utf8.RuneCountInString(user.Username) > UserNameMaxLength { + return fmt.Errorf("username exceeds max length %d", UserNameMaxLength) + } + if utf8.RuneCountInString(user.DisplayName) > 20 { + return errors.New("display_name exceeds max length 20") + } + if utf8.RuneCountInString(user.Email) > 50 { + return errors.New("email exceeds max length 50") + } + return nil +} + func InsertImportedUserWithHashedPasswordTx(tx *gorm.DB, user *User) error { if user == nil { return errors.New("user is nil") @@ -586,6 +606,9 @@ func InsertImportedUserWithHashedPasswordTx(tx *gorm.DB, user *User) error { if user.Id != 0 { return errors.New("InsertImportedUserWithHashedPassword: Id must not be set, let DB auto-assign") } + if err := validateImportedUserFields(user); err != nil { + return err + } if user.AffCode == "" { user.AffCode = common.GetRandomString(4) } diff --git a/router/api-router.go b/router/api-router.go index b254d6c..be15355 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -417,7 +417,7 @@ func SetApiRouter(router *gin.Engine) { } userMigrationRoute := apiRouter.Group("/user-migrations") - userMigrationRoute.Use(middleware.RootAuth()) + userMigrationRoute.Use(middleware.RootAuth(), middleware.MasterNodeOnly()) { userMigrationRoute.GET("/candidate-users", controller.ListMigrationCandidateUsers) userMigrationRoute.POST("/batches", controller.CreateUserMigrationBatch) diff --git a/router/user_migration_router_test.go b/router/user_migration_router_test.go index 65e47e3..3f67572 100644 --- a/router/user_migration_router_test.go +++ b/router/user_migration_router_test.go @@ -61,7 +61,7 @@ func setupUserMigrationRootRouter() *gin.Engine { session.Save() c.Request.Header.Set("New-Api-User", "1") c.Next() - }, middleware.RootAuth()) + }, middleware.RootAuth(), middleware.MasterNodeOnly()) { userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches) userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch) @@ -93,7 +93,7 @@ func setupUserMigrationAdminRouter() *gin.Engine { session.Save() c.Request.Header.Set("New-Api-User", "2") c.Next() - }, middleware.RootAuth()) + }, middleware.RootAuth(), middleware.MasterNodeOnly()) { userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches) userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch) @@ -106,10 +106,13 @@ func setupUserMigrationAdminRouter() *gin.Engine { func TestUserMigrationRoutes_RootCanAccess(t *testing.T) { db := setupUserMigrationRouterDB(t) orig := model.DB + origMaster := common.IsMasterNode model.DB = db defer func() { model.DB = orig + common.IsMasterNode = origMaster }() + common.IsMasterNode = true require.NoError(t, db.Create(&model.UserMigrationBatch{ Id: 1, @@ -141,10 +144,13 @@ func TestUserMigrationRoutes_RootCanAccess(t *testing.T) { func TestUserMigrationRoutes_AdminForbidden(t *testing.T) { db := setupUserMigrationRouterDB(t) orig := model.DB + origMaster := common.IsMasterNode model.DB = db defer func() { model.DB = orig + common.IsMasterNode = origMaster }() + common.IsMasterNode = true router := setupUserMigrationAdminRouter() @@ -166,3 +172,23 @@ func TestUserMigrationRoutes_AdminForbidden(t *testing.T) { require.Equal(t, http.StatusOK, w.Code) require.Contains(t, w.Body.String(), "success\":false") } + +func TestUserMigrationRoutes_RootForbiddenOnSlaveNode(t *testing.T) { + db := setupUserMigrationRouterDB(t) + orig := model.DB + origMaster := common.IsMasterNode + model.DB = db + defer func() { + model.DB = orig + common.IsMasterNode = origMaster + }() + common.IsMasterNode = false + + router := setupUserMigrationRootRouter() + + req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + require.Equal(t, http.StatusForbidden, w.Code) + require.Contains(t, w.Body.String(), "only master node") +} diff --git a/scripts/deploy-test.sh b/scripts/deploy-test.sh index e1cda44..61662c7 100644 --- a/scripts/deploy-test.sh +++ b/scripts/deploy-test.sh @@ -10,15 +10,12 @@ set -euo pipefail -# ==================== 配置 ==================== DOCKER_REGISTRY="registry.cn-hangzhou.aliyuncs.com/fengsilin/new-api" PROJECT_DIR="/mnt/d/code/new-api" -# 测试节点 CN_DIR="/root/new-api-cn" OV_DIR="/root/new-api-ov" -# ==================== 颜色 ==================== RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[1;33m' @@ -38,14 +35,13 @@ section() { echo -e " ${BOLD}=========================================${NC}" } -# ==================== 参数解析 ==================== TARGET="all" NO_BUILD=false for arg in "$@"; do case "$arg" in - cn) TARGET="cn" ;; - ov) TARGET="ov" ;; + cn) TARGET="cn" ;; + ov) TARGET="ov" ;; --no-build) NO_BUILD=true ;; -h|--help) echo "用法: $0 [cn|ov] [--no-build]" @@ -62,7 +58,6 @@ for arg in "$@"; do esac done -# ==================== Step 1: 构建并推送 ==================== if [ "$NO_BUILD" = false ]; then section "Step 1/2: 构建并推送镜像" @@ -99,7 +94,6 @@ else ok "使用镜像: ${IMAGE}" fi -# ==================== Step 2: 部署到测试节点 ==================== section "Step 2/2: 部署到测试环境" FAILED=() @@ -107,6 +101,7 @@ FAILED=() deploy_local() { local NAME="$1" local DIR="$2" + local NODE_TYPE_VALUE="$3" echo "" info "部署到 ${NAME}..." @@ -123,6 +118,11 @@ deploy_local() { echo " [1/3] 更新 image tag..." sed -i "s|${DOCKER_REGISTRY}:.*|${DOCKER_REGISTRY}:${TAG}|" docker-compose.yml + if grep -q "NODE_TYPE=" docker-compose.yml; then + sed -i "s|NODE_TYPE=.*|NODE_TYPE=${NODE_TYPE_VALUE}|" docker-compose.yml + else + sed -i "/SESSION_NAME=/a\\ - NODE_TYPE=${NODE_TYPE_VALUE}" docker-compose.yml + fi echo " [2/3] 拉取新镜像..." docker compose pull new-api @@ -136,15 +136,18 @@ deploy_local() { } case "$TARGET" in - cn) deploy_local "国内测试(CN)" "$CN_DIR" ;; - ov) deploy_local "国外测试(OV)" "$OV_DIR" ;; + cn) + deploy_local "国内测试(CN)" "$CN_DIR" "master" + ;; + ov) + deploy_local "国外测试(OV)" "$OV_DIR" "slave" + ;; all) - deploy_local "国内测试(CN)" "$CN_DIR" - deploy_local "国外测试(OV)" "$OV_DIR" + deploy_local "国内测试(CN)" "$CN_DIR" "master" + deploy_local "国外测试(OV)" "$OV_DIR" "slave" ;; esac -# ==================== 总结 ==================== echo "" section "部署总结" echo -e " 镜像: ${GREEN}${IMAGE}${NC}"