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

package wire

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

	"go.uber.org/zap"

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

// EdgeServer hosts the Control Plane-Edge native TCP wire endpoint. It accepts
// Edge connections over proto-socket, handles the hello handshake, and keeps an
// internal connection registry. It does not connect to or schedule Node
// directly; the Control Plane observes and controls the system through Edge.
type EdgeServer struct {
	server   *proto_socket.TcpServer
	registry *EdgeRegistry
	logger   *zap.Logger
	host     string
	port     int
}

// NewEdgeServer creates an EdgeServer bound to the given TCP listen address.
func NewEdgeServer(listenAddr string, logger *zap.Logger) (*EdgeServer, error) {
	host, portStr, err := net.SplitHostPort(listenAddr)
	if err != nil {
		return nil, fmt.Errorf("invalid edge listen address %q: %w", listenAddr, err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("invalid port in edge listen address %q: %w", listenAddr, err)
	}

	registry := NewEdgeRegistry()
	es := &EdgeServer{
		registry: registry,
		logger:   logger,
		host:     host,
		port:     port,
	}

	es.server = proto_socket.NewTcpServer(host, port, func(conn net.Conn) *proto_socket.TcpClient {
		client := proto_socket.NewTcpClient(conn, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, EdgeParserMap())

		// edgeID and token are set once the hello handshake identifies the
		// connection so the disconnect listener can flip the right registry entry.
		// The token lets a stale connection's late disconnect be ignored after the
		// same edge_id reconnected. The hello and disconnect callbacks may run on
		// different goroutines, so guard them.
		var (
			idMu   sync.Mutex
			edgeID string
			token  uint64
		)

		proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
			now := time.Now()
			if req.GetEdgeId() == "" {
				logger.Warn("edge hello rejected: missing edge_id")
				return &iop.EdgeHelloResponse{
					Accepted:           false,
					Protocol:           Protocol,
					ServerTimeUnixNano: now.UnixNano(),
					Reason:             "edge_id is required",
				}, nil
			}

			connToken := registry.MarkConnected(req, now)
			idMu.Lock()
			edgeID = req.GetEdgeId()
			token = connToken
			idMu.Unlock()

			logger.Info("edge hello accepted",
				zap.String("edge_id", req.GetEdgeId()),
				zap.String("edge_name", req.GetEdgeName()),
				zap.String("version", req.GetVersion()),
			)
			return &iop.EdgeHelloResponse{
				Accepted:           true,
				Protocol:           Protocol,
				ServerTimeUnixNano: now.UnixNano(),
				Message:            "Edge enrolled on IOP Control Plane",
			}, nil
		})

		client.AddDisconnectListener(func(c *proto_socket.TcpClient) {
			idMu.Lock()
			id := edgeID
			tok := token
			idMu.Unlock()
			if id == "" {
				return
			}
			info := c.DisconnectInfo()
			registry.MarkDisconnected(id, tok, time.Now(), info)
			logger.Info("edge disconnected",
				zap.String("edge_id", id),
				zap.String("reason", info.Reason),
			)
		})

		return client
	})

	return es, nil
}

// Start begins accepting Edge connections on the configured TCP endpoint.
func (s *EdgeServer) Start(ctx context.Context) error {
	s.logger.Info("starting edge wire TCP server",
		zap.String("host", s.host),
		zap.Int("port", s.port),
		zap.String("transport", EdgeTransport),
	)
	return s.server.Start(ctx)
}

// Stop stops the Edge wire TCP server and closes active connections.
func (s *EdgeServer) Stop() error {
	s.logger.Info("stopping edge wire TCP server")
	return s.server.Stop()
}

// Registry exposes the internal Edge connection registry.
func (s *EdgeServer) Registry() *EdgeRegistry { return s.registry }

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

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