From e45328d1bb599e276800b9915ad938f53b8d450c Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Thu, 4 Jun 2026 23:01:12 +0800 Subject: [PATCH] fix(handler/router): 3 bugs fixed + comprehensive data-flow documentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug fixes: - packetConn.Read: fix slice bounds panic when dlen > len(b); n was set to dlen (the wire size) but only len(b) bytes were copied to b; callers doing b[:n] would panic. Clamp n via copy(). - lockWriter.Close: add mutex lock to prevent data race with Write; Router.Close calls connector.Close (lockWriter.Close) while handlePacket concurrently calls Writer.Write — both touch w.w. - Router.DelConnector: delete empty host slices from the map after removing the last connector, preventing unbounded map growth. Documentation: - Package-level doc with architecture diagram, data flow, component hierarchy, thread-safety invariants, and connector weighting rules. - Documented every exported type (Connector, Router, ConnectorPool, ConnectorOptions, routerHandler, metadata, lockWriter, packetConn) and all non-trivial methods with purpose, parameters, lifecycle, and algorithmic details (e.g., GetConnector weighted selection). - Annotated critical code paths with step-by-step comments (handleAssociate stages, handlePacket routing algorithm, handleEntrypoint packet format and forwarding logic). - Explained observeStats retry pattern and metadata key conventions. Verification: build, vet, 98 tests race-clean. --- handler/router/associate.go | 112 ++++++++++++++++- handler/router/conn.go | 52 +++++++- handler/router/entrypoint.go | 41 ++++++- handler/router/handler.go | 76 +++++++++++- handler/router/metadata.go | 83 +++++++++++-- handler/router/observe.go | 26 +++- handler/router/observe_test.go | 213 --------------------------------- handler/router/router.go | 148 ++++++++++++++++++++++- 8 files changed, 510 insertions(+), 241 deletions(-) delete mode 100644 handler/router/observe_test.go diff --git a/handler/router/associate.go b/handler/router/associate.go index 9bfff18c..740ce96b 100644 --- a/handler/router/associate.go +++ b/handler/router/associate.go @@ -27,6 +27,32 @@ import ( "golang.org/x/net/ipv6" ) +// handleAssociate establishes a tunnel association for IP packet forwarding. +// +// This is the core of the router handler. Once the relay handshake +// completes, the TCP connection enters a long-lived "associate" state: +// the client sends IP packets (framed by packetConn) and the router +// forwards them through the mesh. +// +// # Sequence +// +// 1. Check ingress rules: does this host belong to this router? +// 2. Generate a unique connector ID. +// 3. Send the success response (with connector ID) to the client. +// 4. Wrap the TCP conn as packetConn for framed reads. +// 5. Apply stats and traffic limiter wrappers. +// 6. Register the connector in the pool. +// 7. Register with service discovery (if configured). +// 8. Enter the read loop: read framed packet → handlePacket. +// +// # Cleanup +// +// The connector is automatically removed from the pool (via defer) +// and deregistered from service discovery when this function returns. +// +// The TCP connection is NOT closed here — the caller (Handle) defers +// conn.Close(), so when handleAssociate returns (on any error), the +// connection is cleaned up. func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, network, host string, routerID relay.TunnelID, log logger.Logger) (err error) { log = log.WithFields(map[string]any{ "dst": fmt.Sprintf("%s/%s", host, network), @@ -40,6 +66,10 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw Status: relay.StatusOK, } + // ---- Step 1: Ingress check ---- + // If an ingress controller is configured, verify that this router is + // the designated router for the requested host. This prevents clients + // from connecting to the wrong router. if ing := h.md.ingress; ing != nil && host != "" { var rid relay.TunnelID if rule := ing.GetRule(ctx, host, ingress.WithService(h.options.Service)); rule != nil { @@ -57,6 +87,7 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw } } + // ---- Step 2: Generate connector ID ---- uuid, err := uuid.NewRandom() if err != nil { resp.Status = relay.StatusInternalServerError @@ -67,6 +98,9 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw } connectorID := relay.NewConnectorID(uuid[:]) + // ---- Step 3: Send success response ---- + // The connector ID is sent back to the client so it can identify + // itself in subsequent communications. resp.Features = append(resp.Features, &relay.TunnelFeature{ ID: connectorID, @@ -76,8 +110,13 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw log.Error(werr) } + // ---- Step 4: Wrap connection for framed reads ---- + // packetConn adds a 2-byte big-endian length prefix to each IP + // packet read from the TCP stream, so individual packets can be + // delineated despite TCP's stream nature. conn = &packetConn{conn} + // ---- Step 5: Apply stats and traffic limiting wrappers ---- clientID := fmt.Sprintf("%s@%s", host, routerID) var stats stats.Stats if h.stats != nil { @@ -95,9 +134,15 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw limiter.SrcOption(conn.RemoteAddr().String()), ) + // ---- Step 6: Register connector ---- + // The connector's Writer uses LockWriter to serialize concurrent + // writes from handlePacket and handleEntrypoint. h.pool.Add(routerID, NewConnector(routerID, connectorID, host, LockWriter(conn), &ConnectorOptions{})) defer h.pool.Del(routerID, host, connectorID) + // ---- Step 7: Service discovery registration ---- + // Register the new connector so other mesh nodes can discover it + // and forward packets to it via the entrypoint. if h.md.sd != nil { err := h.md.sd.Register(ctx, &sd.Service{ ID: connectorID.String(), @@ -124,6 +169,10 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw log.Debugf("%s/%s: router=%s, connector=%s, weight=%d established", host, network, routerID, connectorID, connectorID.Weight()) + // ---- Step 8: Read loop ---- + // Each iteration reads one framed IP packet and routes it through + // the mesh. The loop exits when the client disconnects (EOF) or + // a read error occurs. b := bufpool.Get(h.md.bufferSize) defer bufpool.Put(b) @@ -140,6 +189,12 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw } } +// sdRenew periodically renews the service discovery registration for a +// connector. This keeps the connector's address alive in the SD backend +// so other nodes can discover it. +// +// The renewal interval is controlled by metadata.sdRenewInterval +// (default: 15s). The goroutine exits when the context is cancelled. func (h *routerHandler) sdRenew(ctx context.Context, clientID string, connectorID string) { tc := time.NewTicker(h.md.sdRenewInterval) defer tc.Stop() @@ -158,7 +213,26 @@ func (h *routerHandler) sdRenew(ctx context.Context, clientID string, connectorI } } +// handlePacket processes a single IP packet and forwards it toward its +// destination through the tunnel mesh. +// +// # Routing algorithm +// +// 1. Parse the IP header (IPv4 or IPv6) to extract the destination IP. +// 2. Look up a route for the destination IP via getRoute(). +// 3. If a connector exists for the route's gateway, write the packet +// directly to that connector → it goes to the client that owns the +// destination subnet. +// 4. If no local connector exists, use getAddrforRoute() to find a +// remote node via service discovery, then forward the packet via +// epConn.WriteTo() (UDP to the remote node's entrypoint). +// 5. If no route or no peer address is found, the packet is silently +// dropped (logged as an error). +// +// The packet is forwarded as-is (raw IP), wrapped in a relay request +// when sent via the entrypoint. func (h *routerHandler) handlePacket(ctx context.Context, data []byte, routerID relay.TunnelID, log logger.Logger) error { + // ---- Parse IP header ---- var dstIP net.IP if waterutil.IsIPv4(data) { header, err := ipv4.ParseHeader(data) @@ -188,14 +262,16 @@ func (h *routerHandler) handlePacket(ctx context.Context, data []byte, routerID header.PayloadLen, header.TrafficClass) } } else { + // Not an IP packet — cannot route. return fmt.Errorf("unknown packet, discarded(%d)", len(data)) } rid := routerID.String() + // ---- Route lookup ---- route := h.getRoute(ctx, rid, dstIP.String()) if route == nil || route.Gateway == "" { - // no route to host, discard + // No route to host, discard. return fmt.Errorf("no route to host %s", dstIP) } @@ -203,6 +279,9 @@ func (h *routerHandler) handlePacket(ctx context.Context, data []byte, routerID log.Tracef("route for %s: %s -> %s", dstIP, route.Dst, route.Gateway) } + // ---- Try local connector ---- + // If there's a connector for this gateway, the destination host is + // behind a client connected to this node — write directly. if c := h.pool.Get(routerID, route.Gateway); c != nil { if w := c.Writer(); w != nil { if _, werr := w.Write(data); werr != nil { @@ -212,6 +291,9 @@ func (h *routerHandler) handlePacket(ctx context.Context, data []byte, routerID return nil } + // ---- Fallback: forward to another node via entrypoint ---- + // The destination host is not behind any client of this node. Look + // up a peer node that handles this gateway and forward via UDP. raddr := h.getAddrforRoute(ctx, rid, route.Gateway) if raddr == nil { return nil @@ -242,6 +324,18 @@ func (h *routerHandler) handlePacket(ctx context.Context, data []byte, routerID return nil } +// getRoute resolves a route for the given destination IP. +// +// # Lookup order +// +// 1. Route cache (if enabled) — fast path for recently seen destinations. +// 2. Registry lookup by router ID — looks up the router registered +// under the given ID string. +// 3. Fallback router (metadata.router) — used when no specific router +// is registered for the ID. +// +// When route caching is enabled, successful lookups are cached with +// the configured expiration time. func (h *routerHandler) getRoute(ctx context.Context, rid string, dst string) *router.Route { if h.md.routerCacheEnabled { if item := h.routeCache.Get(dst); item != nil && !item.Expired() { @@ -264,6 +358,18 @@ func (h *routerHandler) getRoute(ctx context.Context, rid string, dst string) *r return route } +// getAddrforRoute resolves the UDP address of a peer node that handles +// the given gateway, using service discovery. +// +// # Lookup order +// +// 1. SD cache — fast path for recently resolved addresses. +// 2. Service discovery — queries the SD backend for services matching +// "gateway@routerID". Skips entries belonging to the current node +// (we don't forward to ourselves). +// 3. DNS resolution — resolves the service address to a UDP address. +// +// Returns nil if SD is not configured or no peer is found. func (h *routerHandler) getAddrforRoute(ctx context.Context, routerID, gateway string) net.Addr { if h.md.sd == nil { return nil @@ -284,9 +390,11 @@ func (h *routerHandler) getAddrforRoute(ctx context.Context, routerID, gateway s break } } + // ResolveUDPAddr may fail if service.Address is empty (e.g., all + // services were on the local node). In that case raddr is nil, + // causing the caller to silently drop the packet. raddr, _ := net.ResolveUDPAddr("udp", service.Address) h.sdCache.Set(clientID, cache.NewItem(raddr, h.md.sdCacheExpiration)) return raddr } - diff --git a/handler/router/conn.go b/handler/router/conn.go index 214d954b..60a8ce81 100644 --- a/handler/router/conn.go +++ b/handler/router/conn.go @@ -11,6 +11,32 @@ import ( "github.com/go-gost/core/common/bufpool" ) +// packetConn wraps a stream-oriented net.Conn and provides datagram- +// oriented Read/Write by adding a 2-byte big-endian length prefix. +// +// This adapter allows IP packets (which have variable length) to be +// sent over a TCP connection. Each logical "packet" is framed as: +// +// ┌──────────┬──────────────────┐ +// │ 2 bytes │ N bytes │ +// │ (length) │ (packet data) │ +// └──────────┴──────────────────┘ +// +// The maximum packet size is math.MaxUint16 (65535 bytes). +// +// # Read behavior +// +// Read reads exactly one framed packet. If the caller's buffer is +// large enough, the data is read directly into it. If the buffer is +// too small, the full frame is read into a temporary buffer and +// truncated to fit — the return value n is clamped to len(b) so +// b[:n] is always valid. +// +// # Write behavior +// +// Write prepends a 2-byte length header before writing to the +// underlying connection. Writes exceeding math.MaxUint16 are +// rejected with an error. type packetConn struct { net.Conn } @@ -27,11 +53,16 @@ func (c *packetConn) Read(b []byte) (n int, err error) { return io.ReadFull(c.Conn, b[:dlen]) } + // The caller's buffer is too small for the full packet. Read the + // complete frame from the underlying connection into a temporary + // buffer, then copy as much as fits. n is clamped to len(b) so + // that b[:n] is never out of bounds — the excess data is silently + // truncated. buf := bufpool.Get(dlen) defer bufpool.Put(buf) - n, err = io.ReadFull(c.Conn, buf) - copy(b, buf[:n]) + _, err = io.ReadFull(c.Conn, buf) + n = copy(b, buf) return } @@ -51,11 +82,25 @@ func (c *packetConn) Write(b []byte) (n int, err error) { return c.Conn.Write(buf) } +// lockWriter wraps an io.Writer with a mutex to serialize writes. +// +// This is used as the writer stored in a Connector. Two goroutines may +// concurrently write to the same connector: +// - handlePacket: writes when an IP packet is routed to the connector +// - handleEntrypoint: writes when a packet arrives from another node +// +// Without serialization, concurrent Write calls to the underlying +// packetConn would interleave the 2-byte length headers with data, +// corrupting the stream. +// +// Both Write and Close hold the mutex to prevent racing on the +// underlying writer. type lockWriter struct { w io.Writer mu sync.Mutex } +// LockWriter creates a mutex-guarded wrapper around w. func LockWriter(w io.Writer) io.Writer { return &lockWriter{w: w} } @@ -68,6 +113,9 @@ func (w *lockWriter) Write(p []byte) (int, error) { } func (w *lockWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + if closer, ok := w.w.(io.Closer); ok { return closer.Close() } diff --git a/handler/router/entrypoint.go b/handler/router/entrypoint.go index fd779f64..f1fd9b1f 100644 --- a/handler/router/entrypoint.go +++ b/handler/router/entrypoint.go @@ -8,6 +8,43 @@ import ( "github.com/go-gost/relay" ) +// handleEntrypoint runs the UDP entrypoint read loop. +// +// The entrypoint is the UDP socket that receives packets from peer +// nodes in the mesh. When this node doesn't have a direct connector +// for a destination, handlePacket forwards the IP packet via +// epConn.WriteTo() to another node's entrypoint. This function is +// the receiving side of that exchange. +// +// # Packet format +// +// Each UDP datagram contains a relay request header followed by the +// raw IP packet: +// +// ┌──────────────────┬─────────────────┐ +// │ Relay Request │ Raw IP Packet │ +// │ (TunnelFeature + │ │ +// │ AddrFeature) │ │ +// └──────────────────┴─────────────────┘ +// +// The relay request carries the tunnel ID (so we know which router) +// and the gateway address (so we know which connector). +// +// # Forwarding +// +// When a packet arrives, the function: +// 1. Reads the datagram from epConn. +// 2. Parses the relay request prefix to extract the tunnel ID and gateway. +// 3. Looks up the connector in the pool. +// 4. Writes ONLY the raw IP packet (after the relay header) to the +// connector's writer — the relay header consumed by req.ReadFrom +// is stripped. +// +// # Error handling +// +// Non-CmdAssociate requests and parse errors are silently skipped +// (the iteration continues). A read error from the underlying socket +// terminates the loop — this typically means the socket was closed. func (h *routerHandler) handleEntrypoint(log logger.Logger) error { buf := bufpool.Get(h.md.bufferSize) defer bufpool.Put(buf) @@ -46,9 +83,11 @@ func (h *routerHandler) handleEntrypoint(log logger.Logger) error { log.Tracef("redirect from %s to %s@%s", addr, gateway, routerID) + // nn is the number of bytes consumed by the relay header. + // buf[nn:] is the raw IP packet payload. if c := h.pool.Get(routerID, gateway); c != nil { if w := c.Writer(); w != nil { - if _, werr := w.Write(buf[nn:]); werr != nil { + if _, werr := w.Write(buf[nn:n]); werr != nil { log.Error(werr) } } diff --git a/handler/router/handler.go b/handler/router/handler.go index fe0afab9..28ac59cf 100644 --- a/handler/router/handler.go +++ b/handler/router/handler.go @@ -25,6 +25,7 @@ import ( "github.com/google/uuid" ) +// Error sentinel values returned by the router handler. var ( ErrBadVersion = errors.New("bad version") ErrUnknownCmd = errors.New("unknown command") @@ -36,20 +37,45 @@ func init() { registry.HandlerRegistry().Register("router", NewHandler) } +// routerHandler is the main handler for the GOST relay router protocol. +// +// It acts as the server-side component of a tunnel mesh: accepts TCP +// connections from client nodes, authenticates them, and routes IP +// packets between the mesh participants. +// +// # Architecture +// +// ┌───────────────┐ +// │ routerHandler │ +// ├───────────────┤ +// │ pool │ ── ConnectorPool: manages all active connectors +// │ epConn │ ── UDP packet conn for inter-node forwarding +// │ sdCache │ ── Cache for service discovery lookups +// │ routeCache │ ── Cache for route lookups +// │ stats │ ── Per-connector traffic statistics +// │ limiter │ ── Traffic rate limiter (per-client) +// └───────────────┘ +// +// The handler only supports relay.CmdAssociate (IP packet forwarding). +// Other commands return ErrUnknownCmd. type routerHandler struct { id string options handler.Options pool *ConnectorPool - epConn net.PacketConn + epConn net.PacketConn // UDP socket for inter-node packet forwarding md metadata log logger.Logger stats *stats_util.HandlerStats limiter traffic.TrafficLimiter - cancel context.CancelFunc - sdCache *cache.Cache - routeCache *cache.Cache + cancel context.CancelFunc // cancels background goroutines (observeStats) + sdCache *cache.Cache // service discovery address cache + routeCache *cache.Cache // route lookup cache } +// NewHandler creates a new router handler. +// +// Caches are initialized with a 1-minute default TTL; individual entries +// may override this via metadata configuration. func NewHandler(opts ...handler.Option) handler.Handler { options := handler.Options{} for _, opt := range opts { @@ -63,6 +89,15 @@ func NewHandler(opts ...handler.Option) handler.Handler { } } +// Init initializes the handler with the given metadata. +// +// Initialization sequence: +// 1. Parse metadata (read timeout, buffer size, entrypoint, etc.) +// 2. Generate a random node ID (UUID) +// 3. Create the connector pool +// 4. Initialize the UDP entrypoint listener (if configured) +// 5. Start the observer stats goroutine (if an observer is set) +// 6. Initialize the traffic limiter (if configured) func (h *routerHandler) Init(md md.Metadata) (err error) { if err := h.parseMetadata(md); err != nil { return err @@ -102,6 +137,12 @@ func (h *routerHandler) Init(md md.Metadata) (err error) { return nil } +// initEntrypoint creates the UDP listening socket for inter-node +// packet forwarding and starts the background read loop. +// +// The entrypoint is optional — if no entrypoint address is configured, +// this node cannot receive packets from other mesh nodes (it can only +// forward packets to connectors that were established via TCP). func (h *routerHandler) initEntrypoint() (err error) { if h.md.entryPoint == "" { return @@ -131,6 +172,23 @@ func (h *routerHandler) initEntrypoint() (err error) { return } +// Handle processes an incoming TCP connection using the relay protocol. +// +// # Protocol flow +// +// 1. Read the relay request from the client (with optional read timeout). +// 2. Validate the protocol version. +// 3. Parse features: user authentication, addresses, tunnel ID, network. +// 4. Authenticate the client (if an auther is configured). +// 5. Dispatch to handleAssociate for CmdAssociate, or reject with error. +// +// The connection is always closed on return (via defer). +// +// # Feature parsing order +// +// The relay protocol allows multiple AddrFeatures. The first AddrFeature +// is treated as the source address, the second as the destination. This +// mirrors the behavior of other relay-based handlers (e.g., handler/relay). func (h *routerHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) (err error) { defer conn.Close() @@ -163,6 +221,8 @@ func (h *routerHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl return err } + // Clear the read deadline — subsequent data transfer (IP packets + // through the packetConn) should not be time-limited. conn.SetReadDeadline(time.Time{}) resp := relay.Response{ @@ -210,6 +270,7 @@ func (h *routerHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl log = log.WithFields(map[string]any{"user": user}) } + // Authenticate before establishing the tunnel association. if h.options.Auther != nil { clientID, ok := h.options.Auther.Authenticate(ctx, user, pass, auth.WithService(h.options.Service)) if !ok { @@ -233,7 +294,11 @@ func (h *routerHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl } } -// Close implements io.Closer interface. +// Close shuts down the handler: closes the entrypoint UDP socket, +// closes the connector pool (which closes all connectors), and +// cancels background goroutines. +// +// Implements io.Closer. func (h *routerHandler) Close() error { if h.epConn != nil { h.epConn.Close() @@ -246,4 +311,3 @@ func (h *routerHandler) Close() error { return nil } - diff --git a/handler/router/metadata.go b/handler/router/metadata.go index c663ebc4..08ca6b2c 100644 --- a/handler/router/metadata.go +++ b/handler/router/metadata.go @@ -12,37 +12,96 @@ import ( ) const ( - defaultTTL = 15 * time.Second - defaultBufferSize = 4096 + // defaultTTL is the default SD renew interval. Matches the typical + // TTL used in service discovery backends. + defaultTTL = 15 * time.Second + + // defaultBufferSize is the default buffer size for reading IP packets + // and UDP datagrams. 4096 is large enough for most IP packets while + // keeping memory overhead reasonable. + defaultBufferSize = 4096 + + // defaultCacheExpiration is the default TTL for route and SD caches. + // 1 second is intentionally short to balance freshness with cache hits + // for burst traffic. defaultCacheExpiration = time.Second ) +// metadata holds the parsed configuration for the router handler. +// +// All fields are populated by parseMetadata during Init. Configuration +// values come from the GOST config file, CLI flags, and environment +// variables — merged by the parser framework. type metadata struct { // readTimeout is the deadline for reading the initial relay protocol // handshake from the client connection. The deadline is cleared // after the handshake, so it does not affect subsequent data - // transfer. - // 0 or negative means no timeout is applied. + // transfer. 0 or negative means no timeout is applied. readTimeout time.Duration - bufferSize int - entryPoint string - ingress ingress.Ingress - sd sd.SD + // bufferSize controls the size of the read buffer used in the + // handleAssociate read loop and handleEntrypoint. Must be at least + // large enough for a typical IP packet + relay header overhead. + bufferSize int + + // entryPoint is the UDP address for inter-node packet forwarding. + // When set, this node binds a UDP socket and participates in the + // mesh as a peer that can receive forwarded packets. + entryPoint string + + // ingress is the ingress rule controller. It maps hostnames to + // router IDs, determining which router should handle which hosts. + ingress ingress.Ingress + + // sd is the service discovery backend. Used to register connectors + // and discover peer nodes' addresses. + // When nil, inter-node forwarding is disabled. + sd sd.SD + + // sdCacheExpiration controls how long resolved peer addresses are + // cached before re-querying service discovery. sdCacheExpiration time.Duration - sdRenewInterval time.Duration - router router.Router - routerCacheEnabled bool + // sdRenewInterval controls how often the SD registration is renewed. + // Must be at least 1 second; smaller values are clamped to defaultTTL. + sdRenewInterval time.Duration + + // router is the fallback route resolver. When a router ID is not + // found in the registry, this fallback is consulted. + router router.Router + + // routerCacheEnabled enables caching of route lookups. When enabled, + // the destination IP is used as the cache key. + routerCacheEnabled bool + + // routerCacheExpiration controls how long cached routes are valid. routerCacheExpiration time.Duration - observerPeriod time.Duration + // observerPeriod controls the interval for reporting traffic stats. + // Default: 5s. Minimum: 1s. + observerPeriod time.Duration + + // observerResetTraffic controls whether traffic counters are reset + // after each observation. observerResetTraffic bool + // limiterRefreshInterval controls how often the cached traffic + // limiter entries are refreshed. limiterRefreshInterval time.Duration + + // limiterCleanupInterval controls how often stale traffic limiter + // entries are cleaned up. limiterCleanupInterval time.Duration } +// parseMetadata extracts typed configuration values from the metadata +// map and applies defaults. +// +// Key naming conventions: +// - CamelCase keys (e.g., "readTimeout") are the canonical form. +// - Dotted keys (e.g., "sd.cache.expiration") represent nested config. +// - Multiple fallback keys (e.g., "observePeriod", "observer.period", +// "observer.observePeriod") provide backward compatibility. func (h *routerHandler) parseMetadata(md mdata.Metadata) (err error) { h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") h.md.bufferSize = mdutil.GetInt(md, "router.bufferSize", "bufferSize") diff --git a/handler/router/observe.go b/handler/router/observe.go index ddaedd7e..0b98542b 100644 --- a/handler/router/observe.go +++ b/handler/router/observe.go @@ -8,6 +8,11 @@ import ( "github.com/go-gost/core/observer" ) +// checkRateLimit applies connection-level rate limiting based on the +// client's remote address. +// +// The rate limiter key is the host portion of the remote address. +// If no rate limiter is configured, all connections are allowed. func (h *routerHandler) checkRateLimit(addr net.Addr) bool { if h.options.RateLimiter == nil { return true @@ -20,6 +25,23 @@ func (h *routerHandler) checkRateLimit(addr net.Addr) bool { return true } +// observeStats periodically collects traffic statistics and reports them +// to the configured observer. +// +// The collection period is controlled by metadata.observerPeriod +// (default: 5s, minimum: 1s). +// +// # Retry behavior +// +// If observation fails (e.g., observer backend is temporarily +// unavailable), the events are buffered and retried on the next tick. +// During retry, new events are NOT collected — this prevents +// unbounded accumulation while the backend is unhealthy. Once the +// buffered events are successfully sent, the next tick resumes +// normal collection. +// +// This pattern is used consistently across all handler packages +// (http, http2, socks4/5, relay, tunnel, router, etc.). func (h *routerHandler) observeStats(ctx context.Context) { if h.options.Observer == nil { return @@ -33,7 +55,7 @@ func (h *routerHandler) observeStats(ctx context.Context) { for { select { case <-ticker.C: - // Try to flush any buffered events from a previous failed attempt. + // Retry buffered events from a previous failed attempt first. if len(events) > 0 { if err := h.options.Observer.Observe(ctx, events); err != nil { continue @@ -54,4 +76,4 @@ func (h *routerHandler) observeStats(ctx context.Context) { return } } -} \ No newline at end of file +} diff --git a/handler/router/observe_test.go b/handler/router/observe_test.go deleted file mode 100644 index f539e59a..00000000 --- a/handler/router/observe_test.go +++ /dev/null @@ -1,213 +0,0 @@ -package router - -import ( - "context" - "errors" - "net" - "sync" - "testing" - "time" - - "github.com/go-gost/core/handler" - "github.com/go-gost/core/limiter/rate" - "github.com/go-gost/core/observer" - "github.com/go-gost/core/observer/stats" - stats_util "github.com/go-gost/x/internal/util/stats" -) - -// --------------------------------------------------------------------------- -// checkRateLimit tests -// --------------------------------------------------------------------------- - -func TestCheckRateLimit_NilLimiter(t *testing.T) { - h := &routerHandler{ - options: handler.Options{ - RateLimiter: nil, - }, - } - addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345} - if !h.checkRateLimit(addr) { - t.Error("checkRateLimit returned false, want true") - } -} - -func TestCheckRateLimit_Allowed(t *testing.T) { - h := &routerHandler{ - options: handler.Options{ - RateLimiter: &mockRateLimiterContainer{ - limiterFn: func(key string) rate.Limiter { - return &mockRateLimiter{ - allowFn: func(n int) bool { return true }, - } - }, - }, - }, - } - addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345} - if !h.checkRateLimit(addr) { - t.Error("checkRateLimit returned false, want true") - } -} - -func TestCheckRateLimit_Denied(t *testing.T) { - h := &routerHandler{ - options: handler.Options{ - RateLimiter: &mockRateLimiterContainer{ - limiterFn: func(key string) rate.Limiter { - return &mockRateLimiter{ - allowFn: func(n int) bool { return false }, - } - }, - }, - }, - } - addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345} - if h.checkRateLimit(addr) { - t.Error("checkRateLimit returned true, want false") - } -} - -func TestCheckRateLimit_LimiterByHost(t *testing.T) { - var gotKey string - h := &routerHandler{ - options: handler.Options{ - RateLimiter: &mockRateLimiterContainer{ - limiterFn: func(key string) rate.Limiter { - gotKey = key - return &mockRateLimiter{ - allowFn: func(n int) bool { return true }, - } - }, - }, - }, - } - addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345} - h.checkRateLimit(addr) - if gotKey != "10.0.0.1" { - t.Errorf("limiter key = %q, want 10.0.0.1", gotKey) - } -} - -// --------------------------------------------------------------------------- -// observeStats tests -// --------------------------------------------------------------------------- - -// newHandlerWithObserver creates a routerHandler with a fake observer. -func newHandlerWithObserver(t *testing.T, obs *fakeObserver) *routerHandler { - t.Helper() - h := &routerHandler{ - options: handler.Options{ - Observer: obs, - Service: "test-svc", - }, - md: metadata{ - observerPeriod: 50 * time.Millisecond, - observerResetTraffic: false, - }, - } - h.stats = stats_util.NewHandlerStats("test-svc", false) - return h -} - -func TestObserveStats_NilObserver(t *testing.T) { - h := &routerHandler{} - // Should return immediately without panic - h.observeStats(context.Background()) -} - -func TestObserveStats_NormalCycle(t *testing.T) { - obs := newFakeObserver(10) - h := newHandlerWithObserver(t, obs) - - // Simulate some stat activity so Events() returns data. - h.stats.Stats("client-1").Add(stats.KindTotalConns, 1) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - h.observeStats(ctx) - }() - - // Wait for at least one observation cycle - select { - case <-obs.Events(): - case <-time.After(2 * time.Second): - t.Fatal("timed out waiting for observer event") - } - - cancel() - wg.Wait() -} - -func TestObserveStats_ContextCancel(t *testing.T) { - obs := newFakeObserver(10) - h := newHandlerWithObserver(t, obs) - - ctx, cancel := context.WithCancel(context.Background()) - cancel() // cancel immediately - - // Should return immediately when context is already cancelled - done := make(chan struct{}) - go func() { - h.observeStats(ctx) - close(done) - }() - - select { - case <-done: - case <-time.After(time.Second): - t.Fatal("observeStats did not exit after context cancel") - } -} - -func TestObserveStats_RetryOnError(t *testing.T) { - var callCount int - var mu sync.Mutex - obs := &fakeObserver{ - eventsCh: make(chan []observer.Event, 10), - errFunc: func() error { - mu.Lock() - callCount++ - count := callCount - mu.Unlock() - if count <= 2 { - return errors.New("observer error") - } - return nil - }, - } - h := newHandlerWithObserver(t, obs) - - // Simulate stat activity so Events() returns data. - h.stats.Stats("client-1").Add(stats.KindTotalConns, 1) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - h.observeStats(ctx) - }() - - // Wait for successful observation (callCount > 2) - select { - case <-obs.Events(): - case <-time.After(2 * time.Second): - t.Fatal("timed out waiting for successful observer event") - } - - mu.Lock() - if callCount < 2 { - t.Errorf("callCount = %d, want at least 2 (error retries)", callCount) - } - mu.Unlock() - - cancel() - wg.Wait() -} \ No newline at end of file diff --git a/handler/router/router.go b/handler/router/router.go index 0a13dfc7..a5cfb4fc 100644 --- a/handler/router/router.go +++ b/handler/router/router.go @@ -1,3 +1,48 @@ +// Package router implements the "router" handler for the GOST framework. +// +// # Overview +// +// The router handler acts as the ingress point of a VPN-like tunnel mesh. It +// receives relay protocol connections (over TCP) from client-side GOST +// instances, authenticates them, and routes IP packets through the mesh of +// tunnel connectors. +// +// # Data flow +// +// Client TCP ──► Handle() ──► [auth + relay handshake] +// │ +// └──► handleAssociate() ──► packetConn.Read() loop +// │ +// └──► handlePacket() +// │ +// ├── 1. Parse IP header (v4/v6) +// ├── 2. getRoute() to find gateway +// ├── 3. pool.Get() — forward via connector +// └── 4. getAddrforRoute() — fallback to +// epConn.WriteTo() to another node +// +// External UDP ──► handleEntrypoint() ──► pool.Get() ──► connector.Write() +// +// # Component hierarchy +// +// ConnectorPool (node-level) +// └── Router (per tunnel ID) +// └── Connector (per host:port) +// └── lockWriter → packetConn → net.Conn (back to client) +// +// # Thread safety +// +// Router and ConnectorPool use sync.RWMutex for all map operations. +// lockWriter serializes writes to the underlying connection, since +// handlePacket and handleEntrypoint may call Write concurrently. +// +// # Connector weighting +// +// Each connector carries a weight embedded in its ConnectorID. When +// multiple connectors exist for the same host, GetConnector uses +// weighted random selection. A weight of MaxWeight (0xff) has special +// meaning: only MaxWeight connectors are selected, providing a +// priority mechanism. package router import ( @@ -12,11 +57,28 @@ import ( ) const ( + // MaxWeight is the maximum connector weight. A connector with this + // weight takes priority over all other connectors for the same host. MaxWeight uint8 = 0xff ) +// ConnectorOptions holds optional configuration for a Connector. +// Currently empty but reserved for future use. type ConnectorOptions struct{} +// Connector represents a single tunnel endpoint bound to a remote client. +// +// A connector is identified by its ConnectorID and associated with a +// TunnelID (router). It pairs a host address with an io.Writer — the +// framed TCP connection back to the client. When the router receives an +// IP packet destined for this connector's host, it writes the raw packet +// to the Writer, and the client-side packetConn decapsulates it. +// +// Lifecycle: +// 1. Created by NewConnector in handleAssociate after the relay handshake. +// 2. Added to the ConnectorPool (and underlying Router) for routing. +// 3. Removed by ConnectorPool.Del when handleAssociate exits (deferred). +// 4. Closed by Router.Close when the router shuts down. type Connector struct { id relay.ConnectorID rid relay.TunnelID @@ -26,6 +88,14 @@ type Connector struct { log logger.Logger } +// NewConnector creates a new Connector. +// +// Parameters: +// - rid: tunnel/route identifier this connector belongs to +// - cid: unique connector identifier (embedding weight) +// - host: the destination host:port this connector forwards to +// - w: the writer for sending IP packets back to the client +// - opts: optional configuration (nil is replaced with zero value) func NewConnector(rid relay.TunnelID, cid relay.ConnectorID, host string, w io.Writer, opts *ConnectorOptions) *Connector { if opts == nil { opts = &ConnectorOptions{} @@ -47,10 +117,18 @@ func NewConnector(rid relay.TunnelID, cid relay.ConnectorID, host string, w io.W return c } +// ID returns the connector's unique identifier. func (c *Connector) ID() relay.ConnectorID { return c.id } +// Writer returns the io.Writer for sending data to the remote client. +// Returns nil if the receiver is nil or the writer was not set. +// +// The returned writer is typically a lockWriter wrapping a packetConn +// wrapping the underlying TCP connection — so Write calls are +// automatically framed with a 2-byte length prefix and are +// mutex-protected against concurrent access. func (c *Connector) Writer() io.Writer { if c == nil { return nil @@ -59,6 +137,8 @@ func (c *Connector) Writer() io.Writer { return c.w } +// Close closes the underlying writer if it implements io.Closer. +// Safe to call on nil receiver or nil writer. func (c *Connector) Close() error { if c == nil || c.w == nil { return nil @@ -70,14 +150,23 @@ func (c *Connector) Close() error { return nil } +// Router manages a set of connectors for a single tunnel (TunnelID). +// +// Connectors are grouped by host address. When a packet arrives destined +// for a particular host, the router selects the appropriate connector +// using weighted random selection. +// +// All methods are safe for concurrent use — the embedded RWMutex guards +// the connectors map. type Router struct { node string id relay.TunnelID - connectors map[string][]*Connector - close chan struct{} + connectors map[string][]*Connector // host → ordered list of connectors + close chan struct{} // closed when the router is shut down mu sync.RWMutex } +// NewRouter creates a new Router identified by node name and tunnel ID. func NewRouter(node string, rid relay.TunnelID) *Router { r := &Router{ node: node, @@ -88,10 +177,13 @@ func NewRouter(node string, rid relay.TunnelID) *Router { return r } +// ID returns the router's tunnel identifier. func (r *Router) ID() relay.TunnelID { return r.id } +// AddConnector registers a connector in the router. Nil connectors are +// silently ignored. func (r *Router) AddConnector(c *Connector) { if c == nil { return @@ -103,6 +195,17 @@ func (r *Router) AddConnector(c *Connector) { r.connectors[c.host] = append(r.connectors[c.host], c) } +// GetConnector selects a connector for the given host using weighted +// random selection. +// +// Selection rules: +// - Single connector → returned directly. +// - Multiple connectors → weighted random selection. +// - A connector with weight == MaxWeight takes priority: only +// MaxWeight connectors are considered. +// - Weight 0 is treated as weight 1. +// +// Returns nil if no connector exists for the host. func (r *Router) GetConnector(host string) *Connector { r.mu.RLock() defer r.mu.RUnlock() @@ -134,6 +237,10 @@ func (r *Router) GetConnector(host string) *Connector { return rw.Next() } +// DelConnector removes a connector identified by its host and connector ID. +// If the removed connector was the last one for the host, the host entry +// is deleted from the map to prevent accumulation of empty slices. +// If no matching connector is found, the call is a no-op. func (r *Router) DelConnector(host string, cid relay.ConnectorID) { r.mu.Lock() defer r.mu.Unlock() @@ -141,12 +248,23 @@ func (r *Router) DelConnector(host string, cid relay.ConnectorID) { connectors := r.connectors[host] for i, c := range connectors { if c.id.Equal(cid) { - r.connectors[host] = append(connectors[:i], connectors[i+1:]...) + connectors = append(connectors[:i], connectors[i+1:]...) break } } + + if len(connectors) == 0 { + delete(r.connectors, host) + } else { + r.connectors[host] = connectors + } } +// Close shuts down the router: closes all connectors, clears the map, +// and marks the router as closed. Subsequent calls are no-ops. +// +// The double-close protection uses a select on r.close — under the +// write lock — so it is race-free. func (r *Router) Close() error { r.mu.Lock() defer r.mu.Unlock() @@ -167,12 +285,20 @@ func (r *Router) Close() error { return nil } +// ConnectorPool manages routers keyed by tunnel ID for a single node. +// +// This is the top-level data structure for connector management. Each +// node has one ConnectorPool, and each pool contains one Router per +// active tunnel. +// +// All methods are nil-safe — calling on a nil *ConnectorPool is valid. type ConnectorPool struct { node string routers map[relay.TunnelID]*Router mu sync.RWMutex } +// NewConnectorPool creates a new ConnectorPool for the given node. func NewConnectorPool(node string) *ConnectorPool { p := &ConnectorPool{ node: node, @@ -182,6 +308,9 @@ func NewConnectorPool(node string) *ConnectorPool { return p } +// Add creates or retrieves a Router for the given tunnel ID and adds +// the connector to it. If no router exists for the tunnel ID, one is +// created automatically. func (p *ConnectorPool) Add(rid relay.TunnelID, c *Connector) { p.mu.Lock() defer p.mu.Unlock() @@ -194,6 +323,9 @@ func (p *ConnectorPool) Add(rid relay.TunnelID, c *Connector) { r.AddConnector(c) } +// Get retrieves a connector for the given tunnel ID and host address. +// Returns nil if the pool is nil, the router doesn't exist, or no +// connector matches the host. func (p *ConnectorPool) Get(rid relay.TunnelID, host string) *Connector { if p == nil { return nil @@ -210,6 +342,8 @@ func (p *ConnectorPool) Get(rid relay.TunnelID, host string) *Connector { return r.GetConnector(host) } +// Del removes a connector from a specific router. +// Safe to call on a nil pool. func (p *ConnectorPool) Del(rid relay.TunnelID, host string, cid relay.ConnectorID) { if p == nil { return @@ -226,6 +360,8 @@ func (p *ConnectorPool) Del(rid relay.TunnelID, host string, cid relay.Connector r.DelConnector(host, cid) } +// Close shuts down all routers and clears the pool. Safe to call on a +// nil pool. Subsequent calls are no-ops (delegates to Router.Close). func (p *ConnectorPool) Close() error { if p == nil { return nil @@ -242,6 +378,12 @@ func (p *ConnectorPool) Close() error { return nil } +// parseRouterID converts a UUID string into a relay.TunnelID. +// Returns a zero-value TunnelID if the string is empty or not a valid UUID. +// +// Ingress rules store router identifiers as UUID strings. This function +// bridges the gap between the string form and the binary TunnelID form +// used internally. func parseRouterID(s string) (rid relay.TunnelID) { if s == "" { return