test(handler/router): add associate_test.go (handleAssociate, handlePacket, getRoute, getAddrforRoute, sdRenew)

2026-06-04 08:46:34 +08:00
parent ca83acc323
commit 7e9d14bb41
2 changed files with 402 additions and 9 deletions
@@ -0,0 +1,385 @@
+package router
+
+import (
+	"bytes"
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/go-gost/core/ingress"
+	"github.com/go-gost/core/router"
+	"github.com/go-gost/core/sd"
+	"github.com/go-gost/relay"
+	"github.com/go-gost/x/internal/util/cache"
+)
+
+// ---------------------------------------------------------------------------
+// handleAssociate tests
+// ---------------------------------------------------------------------------
+
+func TestHandleAssociate_NoIngress(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.ingress = nil
+
+	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
+	reqData := buildRelayAssociateRequest(t, "10.0.0.1:0", rid, "ip")
+
+	conn := &fakeConn{buf: reqData}
+	err := h.handleAssociate(context.Background(), conn, "ip", "10.0.0.1", rid, &testLogger{})
+	// The error comes from the packet read loop when the fake conn has no
+	// more data — it is not an associate failure.
+	if err != nil && err.Error() != "unexpected EOF" {
+		t.Fatalf("handleAssociate: %v", err)
+	}
+
+	// Response should contain StatusOK.
+	resp := readRelayResponse(t, conn.writeBuf.Bytes())
+	if resp.Status != relay.StatusOK {
+		t.Errorf("status = %d, want StatusOK", resp.Status)
+	}
+}
+
+func TestHandleAssociate_IngressMatch(t *testing.T) {
+	h := newInitdHandler(t)
+	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
+
+	// The ingress endpoint UUID must parse to the same bytes as the routerID.
+	// routerID bytes: 0x30,0x31,0x32,0x33,0x34,0x35,0x36,0x37,0x38,0x39,0x61,0x62,0x63,0x64,0x65,0x66
+	h.md.ingress = &mockIngress{
+		getRuleFn: func(ctx context.Context, host string, opts ...ingress.Option) *ingress.Rule {
+			return &ingress.Rule{
+				Hostname: "10.0.0.1",
+				Endpoint: "30313233-3435-3637-3839-616263646566",
+			}
+		},
+	}
+
+	reqData := buildRelayAssociateRequest(t, "10.0.0.1:0", rid, "ip")
+	conn := &fakeConn{buf: reqData}
+	err := h.handleAssociate(context.Background(), conn, "ip", "10.0.0.1", rid, &testLogger{})
+	// The error comes from the packet read loop when the fake conn has no
+	// more data — it is not an associate failure.
+	if err != nil && err.Error() != "unexpected EOF" {
+		t.Fatalf("handleAssociate: %v", err)
+	}
+
+	resp := readRelayResponse(t, conn.writeBuf.Bytes())
+	if resp.Status != relay.StatusOK {
+		t.Errorf("status = %d, want StatusOK", resp.Status)
+	}
+}
+
+func TestHandleAssociate_IngressMismatch(t *testing.T) {
+	h := newInitdHandler(t)
+	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
+
+	// Ingress rule points to a DIFFERENT router ID.
+	h.md.ingress = &mockIngress{
+		getRuleFn: func(ctx context.Context, host string, opts ...ingress.Option) *ingress.Rule {
+			return &ingress.Rule{
+				Hostname: "10.0.0.1",
+				Endpoint: "ffffffff-ffff-ffff-ffff-ffffffffffff",
+			}
+		},
+	}
+
+	reqData := buildRelayAssociateRequest(t, "10.0.0.1:0", rid, "ip")
+	conn := &fakeConn{buf: reqData}
+	err := h.handleAssociate(context.Background(), conn, "ip", "10.0.0.1", rid, &testLogger{})
+	if err == nil {
+		t.Fatal("expected error for ingress mismatch")
+	}
+
+	// Should have written StatusHostUnreachable.
+	resp := readRelayResponse(t, conn.writeBuf.Bytes())
+	if resp.Status != relay.StatusHostUnreachable {
+		t.Errorf("status = %d, want HostUnreachable", resp.Status)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// handlePacket tests
+// ---------------------------------------------------------------------------
+
+func TestHandlePacket_IPv4_ConnectorForward(t *testing.T) {
+	var buf bytes.Buffer
+	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
+	cid := relay.NewConnectorID([]byte("aaaaaaaaaaaaaaaa"))
+	c := NewConnector(rid, cid, "10.0.0.1", &buf, nil)
+
+	h := newInitdHandler(t)
+	h.pool.Add(rid, c)
+	h.md.routerCacheEnabled = false
+
+	// Provide a fallback router so getRoute returns a route whose gateway
+	// matches the connector's host.
+	h.md.router = &mockRouter{
+		getRouteFn: func(ctx context.Context, dst string, opts ...router.Option) *router.Route {
+			return &router.Route{
+				Dst:     "10.0.0.0/24",
+				Gateway: "10.0.0.1",
+			}
+		},
+	}
+
+	pkt := buildIPv4Packet("10.0.0.2", "10.0.0.1", []byte("hello"))
+	err := h.handlePacket(context.Background(), pkt, rid, &testLogger{})
+	if err != nil {
+		t.Fatalf("handlePacket: %v", err)
+	}
+	if buf.Len() == 0 {
+		t.Error("no data written to connector")
+	}
+}
+
+func TestHandlePacket_IPv6_ConnectorForward(t *testing.T) {
+	var buf bytes.Buffer
+	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
+	cid := relay.NewConnectorID([]byte("aaaaaaaaaaaaaaaa"))
+	c := NewConnector(rid, cid, "2001:db8::1", &buf, nil)
+
+	h := newInitdHandler(t)
+	h.pool.Add(rid, c)
+	h.md.routerCacheEnabled = false
+
+	// Provide a fallback router so getRoute returns a route whose gateway
+	// matches the connector's host.
+	h.md.router = &mockRouter{
+		getRouteFn: func(ctx context.Context, dst string, opts ...router.Option) *router.Route {
+			return &router.Route{
+				Dst:     "2001:db8::/32",
+				Gateway: "2001:db8::1",
+			}
+		},
+	}
+
+	pkt := buildIPv6Packet("2001:db8::2", "2001:db8::1", []byte("hello"))
+	err := h.handlePacket(context.Background(), pkt, rid, &testLogger{})
+	if err != nil {
+		t.Fatalf("handlePacket: %v", err)
+	}
+	if buf.Len() == 0 {
+		t.Error("no data written to connector")
+	}
+}
+
+func TestHandlePacket_Unknown(t *testing.T) {
+	h := newInitdHandler(t)
+	err := h.handlePacket(context.Background(), []byte{0, 1, 2, 3}, relay.TunnelID{}, &testLogger{})
+	if err == nil {
+		t.Fatal("expected error for unknown packet type")
+	}
+}
+
+func TestHandlePacket_NoRoute(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.routerCacheEnabled = false
+
+	pkt := buildIPv4Packet("10.0.0.2", "192.168.1.1", []byte("data"))
+	err := h.handlePacket(context.Background(), pkt, relay.NewTunnelID([]byte("0123456789abcdef")), &testLogger{})
+	if err == nil {
+		t.Fatal("expected error for no route")
+	}
+}
+
+func TestHandlePacket_NoConnector_EntrypointFallback(t *testing.T) {
+	// When no connector is found for the destination, handlePacket falls
+	// through to getAddrforRoute and writes via epConn.
+	h := newInitdHandler(t)
+	h.md.routerCacheEnabled = false
+
+	// Provide a fallback router so getRoute returns a route.
+	h.md.router = &mockRouter{
+		getRouteFn: func(ctx context.Context, dst string, opts ...router.Option) *router.Route {
+			return &router.Route{
+				Dst:     "10.0.0.0/24",
+				Gateway: "10.0.0.254",
+			}
+		},
+	}
+
+	// Provide an SD so getAddrforRoute returns an address.
+	h.md.sd = &mockSD{
+		getFn: func(ctx context.Context, name string) ([]*sd.Service, error) {
+			return []*sd.Service{
+				{ID: "svc-1", Name: name, Node: "other-node", Network: "udp", Address: "10.0.0.99:8080"},
+			}, nil
+		},
+	}
+
+	// Create a real epConn so WriteTo succeeds.
+	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("ListenPacket: %v", err)
+	}
+	defer pc.Close()
+	h.epConn = pc
+
+	pkt := buildIPv4Packet("10.0.0.2", "10.0.0.100", []byte("via-ep"))
+	err = h.handlePacket(context.Background(), pkt, relay.NewTunnelID([]byte("0123456789abcdef")), &testLogger{})
+	if err != nil {
+		t.Fatalf("handlePacket: %v", err)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// getRoute tests
+// ---------------------------------------------------------------------------
+
+func TestGetRoute_CacheHit(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.routerCacheEnabled = true
+	h.md.routerCacheExpiration = time.Minute
+	h.routeCache.Set("10.0.0.1", cache.NewItem(&router.Route{
+		Dst:     "10.0.0.0/24",
+		Gateway: "10.0.0.254",
+	}, time.Minute))
+
+	route := h.getRoute(context.Background(), "test-rid", "10.0.0.1")
+	if route == nil {
+		t.Fatal("getRoute returned nil")
+	}
+	if route.Gateway != "10.0.0.254" {
+		t.Errorf("gateway = %s, want 10.0.0.254", route.Gateway)
+	}
+}
+
+func TestGetRoute_CacheDisabled(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.routerCacheEnabled = false
+
+	// With no registry router and no fallback, should return nil.
+	route := h.getRoute(context.Background(), "test-rid", "10.0.0.1")
+	if route != nil {
+		t.Log("route is not nil (may have registry router)")
+	}
+}
+
+func TestGetRoute_FallbackRouter(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.routerCacheEnabled = false
+	h.md.router = &mockRouter{
+		getRouteFn: func(ctx context.Context, dst string, opts ...router.Option) *router.Route {
+			return &router.Route{
+				Dst:     "10.0.0.0/24",
+				Gateway: "10.0.0.254",
+			}
+		},
+	}
+
+	route := h.getRoute(context.Background(), "nonexistent-rid", "10.0.0.1")
+	if route == nil {
+		t.Fatal("getRoute returned nil")
+	}
+	if route.Gateway != "10.0.0.254" {
+		t.Errorf("gateway = %s, want 10.0.0.254", route.Gateway)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// getAddrforRoute tests
+// ---------------------------------------------------------------------------
+
+func TestGetAddrforRoute_NilSD(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.sd = nil
+
+	addr := h.getAddrforRoute(context.Background(), "rid", "gateway")
+	if addr != nil {
+		t.Errorf("getAddrforRoute = %v, want nil", addr)
+	}
+}
+
+func TestGetAddrforRoute_CacheHit(t *testing.T) {
+	h := newInitdHandler(t)
+	h.md.sd = &mockSD{}
+	h.md.sdCacheExpiration = time.Minute
+
+	expectedAddr := &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 8080}
+	h.sdCache.Set("gateway@rid", cache.NewItem(expectedAddr, time.Minute))
+
+	addr := h.getAddrforRoute(context.Background(), "rid", "gateway")
+	if addr == nil {
+		t.Fatal("getAddrforRoute returned nil")
+	}
+	if addr.String() != expectedAddr.String() {
+		t.Errorf("addr = %s, want %s", addr.String(), expectedAddr.String())
+	}
+}
+
+func TestGetAddrforRoute_SDLookup(t *testing.T) {
+	h := newInitdHandler(t)
+	h.id = "self-node"
+	h.md.sdCacheExpiration = time.Minute
+
+	h.md.sd = &mockSD{
+		getFn: func(ctx context.Context, name string) ([]*sd.Service, error) {
+			return []*sd.Service{
+				{ID: "conn-1", Name: name, Node: "self-node", Network: "udp", Address: "10.0.0.1:8080"},
+				{ID: "conn-2", Name: name, Node: "other-node", Network: "udp", Address: "10.0.0.2:8080"},
+			}, nil
+		},
+	}
+
+	addr := h.getAddrforRoute(context.Background(), "rid", "gateway")
+	if addr == nil {
+		t.Fatal("getAddrforRoute returned nil")
+	}
+	// Should have skipped self-node and returned other-node.
+	if addr.String() != "10.0.0.2:8080" {
+		t.Errorf("addr = %s, want 10.0.0.2:8080", addr.String())
+	}
+}
+
+// ---------------------------------------------------------------------------
+// sdRenew tests
+// ---------------------------------------------------------------------------
+
+func TestSdRenew_Cancel(t *testing.T) {
+	h := newInitdHandler(t)
+
+	renewCalled := make(chan struct{}, 1)
+	h.md.sd = &mockSD{
+		renewFn: func(ctx context.Context, svc *sd.Service) error {
+			renewCalled <- struct{}{}
+			return nil
+		},
+	}
+	h.md.sdRenewInterval = 50 * time.Millisecond
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go h.sdRenew(ctx, "client-id", "connector-id")
+
+	// Let it tick once.
+	select {
+	case <-renewCalled:
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for Renew")
+	}
+
+	cancel()
+}
+
+func TestSdRenew_NormalTick(t *testing.T) {
+	h := newInitdHandler(t)
+
+	var tickCount int
+	h.md.sd = &mockSD{
+		renewFn: func(ctx context.Context, svc *sd.Service) error {
+			tickCount++
+			return nil
+		},
+	}
+	h.md.sdRenewInterval = 50 * time.Millisecond
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go h.sdRenew(ctx, "client-id", "connector-id")
+
+	time.Sleep(120 * time.Millisecond)
+	if tickCount == 0 {
+		t.Error("Renew was never called")
+	}
+}