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

package wire

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"sort"
	"strconv"
	"sync"
	"time"

	"go.uber.org/zap"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	"iop/packages/go/auth"
	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

	// activeMu guards activeClients, the live proto-socket client per edge_id
	// the Control Plane uses to send typed requests (e.g. status) to a connected
	// Edge. token mirrors the registry connection token so a stale connection's
	// late disconnect does not evict a newer reconnection's client.
	activeMu      sync.Mutex
	activeClients map[string]activeEdgeClient
	projection    func(context.Context) (*iop.PrincipalProjection, error)
	leaseIssuer   func(context.Context, *iop.AcquireLeaseRequest) (*iop.AcquireLeaseResponse, error)
	requirePeerID bool

	refreshMu       sync.Mutex
	refreshInterval time.Duration
	refreshCancel   context.CancelFunc
	refreshDone     chan struct{}
	broadcastGate   chan struct{}
}

type activeEdgeClient struct {
	client *proto_socket.TcpClient
	token  uint64
}

// projectionClient is a snapshotted broadcast target: the Edge id and the live
// proto-socket client captured under activeMu so a fan-out push never touches
// the shared registry while it runs.
type projectionClient struct {
	edgeID string
	client *proto_socket.TcpClient
}

// projectionPushTimeout bounds a single per-Edge projection push. proto-socket's
// SendRequest starts its own request timer only after QueuePacket returns, and a
// TCP write has no deadline, so a non-reading Edge can stall the write forever.
// This deadline is enforced by BroadcastProjection itself, which closes the
// stalled client to unblock the write when it expires.
const projectionPushTimeout = 10 * time.Second

// NewEdgeServer creates an EdgeServer bound to the given TCP listen address.
func NewEdgeServer(listenAddr string, logger *zap.Logger) (*EdgeServer, error) {
	return NewEdgeServerTLS(listenAddr, nil, logger)
}

// NewEdgeServerTLS creates an Edge wire server. A non-nil TLS configuration
// is retained by the proto-socket server for every accepted connection and
// reconnect; there is no plaintext retry path.
func NewEdgeServerTLS(listenAddr string, tlsConfig *tls.Config, 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,
		activeClients: make(map[string]activeEdgeClient),
		requirePeerID: tlsConfig != nil,
		broadcastGate: make(chan struct{}, 1),
	}

	newClient := 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.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(event *iop.EdgeNodeEvent) {
			idMu.Lock()
			id := edgeID
			tok := token
			idMu.Unlock()
			if id == "" {
				return
			}
			registry.RecordNodeEventFromConnection(id, tok, event, time.Now())
		})

		proto_socket.AddListenerTyped[*iop.EdgeCommandEvent](&client.Communicator, func(event *iop.EdgeCommandEvent) {
			idMu.Lock()
			id := edgeID
			tok := token
			idMu.Unlock()
			if id == "" {
				return
			}
			registry.RecordCommandEventFromConnection(id, tok, event, time.Now())
		})

		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
			}
			if es.requirePeerID {
				identity, identityErr := auth.PeerWorkloadIdentity(conn)
				if identityErr != nil || identity.Role != "edge" || identity.Name != req.GetEdgeId() {
					logger.Warn("edge hello rejected: authenticated identity mismatch",
						zap.String("edge_id", req.GetEdgeId()),
					)
					return &iop.EdgeHelloResponse{
						Accepted:           false,
						Protocol:           Protocol,
						ServerTimeUnixNano: now.UnixNano(),
						Reason:             "authenticated edge identity mismatch",
					}, nil
				}
			}

			var projection *iop.PrincipalProjection
			if es.projection != nil {
				var projectionErr error
				projection, projectionErr = es.projection(context.Background())
				if projectionErr != nil {
					return &iop.EdgeHelloResponse{Accepted: false, Protocol: Protocol, ServerTimeUnixNano: now.UnixNano(), Reason: "credential projection unavailable"}, nil
				}
			}

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

			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",
				PrincipalProjection: projection,
			}, nil
		})

		proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.AcquireLeaseRequest) (*iop.AcquireLeaseResponse, error) {
			idMu.Lock()
			id := edgeID
			idMu.Unlock()
			if es.leaseIssuer == nil || id == "" || req.GetEdgeId() != id {
				return &iop.AcquireLeaseResponse{Error: "credential lease binding rejected"}, nil
			}
			return es.leaseIssuer(context.Background(), req)
		})

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

		return client
	}
	if tlsConfig != nil {
		es.server = proto_socket.NewTcpServerTLS(host, port, tlsConfig, newClient)
	} else {
		es.server = proto_socket.NewTcpServer(host, port, newClient)
	}

	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),
	)
	if err := s.server.Start(ctx); err != nil {
		return err
	}
	s.startProjectionRefresh(ctx)
	return nil
}

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

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

