《MMO 服务器开发》16.2 洪泛与连接耗尽防护

// file: internal/guard/limit.go
// tier: L2
// verify: go test ./internal/guard/ -run 'TestGateEnforcesPerIPQuota|TestGateDropsSlowHandshake' -v

// Package guard 为 aegis 接入层提供连接配额、握手限速与读写超时兜底。
package guard

import (
	"errors"
	"net"
	"sync"
	"time"
)

// 拒绝原因。接入层据此决定回执文案与是否计入风控信号。
var (
	ErrIPQuotaExceeded     = errors.New("guard: per-IP connection quota exceeded")
	ErrGlobalQuotaExceeded = errors.New("guard: global connection quota exceeded")
)

// Limits 是接入层的配额参数。零值字段由 New 补齐为默认值。
type Limits struct {
	MaxConnsPerIP    int           // 单个 IP 的并发连接上限
	MaxPendingPerIP  int           // 单 IP 未登录（未完成握手/登录）连接上限
	MaxConnsGlobal   int           // 全局并发连接上限
	HandshakeBurst   int           // 单 IP 握手令牌桶容量
	HandshakeRefill  time.Duration // 单 IP 补满一个握手令牌所需的间隔
	HandshakeTimeout time.Duration // 从 accept 到握手完成的墙钟上限
	ReadTimeout      time.Duration // 单次读空闲上限（slowloris 兜底）
	WriteTimeout     time.Duration // 单次写上限
}

// Limiter 维护全局与每 IP 的连接计数，以及每 IP 的握手令牌桶。
type Limiter struct {
	lim    Limits
	mu     sync.Mutex
	global int
	perIP  map[string]*ipState
}

type ipState struct {
	conns   int
	pending int
	tokens  float64
	last    time.Time
}

// New 用给定参数构造 Limiter，未填字段取默认值。
func New(lim Limits) *Limiter {
	if lim.MaxConnsPerIP <= 0 {
		lim.MaxConnsPerIP = 64
	}
	if lim.MaxPendingPerIP <= 0 {
		lim.MaxPendingPerIP = 8
	}
	if lim.MaxConnsGlobal <= 0 {
		lim.MaxConnsGlobal = 10000
	}
	if lim.HandshakeBurst <= 0 {
		lim.HandshakeBurst = 5
	}
	if lim.HandshakeRefill <= 0 {
		lim.HandshakeRefill = 200 * time.Millisecond
	}
	return &Limiter{lim: lim, perIP: make(map[string]*ipState)}
}

// Ticket 是一条连接占用的配额凭据。Authed 在握手/登录成功后调用一次，
// 把该连接从「未登录」计数移出；Close 在连接关闭时调用一次，幂等。
type Ticket struct {
	l      *Limiter
	ip     string
	st     *ipState
	authed bool
	once   sync.Once
}

// Acquire 尝试为来自 ip 的一条新连接占用配额。
func (l *Limiter) Acquire(ip string, now time.Time) (*Ticket, error) {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.global >= l.lim.MaxConnsGlobal {
		return nil, ErrGlobalQuotaExceeded
	}
	st := l.perIP[ip]
	if st == nil {
		st = &ipState{tokens: float64(l.lim.HandshakeBurst), last: now}
		l.perIP[ip] = st
	}
	if st.conns >= l.lim.MaxConnsPerIP || st.pending >= l.lim.MaxPendingPerIP {
		return nil, ErrIPQuotaExceeded
	}
	st.conns++
	st.pending++
	l.global++
	return &Ticket{l: l, ip: ip, st: st}, nil
}

// Authed 把这条连接标记为已登录，释放它占用的「未登录」名额。
func (t *Ticket) Authed() {
	t.l.mu.Lock()
	if !t.authed {
		t.authed = true
		t.st.pending--
	}
	t.l.mu.Unlock()
}

// Close 归还配额。幂等：重复调用只生效一次。
func (t *Ticket) Close() {
	t.once.Do(func() {
		t.l.mu.Lock()
		defer t.l.mu.Unlock()
		t.st.conns--
		t.l.global--
		if !t.authed {
			t.st.pending--
		}
		if t.st.conns == 0 && t.st.pending == 0 {
			delete(t.l.perIP, t.ip)
		}
	})
}

