HyperNet/Nexus
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

♻️ Refactor websocket client id

Browse Source
This commit is contained in:
LittleSheep
2025-03-01 14:51:53 +08:00
parent 1f07c0edf4
commit 8ac88413e0
5 changed files with 60 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pkg/internal/grpc/stream.go
View File

@ -20,7 +20,7 @@ func (v *Server) CountStreamConnection(ctx context.Context, request *proto.Count
func (v *Server) PushStream(ctx context.Context, request *proto.PushStreamRequest) (*proto.PushStreamResponse, error) {
var cnt int
var successes []uint64
var successes []string
var errs []error
if request.UserId != nil {
cnt, successes, errs = ws.WebsocketPush(uint(request.GetUserId()), request.GetBody())
@ -32,7 +32,7 @@ func (v *Server) PushStream(ctx context.Context, request *proto.PushStreamReques
success := len(successes)
log.Debug().
Uint64("client_id", request.GetClientId()).
Str("client_id", request.GetClientId()).
Uint64("user_id", request.GetUserId()).
Int("count", cnt).
Int("success", success).
@ -61,7 +61,7 @@ func (v *Server) PushStream(ctx context.Context, request *proto.PushStreamReques
func (v *Server) PushStreamBatch(ctx context.Context, request *proto.PushStreamBatchRequest) (*proto.PushStreamResponse, error) {
var cnt int
var successes []uint64
var successes []string
var errs []error
if len(request.UserId) != 0 {
cnt, successes, errs = ws.WebsocketPushBatch(

41
pkg/internal/web/ws/connections.go
View File

@ -1,11 +1,12 @@
package ws
import (
"math/rand"
"fmt"
"sync"
"git.solsynth.dev/hypernet/nexus/pkg/internal/directory"
"git.solsynth.dev/hypernet/nexus/pkg/nex/sec"
"github.com/google/uuid"
"github.com/rs/zerolog/log"
"github.com/gofiber/contrib/websocket"
@ -13,20 +14,28 @@ import (
var (
wsMutex sync.Mutex
wsConn = make(map[uint]map[uint64]*websocket.Conn)
wsConn = make(map[uint]map[string]*websocket.Conn)
)
func ClientRegister(user sec.UserInfo, conn *websocket.Conn) uint64 {
func ClientRegister(user sec.UserInfo, conn *websocket.Conn) (string, error) {
wsMutex.Lock()
if wsConn[user.ID] == nil {
wsConn[user.ID] = make(map[uint64]*websocket.Conn)
wsConn[user.ID] = make(map[string]*websocket.Conn)
}
var clientId string
if userDefinedId := conn.Query("clientId"); len(userDefinedId) > 0 {
clientId = userDefinedId
} else {
clientId = uuid.NewString()
}
if _, ok := wsConn[user.ID][clientId]; ok {
return clientId, fmt.Errorf("client already conncted")
}
clientId := rand.Uint64()
wsConn[user.ID][clientId] = conn
wsMutex.Unlock()
log.Debug().
Uint64("client_id", clientId).
Str("client_id", clientId).
Uint("user_id", user.ID).
Msg("An client connected to stream endpoint...")
@ -35,19 +44,19 @@ func ClientRegister(user sec.UserInfo, conn *websocket.Conn) uint64 {
"id": clientId,
})
return clientId
return clientId, nil
}
func ClientUnregister(user sec.UserInfo, id uint64) {
func ClientUnregister(user sec.UserInfo, id string) {
wsMutex.Lock()
if wsConn[user.ID] == nil {
wsConn[user.ID] = make(map[uint64]*websocket.Conn)
wsConn[user.ID] = make(map[string]*websocket.Conn)
}
delete(wsConn[user.ID], id)
wsMutex.Unlock()
log.Debug().
Uint64("client_id", id).
Str("client_id", id).
Uint("user_id", user.ID).
Msg("An client disconnected from stream endpoint...")
@ -61,19 +70,19 @@ func ClientCount(uid uint) int {
return len(wsConn[uid])
}
func WebsocketPush(uid uint, body []byte) (count int, successes []uint64, errs []error) {
func WebsocketPush(uid uint, body []byte) (count int, successes []string, errs []error) {
for _, conn := range wsConn[uid] {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
successes = append(successes, uint64(uid))
successes = append(successes, fmt.Sprintf("%d", uid))
}
count++
}
return
}
func WebsocketPushDirect(clientId uint64, body []byte) (count int, successes []uint64, errs []error) {
func WebsocketPushDirect(clientId string, body []byte) (count int, successes []string, errs []error) {
for _, m := range wsConn {
if conn, ok := m[clientId]; ok {
if err := conn.WriteMessage(1, body); err != nil {
@ -87,13 +96,13 @@ func WebsocketPushDirect(clientId uint64, body []byte) (count int, successes []u
return
}
func WebsocketPushBatch(uidList []uint, body []byte) (count int, successes []uint64, errs []error) {
func WebsocketPushBatch(uidList []uint, body []byte) (count int, successes []string, errs []error) {
for _, uid := range uidList {
for _, conn := range wsConn[uid] {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
successes = append(successes, uint64(uid))
successes = append(successes, fmt.Sprintf("%d", uid))
}
count++
}
@ -101,7 +110,7 @@ func WebsocketPushBatch(uidList []uint, body []byte) (count int, successes []uin
return
}
func WebsocketPushBatchDirect(clientIdList []uint64, body []byte) (count int, successes []uint64, errs []error) {
func WebsocketPushBatchDirect(clientIdList []string, body []byte) (count int, successes []string, errs []error) {
for _, clientId := range clientIdList {
for _, m := range wsConn {
if conn, ok := m[clientId]; ok {

23
pkg/internal/web/ws/ws.go
View File

@ -11,24 +11,34 @@ import (
"github.com/gofiber/contrib/websocket"
jsoniter "github.com/json-iterator/go"
"github.com/samber/lo"
"github.com/spf13/viper"
)
func Listen(c *websocket.Conn) {
user, ok := c.Locals("nex_user").(*sec.UserInfo)
if !ok {
_ = c.WriteMessage(1, nex.WebSocketPackage{
Action: "error",
Message: "unauthorized",
}.Marshal())
c.Close()
return
}
// Push connection
clientId := ClientRegister(*user, c)
var err error
clientId, err := ClientRegister(*user, c)
if err != nil {
_ = c.WriteMessage(1, nex.WebSocketPackage{
Action: "error",
Message: "client with this id already connected",
}.Marshal())
c.Close()
return
}
// Event loop
var mt int
var data []byte
var err error
var packet nex.WebSocketPackage
for {
@ -42,11 +52,6 @@ func Listen(c *websocket.Conn) {
continue
}
aliasingMap := viper.GetStringMapString("services.aliases")
if val, ok := aliasingMap[packet.Endpoint]; ok {
packet.Endpoint = val
}
service := directory.GetServiceInstanceByType(packet.Endpoint)
if service == nil {
_ = c.WriteMessage(mt, nex.WebSocketPackage{