package tcp import ( "context" "crypto/tls" "fmt" "net" "strconv" "sync" "sync/atomic" "time" "github.com/panjf2000/gnet/v2" "go.uber.org/zap" authlogic "spark-server/internal/module/auth/logic" msglogic "spark-server/internal/module/message/logic" synclogic "spark-server/internal/module/sync/logic" "spark-server/internal/pkg/ratelimit" "spark-server/internal/pkg/reqctx" "spark-server/internal/transport/tcp/session" ) const ( defaultMaxFrame = 65536 defaultCompressMin = 512 defaultIdleTimeout = 120 * time.Second defaultOutbound = 1024 defaultRatePerSec = 20 defaultHandlerTimeout = 5 * time.Second maxFramesPerEvent = 32 ) // Metrics is a narrow counter hook. EventLoop must not block on exporters. type Metrics interface { Inc(name string) } type nopMetrics struct{} func (nopMetrics) Inc(string) {} // MsgIDGenerator is satisfied by infrastructure/id.Generator. type MsgIDGenerator interface { NextMsgID() (int64, error) } // Server is the gnet event engine. OnTraffic must not call MySQL, Redis, or deliver fanout. type Server struct { gnet.BuiltinEventEngine log *zap.Logger pool TaskPool sessions session.Manager eng gnet.Engine addr string maxFrame int opts []gnet.Option proxyProtocol bool idleTimeout time.Duration outboundSize int ratePerSec int compressMin int handlerTimeout time.Duration auth authlogic.Service ids MsgIDGenerator messages *msglogic.Service sync *synclogic.Service router *Router metrics Metrics authLimit *ratelimit.Limiter conns sync.Map seq atomic.Uint64 ready chan struct{} readyOnce sync.Once listening atomic.Bool } // Options configure the TCP listener. TLS is terminated at Nginx in production. type Options struct { Addr string Multicore bool KeepAlive time.Duration RecvBuffer int SendBuffer int MaxFrameSize int TLS *tls.Config ProxyProtocol bool IdleTimeout time.Duration OutboundSize int RatePerSec int CompressMin int HandlerTimeout time.Duration Auth authlogic.Service IDs MsgIDGenerator Messages *msglogic.Service Sync *synclogic.Service Metrics Metrics AuthLimiter *ratelimit.Limiter } // NewServer builds a gnet IM listener. Call Start to bind. func NewServer(log *zap.Logger, pool TaskPool, sessions session.Manager, o *Options) *Server { if o == nil { o = &Options{} } if log == nil { log = zap.NewNop() } applyTCPDefaults(o) s := &Server{ log: log, pool: pool, sessions: sessions, addr: o.Addr, maxFrame: o.MaxFrameSize, proxyProtocol: o.ProxyProtocol, idleTimeout: o.IdleTimeout, outboundSize: o.OutboundSize, ratePerSec: o.RatePerSec, compressMin: o.CompressMin, handlerTimeout: o.HandlerTimeout, auth: o.Auth, ids: o.IDs, messages: o.Messages, sync: o.Sync, metrics: o.Metrics, authLimit: o.AuthLimiter, router: NewRouter(), ready: make(chan struct{}), } if s.metrics == nil { s.metrics = nopMetrics{} } s.registerDefaultHandlers() if o.TLS != nil { log.Info("tcp tls cert ignored; terminate IM TLS at nginx (gnet v2 has no native TLS)") } s.opts = gnetOptions(o) return s } func applyTCPDefaults(o *Options) { if o.MaxFrameSize <= 0 { o.MaxFrameSize = defaultMaxFrame } if o.CompressMin <= 0 { o.CompressMin = defaultCompressMin } if o.IdleTimeout <= 0 { o.IdleTimeout = defaultIdleTimeout } if o.OutboundSize <= 0 { o.OutboundSize = defaultOutbound } if o.RatePerSec <= 0 { o.RatePerSec = defaultRatePerSec } if o.HandlerTimeout <= 0 { o.HandlerTimeout = defaultHandlerTimeout } } func gnetOptions(o *Options) []gnet.Option { opts := []gnet.Option{ gnet.WithMulticore(o.Multicore), gnet.WithReusePort(true), gnet.WithTicker(true), } if o.KeepAlive > 0 { opts = append(opts, gnet.WithTCPKeepAlive(o.KeepAlive)) } if o.RecvBuffer > 0 { opts = append(opts, gnet.WithSocketRecvBuffer(o.RecvBuffer)) } if o.SendBuffer > 0 { opts = append(opts, gnet.WithSocketSendBuffer(o.SendBuffer)) } return opts } // Router returns the command table so modules can register SEND/SYNC later. func (s *Server) Router() *Router { return s.router } func (s *Server) OnBoot(eng gnet.Engine) gnet.Action { s.eng = eng s.log.Info("tcp listening", zap.String("addr", s.addr)) s.listening.Store(true) s.signalReady() return gnet.None } func (s *Server) OnOpen(c gnet.Conn) ([]byte, gnet.Action) { s.attach(c) return nil, gnet.None } func (s *Server) OnClose(c gnet.Conn, _ error) gnet.Action { st := connFrom(c) if st == nil { return gnet.None } s.conns.Delete(st.sess.ID()) _ = s.sessions.Unbind(st.sess.ID()) st.sess.Detach() return gnet.None } func (s *Server) OnTick() (time.Duration, gnet.Action) { s.conns.Range(func(_, v any) bool { st, ok := v.(*connState) if !ok { return true } if st.sess.Idle(s.idleTimeout) { _ = st.gc.Close() return true } s.flush(st) return true }) return time.Second, gnet.None } func (s *Server) OnTraffic(c gnet.Conn) gnet.Action { st := connFrom(c) if st == nil { return gnet.Close } if action := s.consumeProxy(st); action != gnet.None { return action } if !st.sess.ProxyDone() { s.flush(st) return gnet.None } for i := 0; i < maxFramesPerEvent; i++ { action, more := s.consumeFrame(st) if action != gnet.None { return action } if !more { break } } s.flush(st) return gnet.None } func (s *Server) Start() error { err := gnet.Run(s, "tcp://"+s.addr, s.opts...) s.signalReady() return err } func (s *Server) WaitReady(ctx context.Context) error { select { case <-ctx.Done(): return ctx.Err() case <-s.ready: if s.listening.Load() { return nil } return fmt.Errorf("tcp not listening") } } func (s *Server) signalReady() { s.readyOnce.Do(func() { close(s.ready) }) } func (s *Server) Stop(ctx context.Context) error { if !s.listening.Load() { s.signalReady() return nil } s.listening.Store(false) return s.eng.Stop(ctx) } func (s *Server) attach(c gnet.Conn) { id := strconv.FormatUint(s.seq.Add(1), 10) remote := "unknown" if addr := c.RemoteAddr(); addr != nil { remote = addr.String() } sess := session.NewConn(id, remote, s.outboundSize) sess.SetRequestID(reqctx.NewID()) if !s.proxyProtocol { sess.MarkProxySkipped() } st := &connState{sess: sess, gc: c, limit: newLimiter(s.ratePerSec)} sess.SetIO(func(b []byte) error { return c.AsyncWrite(b, nil) }, func() error { return c.Close() }) sess.SetWake(func() { _ = c.Wake(nil) }) c.SetContext(st) s.conns.Store(id, st) } func (s *Server) allowAuthIP(ip string) bool { if s.authLimit == nil { return true } host := ip if h, _, err := net.SplitHostPort(ip); err == nil { host = h } if host == "" { s.metrics.Inc("realip_missing_total") host = "unknown" } return s.authLimit.Allow(host) }