// AllowHandshake 是每 IP 的握手令牌桶：接受一条新连接时消耗一个令牌，
// 令牌按 HandshakeRefill 的间隔补充到 HandshakeBurst 上限。返回 false
// 表示该 IP 握手太频繁，应当直接关闭连接、不做任何回执。
func (l *Limiter) AllowHandshake(ip string, now time.Time) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	st := l.perIP[ip]
	if st == nil {
		st = &ipState{tokens: float64(l.lim.HandshakeBurst), last: now}
		l.perIP[ip] = st
	}
	if elapsed := now.Sub(st.last); elapsed > 0 {
		st.tokens += elapsed.Seconds() / l.lim.HandshakeRefill.Seconds()
		if st.tokens > float64(l.lim.HandshakeBurst) {
			st.tokens = float64(l.lim.HandshakeBurst)
		}
		st.last = now
	}
	if st.tokens < 1 {
		return false
	}
	st.tokens--
	return true
}

// Gate 把配额与超时套在一个 net.Listener 上：Accept 之后先查握手限速与配额，
// 再给连接设一个握手截止时间。任何一步失败都直接关闭，不发无谓的回执。
type Gate struct {
	Lim    *Limiter
	OnConn func(c net.Conn, tk *Ticket) // 配额通过后交给上层处理
}

// Serve 接受连接直到 listener 关闭。
func (g *Gate) Serve(ln net.Listener) error {
	for {
		c, err := ln.Accept()
		if err != nil {
			return err
		}
		go g.handle(c)
	}
}

func (g *Gate) handle(c net.Conn) {
	defer c.Close()
	host, _, err := net.SplitHostPort(c.RemoteAddr().String())
	if err != nil {
		return
	}
	now := time.Now()
	if !g.Lim.AllowHandshake(host, now) {
		return
	}
	tk, err := g.Lim.Acquire(host, now)
	if err != nil {
		return
	}
	defer tk.Close()
	if d := g.Lim.lim.HandshakeTimeout; d > 0 {
		_ = c.SetDeadline(now.Add(d))
	}
	if g.OnConn != nil {
		g.OnConn(c, tk)
	}
}

// SteadyTimeouts 给一条已完成握手的连接套上稳态的读写超时。slowloris 靠
// 「慢慢发字节」占住连接，所以读超时必须按「两次读之间的空闲」设置，
// 而不是整条连接的总时长；每收到一帧就由上层重新调用一次。
func (g *Gate) SteadyTimeouts(c net.Conn) {
	if d := g.Lim.lim.ReadTimeout; d > 0 {
		_ = c.SetReadDeadline(time.Now().Add(d))
	}
	if d := g.Lim.lim.WriteTimeout; d > 0 {
		_ = c.SetWriteDeadline(time.Now().Add(d))
	}
}

// snippet: internal/guard/limit_test.go（节选，startGate 已略）

// TestGateEnforcesPerIPQuota 是 L2：真连本机端口，超出每 IP 上限的连接被直接关闭。
func TestGateEnforcesPerIPQuota(t *testing.T) {
	const cap = 4
	lim := New(Limits{MaxConnsPerIP: cap, MaxConnsGlobal: 1000, HandshakeBurst: 1000})
	hold := make(chan struct{})
	ln, addr := startGate(t, lim, func(c net.Conn, tk *Ticket) {
		tk.Authed()
		_, _ = c.Write([]byte{'A'}) // 写标记字节 = 这条连接被接纳
		<-hold
	})
	defer ln.Close()
	defer close(hold)

	var accepted int
	var kept []net.Conn
	for i := 0; i < cap+3; i++ {
		c, err := net.Dial("tcp", addr)
		if err != nil {
			t.Fatalf("dial %d: %v", i, err)
		}
		_ = c.SetReadDeadline(time.Now().Add(2 * time.Second))
		buf := make([]byte, 1)
		if n, err := c.Read(buf); err == nil && n == 1 && buf[0] == 'A' {
			accepted++
			kept = append(kept, c) // 必须持有引用，否则 GC 会提前关掉它
			continue
		}
		_ = c.Close() // 读到 EOF/错误 = 被拒
	}
	for _, c := range kept {
		_ = c.Close()
	}
	if accepted != cap {
		t.Fatalf("accepted = %d, want exactly %d", accepted, cap)
	}
}