// SetCredentialPlane installs the secret-free projection source and lease
// issuer before Start. Both callbacks are reached only from authenticated Edge
// connections when the server itself is configured with mTLS.
func (s *EdgeServer) SetCredentialPlane(projection func(context.Context) (*iop.PrincipalProjection, error), issuer func(context.Context, *iop.AcquireLeaseRequest) (*iop.AcquireLeaseResponse, error), refreshInterval time.Duration) {
	s.projection = projection
	s.leaseIssuer = issuer
	s.refreshInterval = refreshInterval
}

func (s *EdgeServer) startProjectionRefresh(parent context.Context) {
	if s == nil || s.projection == nil || s.refreshInterval <= 0 {
		return
	}
	s.refreshMu.Lock()
	defer s.refreshMu.Unlock()
	if s.refreshCancel != nil {
		return
	}
	ctx, cancel := context.WithCancel(parent)
	done := make(chan struct{})
	s.refreshCancel = cancel
	s.refreshDone = done
	interval := s.refreshInterval
	go func() {
		defer close(done)
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				if err := s.BroadcastProjection(ctx); err != nil && ctx.Err() == nil {
					s.logger.Warn("periodic credential projection refresh failed", zap.Error(err))
				}
			}
		}
	}()
}

func (s *EdgeServer) stopProjectionRefresh() {
	if s == nil {
		return
	}
	s.refreshMu.Lock()
	cancel := s.refreshCancel
	done := s.refreshDone
	s.refreshMu.Unlock()
	if cancel != nil {
		cancel()
	}
	if done != nil {
		<-done
	}
	s.refreshMu.Lock()
	if s.refreshDone == done {
		s.refreshCancel = nil
		s.refreshDone = nil
	}
	s.refreshMu.Unlock()
}

// BroadcastProjection pushes a freshly committed snapshot to every currently
// authenticated Edge. A failed Edge remains on its old generation, where both
// its local fence and the issuer's durable generation check fail closed.
func (s *EdgeServer) BroadcastProjection(ctx context.Context) error {
	return s.broadcastProjection(ctx, projectionPushTimeout)
}

func (s *EdgeServer) acquireProjectionBatch(ctx context.Context) error {
	select {
	case s.broadcastGate <- struct{}{}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *EdgeServer) releaseProjectionBatch() {
	<-s.broadcastGate
}

// broadcastProjection snapshots the active clients under activeMu, then pushes
// the projection to every target concurrently so no Edge waits behind another.
// A private context-aware batch gate still serializes whole batches. Each push
// is bounded by pushTimeout and by ctx: a stalled or cancelled push closes its
// client to unblock the socket write, waits for the request goroutine, and
// reports an Edge-qualified error, so neither the batch nor EdgeServer.Stop can
// be held open by a non-reading Edge.
func (s *EdgeServer) broadcastProjection(ctx context.Context, pushTimeout time.Duration) error {
	if s == nil || s.projection == nil {
		return fmt.Errorf("credential projection unavailable")
	}
	if err := s.acquireProjectionBatch(ctx); err != nil {
		return err
	}
	defer s.releaseProjectionBatch()
	if err := ctx.Err(); err != nil {
		return err
	}
	projection, err := s.projection(ctx)
	if err != nil {
		return err
	}
	s.activeMu.Lock()
	clients := make([]projectionClient, 0, len(s.activeClients))
	for edgeID, entry := range s.activeClients {
		clients = append(clients, projectionClient{edgeID: edgeID, client: entry.client})
	}
	s.activeMu.Unlock()
	sort.Slice(clients, func(i, j int) bool {
		return clients[i].edgeID < clients[j].edgeID
	})

	failures := make([]error, len(clients))
	var sends sync.WaitGroup
	for index, target := range clients {
		sends.Add(1)
		go func() {
			defer sends.Done()
			failures[index] = s.pushProjection(ctx, target, projection, pushTimeout)
		}()
	}
	sends.Wait()
	return errors.Join(failures...)
}

// pushProjection sends the projection to one Edge and bounds it. The typed
// request runs on its own goroutine; the caller selects on that result, ctx, and
// a per-client deadline. On cancellation or deadline it closes the target client
// so a blocked queue/write returns, then waits for the request goroutine to exit
// before returning an Edge-qualified error. It never leaves a send goroutine or a
// pending request behind. A closed client is removed by the token-fenced
// disconnect listener.
func (s *EdgeServer) pushProjection(ctx context.Context, target projectionClient, projection *iop.PrincipalProjection, pushTimeout time.Duration) error {
	result := make(chan error, 1)
	go func() {
		resp, sendErr := proto_socket.SendRequestTyped[*iop.PrincipalProjectionApplyRequest, *iop.PrincipalProjectionApplyResponse](
			&target.client.Communicator, &iop.PrincipalProjectionApplyRequest{Projection: projection}, pushTimeout)
		switch {
		case sendErr != nil:
			result <- fmt.Errorf("push credential projection to edge %q: %w", target.edgeID, sendErr)
		case !resp.GetApplied():
			result <- fmt.Errorf("push credential projection to edge %q rejected", target.edgeID)
		default:
			result <- nil
		}
	}()

	timer := time.NewTimer(pushTimeout)
	defer timer.Stop()

	select {
	case err := <-result:
		return err
	case <-ctx.Done():
		_ = target.client.Close()
		<-result
		return fmt.Errorf("push credential projection to edge %q: %w", target.edgeID, ctx.Err())
	case <-timer.C:
		_ = target.client.Close()
		<-result
		return fmt.Errorf("push credential projection to edge %q: projection push timed out after %s", target.edgeID, pushTimeout)
	}
}

// setActiveClient records the live proto-socket client for a connected Edge so
// the Control Plane can send typed requests back to it. A newer connection
// always supersedes an older one for the same edge_id.
func (s *EdgeServer) setActiveClient(edgeID string, token uint64, client *proto_socket.TcpClient) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	s.activeClients[edgeID] = activeEdgeClient{client: client, token: token}
}

