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

package wire

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"time"

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

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

// ClientServer wraps the proto-socket WebSocket server for Client-Control Plane communication.
type ClientServer struct {
	server *proto_socket.WsServer
	logger *zap.Logger
	host   string
	port   int
}

// NewClientServer creates a new ClientServer instance listening on the specified address.
func NewClientServer(listenAddr string, logger *zap.Logger) (*ClientServer, error) {
	host, portStr, err := net.SplitHostPort(listenAddr)
	if err != nil {
		return nil, fmt.Errorf("invalid listen address %q: %w", listenAddr, err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("invalid port in listen address %q: %w", listenAddr, err)
	}

	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)
		},
	}

	opts := proto_socket.WsServerOptions{
		AcceptOptions: &websocket.AcceptOptions{
			OriginPatterns: []string{
				"localhost",
				"localhost:*",
				"127.0.0.1",
				"127.0.0.1:*",
			},
		},
	}

	wsServer := proto_socket.NewWsServerWithOptions(host, port, "/client", opts, func(conn *websocket.Conn) *proto_socket.WsClient {
		client := proto_socket.NewWsClient(conn, proto_socket.DefaultHeartbeatIntervalSec, proto_socket.DefaultHeartbeatWaitSec, parserMap)

		proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.ClientHelloRequest) (*iop.ClientHelloResponse, error) {
			logger.Info("client hello request received",
				zap.String("client_id", req.ClientId),
				zap.String("client_version", req.ClientVersion),
			)
			return &iop.ClientHelloResponse{
				Ready:              true,
				Protocol:           Protocol,
				ServerTimeUnixNano: time.Now().UnixNano(),
				Message:            "Hello from IOP Control Plane",
			}, nil
		})

		return client
	})

	wsServer.OnClientConnected = func(client *proto_socket.WsClient) {
		logger.Info("client connected via proto-socket WS")
	}

	return &ClientServer{
		server: wsServer,
		logger: logger,
		host:   host,
		port:   port,
	}, nil
}

// Start starts the client wire WebSocket server.
func (s *ClientServer) Start(ctx context.Context) error {
	s.logger.Info("starting client wire WS server",
		zap.String("host", s.host),
		zap.Int("port", s.port),
		zap.String("path", "/client"),
	)
	return s.server.Start(ctx)
}

// Stop stops the client wire WebSocket server.
func (s *ClientServer) Stop() error {
	s.logger.Info("stopping client wire WS server")
	return s.server.Stop()
}

// Host returns the host the server is listening on.
func (s *ClientServer) Host() string {
	return s.host
}

// Port returns the port the server is listening on.
func (s *ClientServer) Port() int {
	return s.port
}