96 lines
1.9 KiB
Go
96 lines
1.9 KiB
Go
// Package session keeps TCP connections in process memory.
|
|
package session
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
)
|
|
|
|
type Packet struct {
|
|
Command uint16
|
|
Flags uint16
|
|
RequestID uint64
|
|
SequenceID uint32
|
|
Body []byte
|
|
}
|
|
|
|
type Session interface {
|
|
ID() string
|
|
UserID() int64
|
|
DeviceID() uint64
|
|
RealIP() string
|
|
Send(packet *Packet) error
|
|
Close() error
|
|
}
|
|
|
|
type Manager interface {
|
|
Bind(session Session) error
|
|
Unbind(sessionID string) error
|
|
Get(sessionID string) (Session, bool)
|
|
GetUserSessions(userID int64) []Session
|
|
SendToUserExceptDevice(ctx context.Context, userID int64, exceptDeviceID uint64, packet *Packet) error
|
|
}
|
|
|
|
type memManager struct {
|
|
mu sync.RWMutex
|
|
byID map[string]Session
|
|
byUser map[int64]map[string]Session
|
|
}
|
|
|
|
func NewManager() Manager {
|
|
return &memManager{
|
|
byID: make(map[string]Session),
|
|
byUser: make(map[int64]map[string]Session),
|
|
}
|
|
}
|
|
|
|
func (m *memManager) Bind(s Session) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if old, ok := m.byID[s.ID()]; ok {
|
|
_ = old.Close()
|
|
}
|
|
m.kickSameDevice(s)
|
|
if m.byUser[s.UserID()] == nil {
|
|
m.byUser[s.UserID()] = make(map[string]Session)
|
|
}
|
|
m.byID[s.ID()] = s
|
|
m.byUser[s.UserID()][s.ID()] = s
|
|
return nil
|
|
}
|
|
|
|
func (m *memManager) Unbind(sessionID string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
s, ok := m.byID[sessionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
delete(m.byID, sessionID)
|
|
if um := m.byUser[s.UserID()]; um != nil {
|
|
delete(um, sessionID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *memManager) GetUserSessions(userID int64) []Session {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
um := m.byUser[userID]
|
|
out := make([]Session, 0, len(um))
|
|
for _, s := range um {
|
|
out = append(out, s)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (m *memManager) SendToUserExceptDevice(_ context.Context, userID int64, exceptDeviceID uint64, packet *Packet) error {
|
|
for _, s := range m.GetUserSessions(userID) {
|
|
if s.DeviceID() == exceptDeviceID {
|
|
continue
|
|
}
|
|
_ = s.Send(packet)
|
|
}
|
|
return nil
|
|
}
|