nezha_domains/cmd/dashboard/controller/jwt_test.go

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)
	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.CronShared = singleton.NewCronClass()
	singleton.ServerShared = singleton.NewServerClass()
	singleton.UserLock.Lock()
	singleton.UserInfoMap = map[uint64]model.UserInfo{
		1:   {Role: model.RoleAdmin},
		200: {Role: model.RoleMember},
	}
	singleton.UserLock.Unlock()
	t.Cleanup(func() {
		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
}