iop/apps/edge/internal/controlplane/connector_test.go

package controlplane

import (
	"context"
	"net"
	"sync/atomic"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	"iop/packages/config"
	iop "iop/proto/gen/iop"
)

// fakeCPParserMap returns a parser for EdgeHelloRequest on the server side.
func fakeCPParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.EdgeHelloRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeHelloRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// startFakeCPServer starts a fake Control Plane TCP server on a free port and
// calls onHello for each EdgeHelloRequest, returning the provided response.
// It returns the port and a stop function.
func startFakeCPServer(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServer freeTCPPort: %v", err)
	}
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return onHello(req), nil
			},
		)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServer Start: %v", err)
	}
	return port, func() {
		cancel()
		_ = srv.Stop()
	}
}

// startFakeCPServerWithClientAccess starts a fake server that also exposes
// each connected client via a channel, so the test can close clients directly.
func startFakeCPServerWithClientAccess(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, <-chan *toki.TcpClient, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServerWithClientAccess freeTCPPort: %v", err)
	}
	clients := make(chan *toki.TcpClient, 16)
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		clients <- client
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return onHello(req), nil
			},
		)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServerWithClientAccess Start: %v", err)
	}
	return port, clients, func() {
		cancel()
		_ = srv.Stop()
	}
}

// wireAddr builds a host:port address string from a port int.
func wireAddr(port int) string {
	return net.JoinHostPort("127.0.0.1", tcpPortStr(port))
}

// TestConnectorDisabledNoOp verifies that Start returns immediately when
// enabled=false, without touching the network.
func TestConnectorDisabledNoOp(t *testing.T) {
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: false, WireAddr: "127.0.0.1:9999"},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start disabled connector: %v", err)
	}
	if c.IsEnabled() {
		t.Fatal("expected IsEnabled()==false")
	}
	c.Stop() // must not panic
}

// TestConnectorEmptyAddrNoOp verifies that Start returns immediately when
// wire_addr is empty even if enabled=true.
func TestConnectorEmptyAddrNoOp(t *testing.T) {
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: true, WireAddr: ""},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start empty-addr connector: %v", err)
	}
	if c.IsEnabled() {
		t.Fatal("expected IsEnabled()==false for empty wire_addr")
	}
	c.Stop()
}

// TestConnectorSendsEdgeHello verifies that the connector dials and sends the
// correct EdgeHelloRequest fields.
func TestConnectorSendsEdgeHello(t *testing.T) {
	helloReceived := make(chan *iop.EdgeHelloRequest, 1)
	port, stop := startFakeCPServer(t, func(req *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		select {
		case helloReceived <- req:
		default:
		}
		return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.2.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	select {
	case req := <-helloReceived:
		if req.GetEdgeId() != "edge-dgx" {
			t.Errorf("edge_id: got %q, want %q", req.GetEdgeId(), "edge-dgx")
		}
		if req.GetEdgeName() != "DGX Edge" {
			t.Errorf("edge_name: got %q, want %q", req.GetEdgeName(), "DGX Edge")
		}
		if req.GetVersion() != "0.2.0" {
			t.Errorf("version: got %q, want %q", req.GetVersion(), "0.2.0")
		}
		if len(req.GetCapabilities()) == 0 {
			t.Error("expected non-empty capabilities")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for EdgeHelloRequest")
	}
}

// TestConnectorRejectsRejectedHello verifies that the connector treats a
// rejected hello as a connect failure and does not stay connected.
func TestConnectorRejectsRejectedHello(t *testing.T) {
	port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		return &iop.EdgeHelloResponse{Accepted: false, Message: "not authorized"}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Give connect time to attempt and fail.
	time.Sleep(400 * time.Millisecond)
	if c.CurrentState() == StateConnected {
		t.Error("expected connector not to be in Connected state after rejected hello")
	}
}

// TestConnectorStopClosesConnection verifies that Stop() terminates the loop
// and closes an active connection.
func TestConnectorStopClosesConnection(t *testing.T) {
	connected := make(chan struct{}, 1)
	port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		select {
		case connected <- struct{}{}:
		default:
		}
		return &iop.EdgeHelloResponse{Accepted: true}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}

	select {
	case <-connected:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for connection")
	}

	c.Stop()
	if c.CurrentState() != StateStopped {
		t.Errorf("expected StateStopped after Stop(), got %v", c.CurrentState())
	}
}

// TestConnectorReconnectsAfterDisconnect verifies that after the server closes
// the connection, the connector detects the disconnect and re-dials, sending a
// second EdgeHello. This directly exercises the reconnect loop.
func TestConnectorReconnectsAfterDisconnect(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1, // fast reconnect for test
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Wait for first hello / connection.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}

	// Server closes the first client connection to trigger a disconnect.
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting server-side client to close")
	}

	// Wait for a second hello; this proves the connector reconnected.
	select {
	case <-helloSeen:
		if helloCount.Load() < 2 {
			t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
		}
	case <-time.After(5 * time.Second):
		t.Fatalf("timeout waiting for reconnect (second EdgeHello); got %d hellos", helloCount.Load())
	}
}

// TestConnectorReconnectsAfterAcceptedThenClosed tests the race/regression
// where the Control Plane closes the connection immediately after sending an
// accepted EdgeHelloResponse. The connector must detect the disconnect even
// though it races with the hello response, and must reconnect.
func TestConnectorReconnectsAfterAcceptedThenClosed(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Wait for first hello, then immediately close the server-side client.
	// This simulates Control Plane closing the connection right after accepting.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting server-side client to close immediately after hello")
	}

	// Connector must reconnect and send a second hello. If the disconnect
	// listener was not pre-registered, connect() would block forever here.
	select {
	case <-helloSeen:
		if helloCount.Load() < 2 {
			t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
		}
	case <-time.After(5 * time.Second):
		t.Fatalf("connector did not reconnect after accepted-then-closed; got %d hellos", helloCount.Load())
	}
}

// TestConnectorStopPreventsReconnect verifies that after Stop() the connector
// does not send additional hellos within a short window.
func TestConnectorStopPreventsReconnect(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}

	// Wait for first hello.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}

	// Stop the connector before closing the server-side client.
	c.Stop()

	// Now close the server-side client.
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	default:
		// client may already be gone after Stop
	}

	// Give a window longer than ReconnectIntervalSec; no new hello should arrive.
	countBefore := helloCount.Load()
	select {
	case <-helloSeen:
		t.Errorf("connector sent hello after Stop(); count went from %d to %d", countBefore, helloCount.Load())
	case <-time.After(2 * time.Second):
		// good; no extra hello
	}
}