iop/apps/control-plane/internal/wire/client_test.go

package wire

import (
	"context"
	"fmt"
	"net"
	"testing"
	"time"

	"go.uber.org/zap/zaptest"
	"google.golang.org/protobuf/proto"
	"nhooyr.io/websocket"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	iop "iop/proto/gen/iop"
)

func getFreePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("failed to find free port: %v", err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func TestClientServerHandshake(t *testing.T) {
	logger := zaptest.NewLogger(t)
	port := getFreePort(t)
	listenAddr := fmt.Sprintf("127.0.0.1:%d", port)

	server, err := NewClientServer(listenAddr, logger)
	if err != nil {
		t.Fatalf("NewClientServer failed: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if err := server.Start(ctx); err != nil {
		t.Fatalf("server Start failed: %v", err)
	}
	defer func() {
		_ = server.Stop()
	}()

	// Give a tiny moment for server to bind
	time.Sleep(50 * time.Millisecond)

	// Dial to the client endpoint
	parserMap := proto_socket.ParserMap{
		proto_socket.TypeNameOf(&iop.ClientHelloRequest{}): func(b []byte) (proto.Message, error) {
			req := &iop.ClientHelloRequest{}
			return req, proto.Unmarshal(b, req)
		},
		proto_socket.TypeNameOf(&iop.ClientHelloResponse{}): func(b []byte) (proto.Message, error) {
			res := &iop.ClientHelloResponse{}
			return res, proto.Unmarshal(b, res)
		},
	}

	// Dial using WS client
	client, err := proto_socket.DialWsWithHeartbeat(
		ctx,
		"127.0.0.1",
		port,
		"/client",
		proto_socket.DefaultHeartbeatIntervalSec,
		proto_socket.DefaultHeartbeatWaitSec,
		parserMap,
	)
	if err != nil {
		t.Fatalf("failed to DialWs: %v", err)
	}
	defer client.Close()

	// Send ClientHelloRequest and wait for ClientHelloResponse
	req := &iop.ClientHelloRequest{
		ClientId:      "test-client-id",
		ClientVersion: "1.0.0",
	}

	res, err := proto_socket.SendRequestTyped[*iop.ClientHelloRequest, *iop.ClientHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if !res.Ready {
		t.Errorf("expected response.Ready to be true, got false")
	}
	if res.Protocol != Protocol {
		t.Errorf("expected response.Protocol to be %q, got %q", Protocol, res.Protocol)
	}
	if res.ServerTimeUnixNano <= 0 {
		t.Errorf("expected server time unix nano to be positive, got %d", res.ServerTimeUnixNano)
	}
	if res.Message != "Hello from IOP Control Plane" {
		t.Errorf("expected message to be %q, got %q", "Hello from IOP Control Plane", res.Message)
	}
}

func TestClientServerHandshakeWithBrowserOrigin(t *testing.T) {
	logger := zaptest.NewLogger(t)
	port := getFreePort(t)
	listenAddr := fmt.Sprintf("127.0.0.1:%d", port)

	server, err := NewClientServer(listenAddr, logger)
	if err != nil {
		t.Fatalf("NewClientServer failed: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if err := server.Start(ctx); err != nil {
		t.Fatalf("server Start failed: %v", err)
	}
	defer func() {
		_ = server.Stop()
	}()

	time.Sleep(50 * time.Millisecond)

	parserMap := proto_socket.ParserMap{
		proto_socket.TypeNameOf(&iop.ClientHelloRequest{}): func(b []byte) (proto.Message, error) {
			req := &iop.ClientHelloRequest{}
			return req, proto.Unmarshal(b, req)
		},
		proto_socket.TypeNameOf(&iop.ClientHelloResponse{}): func(b []byte) (proto.Message, error) {
			res := &iop.ClientHelloResponse{}
			return res, proto.Unmarshal(b, res)
		},
	}

	// Dial specifying browser Origin
	dialOpts := &websocket.DialOptions{
		HTTPHeader: map[string][]string{
			"Origin": {"http://localhost:3000"},
		},
	}
	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://127.0.0.1:%d/client", port), dialOpts)
	if err != nil {
		t.Fatalf("failed to Dial with Origin: %v", err)
	}

	// Wrap in WsClient
	client := proto_socket.NewWsClient(conn, proto_socket.DefaultHeartbeatIntervalSec, proto_socket.DefaultHeartbeatWaitSec, parserMap)
	defer client.Close()

	// Send ClientHelloRequest
	req := &iop.ClientHelloRequest{
		ClientId:      "browser-client",
		ClientVersion: "1.0.0",
	}

	res, err := proto_socket.SendRequestTyped[*iop.ClientHelloRequest, *iop.ClientHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if !res.Ready {
		t.Errorf("expected response.Ready to be true, got false")
	}
	if res.Protocol != Protocol {
		t.Errorf("expected response.Protocol to be %q, got %q", Protocol, res.Protocol)
	}
}