《游戏服务端编程实践》5.1.2 连接管理与认证中转

// file: ch05/gateway/conn.go
// tier: L1
// verify: go test ./ch05/gateway/

package gateway

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

// ConnID is unique inside one gateway process. It is not a session ID: the
// login service issues sessions, the gateway only tracks sockets.
type ConnID string

var (
	ErrConnExists   = errors.New("gateway: connection already registered")
	ErrConnNotFound = errors.New("gateway: connection not found")
	ErrUnauthLimit  = errors.New("gateway: unauthenticated connection limit reached")
	ErrAlreadyBound = errors.New("gateway: connection already bound to another player")
	ErrInvalidBind  = errors.New("gateway: account and player IDs are required")
	ErrInvalidConn  = errors.New("gateway: connection ID is required")
)

// Conn is the gateway-side view of one client connection.
type Conn struct {
	ID            ConnID
	AccountID     uint64
	PlayerID      uint64
	GatewayID     string
	Authenticated bool
	CreatedAt     time.Time
	LastSeen      time.Time
}

// ConnRegistry owns connection facts for one gateway process. It is not a
// session store; authentication still comes from the login service.
type ConnRegistry struct {
	mu          sync.RWMutex
	conns       map[ConnID]*Conn
	maxUnauth   int
	authTimeout time.Duration
	now         func() time.Time
}

// NewConnRegistry builds a registry. The defaults (1024 unauthenticated
// connections, 5s first-packet timeout) are conservative starting values.
func NewConnRegistry(maxUnauth int, authTimeout time.Duration) *ConnRegistry {
	if maxUnauth <= 0 {
		maxUnauth = 1024
	}
	if authTimeout <= 0 {
		authTimeout = 5 * time.Second
	}
	return &ConnRegistry{
		conns:       make(map[ConnID]*Conn),
		maxUnauth:   maxUnauth,
		authTimeout: authTimeout,
		now:         time.Now,
	}
}

// Add registers a freshly accepted socket. Sockets above the unauthenticated
// limit are refused before they can consume a buffer or a handler goroutine.
func (r *ConnRegistry) Add(id ConnID, gatewayID string) error {
	if id == "" {
		return ErrInvalidConn
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.conns[id]; ok {
		return ErrConnExists
	}
	if r.unauthenticatedLocked() >= r.maxUnauth {
		return ErrUnauthLimit
	}
	now := r.now()
	r.conns[id] = &Conn{ID: id, GatewayID: gatewayID, CreatedAt: now, LastSeen: now}
	return nil
}

// Bind promotes a socket to authenticated and attaches the player. Repeating
// the same binding is idempotent; switching identity on a live socket is not.
func (r *ConnRegistry) Bind(id ConnID, accountID, playerID uint64) error {
	if accountID == 0 || playerID == 0 {
		return ErrInvalidBind
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	c, ok := r.conns[id]
	if !ok {
		return ErrConnNotFound
	}
	if c.Authenticated {
		if c.AccountID == accountID && c.PlayerID == playerID {
			c.LastSeen = r.now()
			return nil
		}
		return ErrAlreadyBound
	}
	c.AccountID = accountID
	c.PlayerID = playerID
	c.Authenticated = true
	c.LastSeen = r.now()
	return nil
}

// Remove drops a connection from the registry. It reports whether the
// connection existed so close callbacks can be idempotent.
func (r *ConnRegistry) Remove(id ConnID) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.conns[id]; !ok {
		return false
	}
	delete(r.conns, id)
	return true
}

// Lookup returns a snapshot; callers must not hold a pointer across calls.
func (r *ConnRegistry) Lookup(id ConnID) (Conn, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	c, ok := r.conns[id]
	if !ok {
		return Conn{}, false
	}
	return *c, true
}

// Touch records heartbeat activity for an authenticated connection.
func (r *ConnRegistry) Touch(id ConnID) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	c, ok := r.conns[id]
	if !ok || !c.Authenticated {
		return false
	}
	c.LastSeen = r.now()
	return true
}

// ExpireUnauthenticated removes sockets that did not finish authentication in
// time. Authenticated connections are never removed here; they close through
// heartbeat, logout or an explicit kick.
func (r *ConnRegistry) ExpireUnauthenticated() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	removed := 0
	for id, c := range r.conns {
		if !c.Authenticated && now.Sub(c.CreatedAt) >= r.authTimeout {
			delete(r.conns, id)
			removed++
		}
	}
	return removed
}

// Len returns the total number of live sockets.
func (r *ConnRegistry) Len() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.conns)
}

// Unauthenticated returns how many sockets are still waiting for a token.
func (r *ConnRegistry) Unauthenticated() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.unauthenticatedLocked()
}

func (r *ConnRegistry) unauthenticatedLocked() int {
	n := 0
	for _, c := range r.conns {
		if !c.Authenticated {
			n++
		}
	}
	return n
}