diff --git a/cmd/dashboard/controller/controller.go b/cmd/dashboard/controller/controller.go index 898b4e6..0e318dd 100644 --- a/cmd/dashboard/controller/controller.go +++ b/cmd/dashboard/controller/controller.go @@ -132,6 +132,9 @@ func routers(r *gin.Engine, frontendDist fs.FS) { auth.GET("/waf", pCommonHandler(listBlockedAddress)) auth.POST("/batch-delete/waf", adminHandler(batchDeleteBlockedAddress)) + auth.GET("/online-user", pCommonHandler(listOnlineUser)) + auth.GET("/online-user/batch-block", adminHandler(batchBlockOnlineUser)) + auth.PATCH("/setting", adminHandler(updateConfig)) r.NoRoute(fallbackToFrontend(frontendDist)) diff --git a/cmd/dashboard/controller/setting.go b/cmd/dashboard/controller/setting.go index 531d64d..90dea05 100644 --- a/cmd/dashboard/controller/setting.go +++ b/cmd/dashboard/controller/setting.go @@ -50,7 +50,7 @@ func listConfig(c *gin.Context) (model.SettingResponse, error) { // @Security BearerAuth // @Schemes // @Description Edit config -// @Tags auth required +// @Tags admin required // @Accept json // @Param body body model.SettingForm true "SettingForm" // @Produce json diff --git a/cmd/dashboard/controller/user.go b/cmd/dashboard/controller/user.go index 5b5c019..95d0ccd 100644 --- a/cmd/dashboard/controller/user.go +++ b/cmd/dashboard/controller/user.go @@ -2,6 +2,7 @@ package controller import ( "slices" + "strconv" "github.com/gin-gonic/gin" "golang.org/x/crypto/bcrypt" @@ -76,7 +77,7 @@ func updateProfile(c *gin.Context) (any, error) { // @Security BearerAuth // @Schemes // @Description List user -// @Tags auth required +// @Tags admin required // @Produce json // @Success 200 {object} model.CommonResponse[[]model.User] // @Router /user [get] @@ -93,7 +94,7 @@ func listUser(c *gin.Context) ([]model.User, error) { // @Security BearerAuth // @Schemes // @Description Create user -// @Tags auth required +// @Tags admin required // @Accept json // @param request body model.UserForm true "User Request" // @Produce json @@ -135,7 +136,7 @@ func createUser(c *gin.Context) (uint64, error) { // @Security BearerAuth // @Schemes // @Description Batch delete users -// @Tags auth required +// @Tags admin required // @Accept json // @param request body []uint true "id list" // @Produce json @@ -154,3 +155,59 @@ func batchDeleteUser(c *gin.Context) (any, error) { err := singleton.OnUserDelete(ids, newGormError) return nil, err } + +// List online users +// @Summary List online users +// @Security BearerAuth +// @Schemes +// @Description List online users +// @Tags auth required +// @Param limit query uint false "Page limit" +// @Param offset query uint false "Page offset" +// @Produce json +// @Success 200 {object} model.PaginatedResponse[[]model.OnlineUser, model.OnlineUser] +// @Router /online-user [get] +func listOnlineUser(c *gin.Context) (*model.Value[[]*model.OnlineUser], error) { + limit, err := strconv.Atoi(c.Query("limit")) + if err != nil || limit < 1 { + limit = 25 + } + + offset, err := strconv.Atoi(c.Query("offset")) + if err != nil || offset < 0 { + offset = 0 + } + + return &model.Value[[]*model.OnlineUser]{ + Value: singleton.GetOnlineUsers(limit, offset), + Pagination: model.Pagination{ + Offset: offset, + Limit: limit, + Total: int64(singleton.GetOnlineUserCount()), + }, + }, nil +} + +// Batch block online user +// @Summary Batch block online user +// @Security BearerAuth +// @Schemes +// @Description Batch block online user +// @Tags admin required +// @Accept json +// @Param request body []string true "block list" +// @Produce json +// @Success 200 {object} model.CommonResponse[any] +// @Router /online-user/batch-block [patch] +func batchBlockOnlineUser(c *gin.Context) (any, error) { + var list []string + if err := c.ShouldBindJSON(&list); err != nil { + return nil, err + } + + if err := singleton.BlockByIPs(list); err != nil { + return nil, newGormError("%v", err) + } + + return nil, nil +} diff --git a/cmd/dashboard/controller/waf.go b/cmd/dashboard/controller/waf.go index de90faa..a7dd375 100644 --- a/cmd/dashboard/controller/waf.go +++ b/cmd/dashboard/controller/waf.go @@ -56,7 +56,7 @@ func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAF], error) { // @Security BearerAuth // @Schemes // @Description Edit server -// @Tags auth required +// @Tags admin required // @Accept json // @Param request body []string true "block list" // @Produce json diff --git a/cmd/dashboard/controller/ws.go b/cmd/dashboard/controller/ws.go index 9fdf92e..0bbc440 100644 --- a/cmd/dashboard/controller/ws.go +++ b/cmd/dashboard/controller/ws.go @@ -10,6 +10,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + "github.com/hashicorp/go-uuid" "golang.org/x/sync/singleflight" "github.com/nezhahq/nezha/model" @@ -102,13 +103,29 @@ func checkSameOrigin(r *http.Request) bool { // @Success 200 {object} model.StreamServerData // @Router /ws/server [get] func serverStream(c *gin.Context) (any, error) { + connId, err := uuid.GenerateUUID() + if err != nil { + return nil, newWsError("%v", err) + } + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { return nil, newWsError("%v", err) } defer conn.Close() - singleton.OnlineUsers.Add(1) - defer singleton.OnlineUsers.Add(^uint64(0)) + + userIp := c.GetString(model.CtxKeyRealIPStr) + if userIp == "" { + userIp = c.RemoteIP() + } + + singleton.AddOnlineUser(connId, &model.OnlineUser{ + IP: userIp, + ConnectedAt: time.Now(), + Conn: conn, + }) + defer singleton.RemoveOnlineUser(connId) + count := 0 for { stat, err := getServerStat(c, count == 0) @@ -166,7 +183,7 @@ func getServerStat(c *gin.Context, withPublicNote bool) ([]byte, error) { return utils.Json.Marshal(model.StreamServerData{ Now: time.Now().Unix() * 1000, - Online: singleton.OnlineUsers.Load(), + Online: singleton.GetOnlineUserCount(), Servers: servers, }) }) diff --git a/model/server_api.go b/model/server_api.go index 5328e73..de2df16 100644 --- a/model/server_api.go +++ b/model/server_api.go @@ -16,7 +16,7 @@ type StreamServer struct { type StreamServerData struct { Now int64 `json:"now,omitempty"` - Online uint64 `json:"online,omitempty"` + Online int `json:"online,omitempty"` Servers []StreamServer `json:"servers,omitempty"` } diff --git a/model/user.go b/model/user.go index 29dfe4f..bff8d54 100644 --- a/model/user.go +++ b/model/user.go @@ -1,6 +1,9 @@ package model import ( + "time" + + "github.com/gorilla/websocket" "github.com/nezhahq/nezha/pkg/utils" "gorm.io/gorm" ) @@ -20,7 +23,6 @@ type User struct { type UserInfo struct { Role uint8 - _ [3]byte AgentSecret string } @@ -42,3 +44,11 @@ type Profile struct { User LoginIP string `json:"login_ip,omitempty"` } + +type OnlineUser struct { + UserID uint64 `json:"user_id,omitempty"` + ConnectedAt time.Time `json:"connected_at,omitempty"` + IP string `json:"ip,omitempty"` + + Conn *websocket.Conn `json:"-"` +} diff --git a/model/waf.go b/model/waf.go index a0a62a7..59815a6 100644 --- a/model/waf.go +++ b/model/waf.go @@ -14,12 +14,14 @@ const ( WAFBlockReasonTypeLoginFail WAFBlockReasonTypeBruteForceToken WAFBlockReasonTypeAgentAuthFail + WAFBlockReasonTypeManual ) const ( BlockIDgRPC = -127 + iota BlockIDToken BlockIDUnknownUser + BlockIDManual ) type WAFApiMock struct { diff --git a/service/singleton/online_user.go b/service/singleton/online_user.go new file mode 100644 index 0000000..5126864 --- /dev/null +++ b/service/singleton/online_user.go @@ -0,0 +1,68 @@ +package singleton + +import ( + "slices" + "sync" + + "github.com/nezhahq/nezha/model" +) + +var ( + OnlineUserMap = make(map[string]*model.OnlineUser) + OnlineUserMapLock = new(sync.Mutex) +) + +func AddOnlineUser(connId string, user *model.OnlineUser) { + OnlineUserMapLock.Lock() + defer OnlineUserMapLock.Unlock() + OnlineUserMap[connId] = user +} + +func RemoveOnlineUser(connId string) { + OnlineUserMapLock.Lock() + defer OnlineUserMapLock.Unlock() + delete(OnlineUserMap, connId) +} + +func BlockByIPs(ipList []string) error { + OnlineUserMapLock.Lock() + defer OnlineUserMapLock.Unlock() + + for _, ip := range ipList { + if err := model.BlockIP(DB, ip, model.WAFBlockReasonTypeManual, model.BlockIDManual); err != nil { + return err + } + for _, user := range OnlineUserMap { + if user.IP == ip && user.Conn != nil { + user.Conn.Close() + } + } + } + + return nil +} + +func GetOnlineUsers(limit, offset int) []*model.OnlineUser { + OnlineUserMapLock.Lock() + defer OnlineUserMapLock.Unlock() + var users []*model.OnlineUser + for _, user := range OnlineUserMap { + users = append(users, user) + } + slices.SortFunc(users, func(i, j *model.OnlineUser) int { + return i.ConnectedAt.Compare(j.ConnectedAt) + }) + if offset > len(users) { + return nil + } + if offset+limit > len(users) { + return users[offset:] + } + return users[offset : offset+limit] +} + +func GetOnlineUserCount() int { + OnlineUserMapLock.Lock() + defer OnlineUserMapLock.Unlock() + return len(OnlineUserMap) +} diff --git a/service/singleton/singleton.go b/service/singleton/singleton.go index d1a5f0f..c284e91 100644 --- a/service/singleton/singleton.go +++ b/service/singleton/singleton.go @@ -3,7 +3,6 @@ package singleton import ( _ "embed" "log" - "sync/atomic" "time" "github.com/patrickmn/go-cache" @@ -24,7 +23,6 @@ var ( Loc *time.Location FrontendTemplates []model.FrontendTemplate DashboardBootTime = uint64(time.Now().Unix()) - OnlineUsers = new(atomic.Uint64) ) //go:embed frontend-templates.yaml