package controller import ( "net/http/httptest" "testing" "time" jwt "github.com/appleboy/gin-jwt/v2" "github.com/gin-gonic/gin" "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/i18n" "github.com/nezhahq/nezha/service/singleton" "github.com/stretchr/testify/assert" "gorm.io/driver/sqlite" "gorm.io/gorm" ) func TestPayloadFunc(t *testing.T) { payloadFn := payloadFunc() // 测试包含IP的格式 t.Run("format with IP", func(t *testing.T) { data := map[string]interface{}{ "user_id": "123", "ip": "192.168.1.1", } claims := payloadFn(data) assert.Equal(t, "123", claims["user_id"]) assert.Equal(t, "192.168.1.1", claims["ip"]) }) // 测试不包含IP的格式 t.Run("format without IP", func(t *testing.T) { data := map[string]interface{}{ "user_id": "123", } claims := payloadFn(data) assert.Equal(t, "123", claims["user_id"]) assert.Nil(t, claims["ip"]) }) // 测试无效数据格式 t.Run("invalid data format", func(t *testing.T) { claims := payloadFn("123") // 字符串类型不再支持 assert.Empty(t, claims) }) // 测试空的map t.Run("empty map", func(t *testing.T) { data := map[string]interface{}{} claims := payloadFn(data) assert.Empty(t, claims) }) } func TestIPBinding(t *testing.T) { // 创建测试用的gin context gin.SetMode(gin.TestMode) t.Run("IP mismatch should invalidate token", func(t *testing.T) { // 模拟JWT claims包含IP绑定 claims := jwt.MapClaims{ "user_id": "123", "ip": "192.168.1.1", "exp": float64(time.Now().Add(time.Hour).Unix()), } // 这里需要实际的数据库和用户设置来完全测试 // 但可以测试claims的基本结构 assert.Equal(t, "123", claims["user_id"]) assert.Equal(t, "192.168.1.1", claims["ip"]) }) t.Run("no IP in token should deny access", func(t *testing.T) { // 没有IP绑定的token应该被拒绝 claims := jwt.MapClaims{ "user_id": "123", "exp": float64(time.Now().Add(time.Hour).Unix()), } // 验证token结构 assert.Equal(t, "123", claims["user_id"]) assert.Nil(t, claims["ip"]) }) } func TestValidateRuleRejectsForeignTriggerTasks(t *testing.T) { ctx := newMemberValidationContext(t) alertRule := &model.AlertRule{ Common: model.Common{UserID: 200}, Name: "member alert", Rules: []*model.Rule{{Type: "offline", Duration: 3}}, FailTriggerTasks: []uint64{42}, RecoverTriggerTasks: []uint64{42}, } assert.Error(t, validateRule(ctx, alertRule)) } func TestValidateServersRejectsForeignTriggerTasks(t *testing.T) { ctx := newMemberValidationContext(t) service := &model.Service{ Common: model.Common{UserID: 200}, Name: "member service", EnableTriggerTask: true, FailTriggerTasks: []uint64{42}, RecoverTriggerTasks: []uint64{42}, SkipServers: map[uint64]bool{}, } assert.Error(t, validateServers(ctx, service)) } func newMemberValidationContext(t *testing.T) *gin.Context { t.Helper() return newValidationContext(t, 200, model.RoleMember) } func newAdminValidationContext(t *testing.T) *gin.Context { t.Helper() return newValidationContext(t, 1, model.RoleAdmin) } func newValidationContext(t *testing.T, userID uint64, role model.Role) *gin.Context { t.Helper() originalDB := singleton.DB originalLoc := singleton.Loc originalLocalizer := singleton.Localizer originalCronShared := singleton.CronShared originalServerShared := singleton.ServerShared originalUserInfo := singleton.UserInfoMap db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) assert.NoError(t, err) sqlDB, err := db.DB() assert.NoError(t, err) assert.NoError(t, db.AutoMigrate( &model.Cron{}, &model.Server{}, &model.NotificationGroup{}, &model.NotificationGroupNotification{}, &model.ServerGroup{}, &model.ServerGroupServer{}, )) assert.NoError(t, db.Create(&model.Cron{ Common: model.Common{ID: 42, UserID: 1}, Name: "foreign trigger task", Command: "admin-maintenance", TaskType: model.CronTypeTriggerTask, Cover: model.CronCoverAlertTrigger, }).Error) assert.NoError(t, db.Create(&model.Cron{ Common: model.Common{ID: 43, UserID: 200}, Name: "member trigger task", Command: "member-task", TaskType: model.CronTypeTriggerTask, Cover: model.CronCoverAlertTrigger, }).Error) assert.NoError(t, db.Create(&model.NotificationGroup{ Common: model.Common{ID: 7, UserID: 1}, Name: "admin group", }).Error) assert.NoError(t, db.Create(&model.NotificationGroup{ Common: model.Common{ID: 8, UserID: 200}, Name: "member group", }).Error) singleton.DB = db singleton.Loc = time.Local singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) singleton.ServerShared = singleton.NewServerClass() singleton.CronShared = singleton.NewCronClass() singleton.UserLock.Lock() singleton.UserInfoMap = map[uint64]model.UserInfo{ 1: {Role: model.RoleAdmin}, 200: {Role: model.RoleMember}, } singleton.UserLock.Unlock() t.Cleanup(func() { singleton.CronShared.Close() _ = sqlDB.Close() singleton.DB = originalDB singleton.Loc = originalLoc singleton.Localizer = originalLocalizer singleton.CronShared = originalCronShared singleton.ServerShared = originalServerShared singleton.UserLock.Lock() singleton.UserInfoMap = originalUserInfo singleton.UserLock.Unlock() }) gin.SetMode(gin.TestMode) ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) ctx.Set(model.CtxKeyAuthorizedUser, &model.User{ Common: model.Common{ID: userID}, Role: role, }) return ctx }