package session import ( "sync/atomic" "time" "spark-server/internal/pkg/errs" "spark-server/internal/pkg/reqctx" "spark-server/internal/transport/tcp/protocol" ) const ( StateConnecting uint32 = iota StateConnected StateAuthenticated StateActive StateClosing StateClosed ) type Conn struct { id string userID atomic.Int64 deviceID atomic.Uint64 realIP atomic.Value // string remote string state atomic.Uint32 lastActive atomic.Int64 drops atomic.Uint64 proxyDone atomic.Bool closed atomic.Bool outbound chan *Packet write func([]byte) error closeFn func() error wake func() } func NewConn(id, remote string, outboundSize int) *Conn { if outboundSize < 1 { outboundSize = 1024 } c := &Conn{ id: id, remote: remote, outbound: make(chan *Packet, outboundSize), } c.state.Store(StateConnected) c.lastActive.Store(time.Now().UnixNano()) c.realIP.Store("") return c } func (c *Conn) SetIO(write func([]byte) error, closeFn func() error) { c.write = write c.closeFn = closeFn } // SetWake registers EventLoop wakeup after outbound enqueue. func (c *Conn) SetWake(wake func()) { c.wake = wake } func (c *Conn) ID() string { return c.id } func (c *Conn) UserID() int64 { return c.userID.Load() } func (c *Conn) DeviceID() uint64 { return c.deviceID.Load() } func (c *Conn) RemoteAddr() string { return c.remote } func (c *Conn) Drops() uint64 { return c.drops.Load() } func (c *Conn) RealIP() string { v, _ := c.realIP.Load().(string) if v != "" { return v } return c.remote } func (c *Conn) SetRealIP(ip string) { c.realIP.Store(ip) c.proxyDone.Store(true) } func (c *Conn) MarkProxySkipped() { c.proxyDone.Store(true) } func (c *Conn) ProxyDone() bool { return c.proxyDone.Load() } func (c *Conn) Touch() { c.lastActive.Store(time.Now().UnixNano()) } func (c *Conn) Idle(timeout time.Duration) bool { if timeout <= 0 { return false } return time.Since(time.Unix(0, c.lastActive.Load())) > timeout } func (c *Conn) State() uint32 { return c.state.Load() } func (c *Conn) Authenticated() bool { st := c.state.Load() return st == StateAuthenticated || st == StateActive } func (c *Conn) Authenticate(userID int64, deviceID uint64) { c.userID.Store(userID) c.deviceID.Store(deviceID) c.state.Store(StateAuthenticated) } func (c *Conn) Activate() { if c.state.Load() == StateAuthenticated { c.state.Store(StateActive) } } func (c *Conn) Send(packet *Packet) error { if c.closed.Load() { return errs.ErrClosed } select { case c.outbound <- packet: c.notifyWake() return nil default: c.drops.Add(1) return errs.ErrOutboundDropped } } func (c *Conn) notifyWake() { if c.wake != nil { c.wake() } } func (c *Conn) Close() error { if !c.closed.CompareAndSwap(false, true) { return nil } c.state.Store(StateClosing) c.flushDirect() c.state.Store(StateClosed) if c.closeFn != nil { return c.closeFn() } return nil } // Detach marks the session closed without touching the socket (OnClose). func (c *Conn) Detach() { c.closed.Store(true) c.state.Store(StateClosed) c.write = nil c.closeFn = nil c.wake = nil } func (c *Conn) flushDirect() { for _, p := range c.DrainOutbound(64) { c.writePacket(p) } } func (c *Conn) writePacket(p *Packet) { if c.write == nil || p == nil { return } raw, err := protocol.Encode(protocol.Frame{ Magic: protocol.Magic, Version: protocol.CurrentVersion, Command: p.Command, Flags: p.Flags, RequestID: p.RequestID, SequenceID: p.SequenceID, Body: p.Body, }) if err != nil { return } _ = c.write(raw) } func (c *Conn) DrainOutbound(limit int) []*Packet { if limit < 1 { limit = 32 } out := make([]*Packet, 0, limit) for len(out) < limit { select { case p := <-c.outbound: if p != nil { out = append(out, p) } default: return out } } return out } func (c *Conn) RequestContext() *reqctx.RequestContext { return &reqctx.RequestContext{ UserID: c.UserID(), DeviceID: c.DeviceID(), ConnectionID: c.id, RemoteAddr: c.remote, RealIP: c.RealIP(), } }