312 lines
7.0 KiB
Go
312 lines
7.0 KiB
Go
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)
|
|
}
|