105 lines
2.0 KiB
Go
105 lines
2.0 KiB
Go
package gateway
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
)
|
|
|
|
type Client struct {
|
|
ID string
|
|
SendFunc func(data []byte) error
|
|
}
|
|
|
|
type Gateway struct {
|
|
mu sync.RWMutex
|
|
clients map[string]*Client // clientID -> Client
|
|
uidMap map[string][]string // uid -> []clientID
|
|
}
|
|
|
|
func NewGateway() *Gateway {
|
|
return &Gateway{
|
|
clients: make(map[string]*Client),
|
|
uidMap: make(map[string][]string),
|
|
}
|
|
}
|
|
|
|
func (g *Gateway) AddClient(c *Client) {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
g.clients[c.ID] = c
|
|
}
|
|
|
|
func (g *Gateway) RemoveClient(clientID string) {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
delete(g.clients, clientID)
|
|
for uid, list := range g.uidMap {
|
|
newList := []string{}
|
|
for _, cid := range list {
|
|
if cid != clientID {
|
|
newList = append(newList, cid)
|
|
}
|
|
}
|
|
g.uidMap[uid] = newList
|
|
}
|
|
}
|
|
|
|
func (g *Gateway) SendToAll(msg []byte) {
|
|
g.mu.RLock()
|
|
defer g.mu.RUnlock()
|
|
for _, c := range g.clients {
|
|
_ = c.SendFunc(msg)
|
|
}
|
|
}
|
|
|
|
func (g *Gateway) SendToClient(clientID string, msg []byte) error {
|
|
g.mu.RLock()
|
|
defer g.mu.RUnlock()
|
|
if c, ok := g.clients[clientID]; ok {
|
|
return c.SendFunc(msg)
|
|
}
|
|
return errors.New("client not found")
|
|
}
|
|
|
|
func (g *Gateway) BindUid(clientID, uid string) error {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
if _, ok := g.clients[clientID]; !ok {
|
|
return errors.New("client not found")
|
|
}
|
|
g.uidMap[uid] = append(g.uidMap[uid], clientID)
|
|
return nil
|
|
}
|
|
|
|
func (g *Gateway) SendToUid(uid string, msg []byte) {
|
|
g.mu.RLock()
|
|
defer g.mu.RUnlock()
|
|
if list, ok := g.uidMap[uid]; ok {
|
|
for _, cid := range list {
|
|
if c, ok := g.clients[cid]; ok {
|
|
_ = c.SendFunc(msg)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (g *Gateway) ListClients() []string {
|
|
g.mu.RLock()
|
|
defer g.mu.RUnlock()
|
|
ids := make([]string, 0, len(g.clients))
|
|
for id := range g.clients {
|
|
ids = append(ids, id)
|
|
}
|
|
return ids
|
|
}
|
|
|
|
func (g *Gateway) ListUids() map[string][]string {
|
|
g.mu.RLock()
|
|
defer g.mu.RUnlock()
|
|
result := make(map[string][]string, len(g.uidMap))
|
|
for uid, list := range g.uidMap {
|
|
result[uid] = append([]string(nil), list...)
|
|
}
|
|
return result
|
|
}
|