// clearActiveClient removes the active client for an Edge, but only when token
// still matches the tracked connection so a stale disconnect cannot evict a
// newer reconnection.
func (s *EdgeServer) clearActiveClient(edgeID string, token uint64) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	if entry, ok := s.activeClients[edgeID]; ok && entry.token == token {
		delete(s.activeClients, edgeID)
	}
}

// activeClient returns the live proto-socket client for a connected Edge.
func (s *EdgeServer) activeClient(edgeID string) (*proto_socket.TcpClient, bool) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	entry, ok := s.activeClients[edgeID]
	if !ok {
		return nil, false
	}
	return entry.client, true
}

// RequestStatus asks a connected Edge for its Edge-owned node snapshot over the
// wire and returns the typed response. It observes only what the Edge reports;
// it does not connect to or schedule Node directly. It fails when the Edge is
// not currently connected.
func (s *EdgeServer) RequestStatus(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
	client, ok := s.activeClient(edgeID)
	if !ok {
		return nil, fmt.Errorf("edge %q is not connected", edgeID)
	}
	return proto_socket.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
		&client.Communicator,
		&iop.EdgeStatusRequest{RequestId: fmt.Sprintf("edge-status-%d", time.Now().UnixNano())},
		timeout,
	)
}

// SendCommand sends a command request to a connected Edge over the wire, records
// the request and result in the audit registry, and returns the Edge's typed
// response. It routes only to the target Edge's active connection and fails when
// that Edge is not currently connected. Command lifecycle events the Edge relays
// during execution are recorded asynchronously by the connection's event
// listener. It observes and controls through the Edge; it does not reach Node
// directly.
func (s *EdgeServer) SendCommand(edgeID, operation, targetSelector string, parameters map[string]string, timeout time.Duration) (*iop.EdgeCommandResponse, error) {
	client, ok := s.activeClient(edgeID)
	if !ok {
		return nil, fmt.Errorf("edge %q is not connected", edgeID)
	}
	now := time.Now()
	req := &iop.EdgeCommandRequest{
		RequestId:      fmt.Sprintf("edge-command-%d", now.UnixNano()),
		CommandId:      fmt.Sprintf("cmd-%d", now.UnixNano()),
		Operation:      operation,
		TargetSelector: targetSelector,
		Parameters:     parameters,
	}
	s.registry.RecordCommandRequest(edgeID, req, now)

	resp, err := proto_socket.SendRequestTyped[*iop.EdgeCommandRequest, *iop.EdgeCommandResponse](
		&client.Communicator,
		req,
		timeout,
	)
	if err != nil {
		s.registry.RecordCommandResult(edgeID, &iop.EdgeCommandResponse{
			CommandId: req.GetCommandId(),
			EdgeId:    edgeID,
			Status:    "error",
			Error:     err.Error(),
		}, time.Now())
		return nil, err
	}
	if resp.GetCommandId() == "" {
		resp.CommandId = req.GetCommandId()
	}
	if resp.GetEdgeId() == "" {
		resp.EdgeId = edgeID
	}
	s.registry.RecordCommandResult(edgeID, resp, time.Now())
	return resp, nil
}

// 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 }