iop/apps/edge/internal/transport/server.go

package transport

import (
	"context"
	"crypto/tls"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

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

	"iop/apps/edge/internal/configrefresh"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

const (
	// Match the Node's tunnel-tolerant liveness profile. Long provider prefill and
	// streamed-response backpressure must not fence a healthy connection before
	// it can service the next heartbeat.
	heartbeatIntervalSec = 30
	// heartbeatWaitSec mirrors the node side. See the comment in
	// apps/node/internal/transport/client.go for rationale.
	heartbeatWaitSec = 45
)

const configRefreshTimeoutSec = 30

func edgeParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeNodeEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeNodeEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.ProviderTunnelFrame{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelFrame{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RegisterRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeReadyRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeReadyRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeConfigRefreshResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// Server wraps proto-socket TcpServer and manages node connections.
type Server struct {
	tcp           *toki.TcpServer
	listen        string
	registry      *edgenode.Registry
	nodeStoreMu   sync.RWMutex
	nodeStore     *edgenode.NodeStore
	logger        *zap.Logger
	handlerMu     sync.RWMutex
	onRunEvent    func(*iop.RunEvent)
	onNodeEvent   func(*iop.EdgeNodeEvent)
	onTunnelFrame func(*iop.ProviderTunnelFrame)
	// onRunLifecycle, onNodeConnect, and onNodeDisconnect are the authoritative
	// lifecycle hooks. They run synchronously, ahead of the observability fanout,
	// so resource accounting never depends on a bus delivery that is allowed to
	// drop.
	onRunLifecycle    func(*iop.RunEvent)
	onNodeConnect     func(nodeID string, generation uint64)
	onNodeDisconnect  func(nodeID string, generation uint64, reason string)
	peerMu            sync.RWMutex
	peerConnections   map[*toki.TcpClient]net.Conn
	requirePeerID     bool
	stopping          atomic.Bool
	HeartbeatInterval int
	HeartbeatWait     int
}

func NewServer(listen string, registry *edgenode.Registry, nodeStore *edgenode.NodeStore, logger *zap.Logger) (*Server, error) {
	return NewServerTLS(listen, nil, registry, nodeStore, logger)
}

// NewServerTLS creates the Edge-Node server with optional mTLS. Once supplied,
// the TLS configuration is used for every accepted connection; the server does
// not maintain a plaintext fallback listener.
func NewServerTLS(listen string, tlsConfig *tls.Config, registry *edgenode.Registry, nodeStore *edgenode.NodeStore, logger *zap.Logger) (*Server, error) {
	host, portStr, err := net.SplitHostPort(listen)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}

	s := &Server{
		listen:            listen,
		registry:          registry,
		nodeStore:         nodeStore,
		logger:            logger,
		peerConnections:   make(map[*toki.TcpClient]net.Conn),
		requirePeerID:     tlsConfig != nil,
		HeartbeatInterval: heartbeatIntervalSec,
		HeartbeatWait:     heartbeatWaitSec,
	}
	newClient := func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, s.HeartbeatInterval, s.HeartbeatWait, edgeParserMap())
		if tlsConfig != nil {
			s.peerMu.Lock()
			s.peerConnections[client] = conn
			s.peerMu.Unlock()
			client.AddDisconnectListener(func(disconnected *toki.TcpClient) {
				s.peerMu.Lock()
				delete(s.peerConnections, disconnected)
				s.peerMu.Unlock()
			})
			// NewTcpClient starts its read loop immediately. If the peer closed in
			// the tiny window before the listener above was attached, remove the
			// retained identity source here as well.
			if !client.IsAlive() {
				s.peerMu.Lock()
				delete(s.peerConnections, client)
				s.peerMu.Unlock()
			}
		}
		return client
	}
	if tlsConfig != nil {
		s.tcp = toki.NewTcpServerTLS(host, port, tlsConfig, newClient)
	} else {
		s.tcp = toki.NewTcpServer(host, port, newClient)
	}
	s.tcp.OnClientConnected = s.onNodeConnected
	return s, nil
}

func (s *Server) Start(ctx context.Context) error {
	s.stopping.Store(false)
	if err := s.tcp.Start(ctx); err != nil {
		return err
	}
	s.logger.Info("edge listening for nodes", zap.String("addr", s.listen))
	return nil
}

func (s *Server) SetNodeStore(store *edgenode.NodeStore) {
	s.nodeStoreMu.Lock()
	s.nodeStore = store
	s.nodeStoreMu.Unlock()
}

func (s *Server) nodeStoreSnapshot() *edgenode.NodeStore {
	s.nodeStoreMu.RLock()
	defer s.nodeStoreMu.RUnlock()
	return s.nodeStore
}

func (s *Server) Stop() error {
	s.stopping.Store(true)
	err := s.tcp.Stop()
	s.peerMu.Lock()
	clear(s.peerConnections)
	s.peerMu.Unlock()
	return err
}

func (s *Server) peerConnection(client *toki.TcpClient) (net.Conn, bool) {
	s.peerMu.RLock()
	defer s.peerMu.RUnlock()
	conn, ok := s.peerConnections[client]
	return conn, ok
}

func (s *Server) SetRunEventHandler(handler func(*iop.RunEvent)) {
	s.handlerMu.Lock()
	s.onRunEvent = handler
	s.handlerMu.Unlock()
}

func (s *Server) SetNodeEventHandler(handler func(*iop.EdgeNodeEvent)) {
	s.handlerMu.Lock()
	s.onNodeEvent = handler
	s.handlerMu.Unlock()
}

// SetTunnelFrameHandler registers the request-bound ProviderTunnelFrame
// handler. Tunnel frames carry raw provider passthrough bytes and are routed
// to a per-request channel by the handler; they must never be published to
// the run event bus.
func (s *Server) SetTunnelFrameHandler(handler func(*iop.ProviderTunnelFrame)) {
	s.handlerMu.Lock()
	s.onTunnelFrame = handler
	s.handlerMu.Unlock()
}

// SetRunLifecycleHandler registers the authoritative run lifecycle handler. It
// is invoked for every run event, before the observability handler, so the
// service can settle terminal accounting regardless of event bus delivery.
func (s *Server) SetRunLifecycleHandler(handler func(*iop.RunEvent)) {
	s.handlerMu.Lock()
	s.onRunLifecycle = handler
	s.handlerMu.Unlock()
}

// SetNodeConnectHandler registers the authoritative node connect handler. It is
// invoked only for an accepted (non-duplicate) registration, carrying that
// accepted connection's generation, before the connected event is emitted, so a
// reconnect re-activates the node's provider resources and pumps its stranded
// waiters regardless of event bus delivery.
func (s *Server) SetNodeConnectHandler(handler func(nodeID string, generation uint64)) {
	s.handlerMu.Lock()
	s.onNodeConnect = handler
	s.handlerMu.Unlock()
}

// SetNodeDisconnectHandler registers the authoritative node disconnect handler.
// It is invoked only for the client that still owned the registry entry, carrying
// that owner's connection generation, before the disconnected event is emitted.
func (s *Server) SetNodeDisconnectHandler(handler func(nodeID string, generation uint64, reason string)) {
	s.handlerMu.Lock()
	s.onNodeDisconnect = handler
	s.handlerMu.Unlock()
}

func (s *Server) HasRunLifecycleHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onRunLifecycle != nil
}

func (s *Server) HasNodeConnectHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onNodeConnect != nil
}

func (s *Server) HasNodeDisconnectHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onNodeDisconnect != nil
}

func (s *Server) HasTunnelFrameHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onTunnelFrame != nil
}

func (s *Server) HasRunEventHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onRunEvent != nil
}

func (s *Server) HasNodeEventHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onNodeEvent != nil
}

// onNodeConnected wires the per-connection listeners in a fixed order: run
// events, provider tunnel frames, the registration request handler, then the
// dispatch-ready request handler. The ready listener is bound at connect time so
// it is already installed when the node sends NodeReadyRequest after applying its
// config. The concrete listener bodies live in connection_handlers.go.
func (s *Server) onNodeConnected(client *toki.TcpClient) {
	s.logger.Info("node connection established")
	s.registerRunEventListener(client)
	s.registerTunnelFrameListener(client)
	s.registerNodeRequestListener(client)
	s.registerNodeReadyListener(client)
}

// PushConfigRefresh sends a NodeConfigRefreshRequest to every configured node
// and collects per-node results. Nodes that are not in the live registry or
// have no alive client are recorded as skipped. Connected nodes receive a
// node-specific payload built from the refreshed NodeStore record. Transport
// failures are recorded as failed and do not roll back the local config commit.
func (s *Server) PushConfigRefresh(ctx context.Context, req *iop.NodeConfigRefreshRequest) []configrefresh.NodeResult {
	ns := s.nodeStoreSnapshot()
	var records []*edgenode.NodeRecord
	if ns != nil {
		records = ns.All()
	}
	results := make([]configrefresh.NodeResult, 0, len(records))
	timeout := time.Duration(configRefreshTimeoutSec) * time.Second
	for _, rec := range records {
		nr := configrefresh.NodeResult{
			NodeID: rec.ID,
			Alias:  rec.Alias,
		}
		// GetReady, not Get: a pending accepted connection has not installed its
		// handler, so a config-refresh push would be dropped by the node. Skip it
		// until it signals readiness — the node applies the latest config from its
		// RegisterResponse on connect anyway.
		entry, ok := s.registry.GetReady(rec.ID)
		if !ok || entry.Client == nil || !entry.Client.IsAlive() {
			nr.Status = "skipped"
			results = append(results, nr)
			continue
		}
		nodeReq := &iop.NodeConfigRefreshRequest{
			RequestId:    req.GetRequestId(),
			ChangedPaths: req.GetChangedPaths(),
		}
		if payload, err := edgenode.BuildConfigPayload(rec); err == nil {
			nodeReq.Config = payload
		}
		resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
			&entry.Client.Communicator,
			nodeReq,
			timeout,
		)
		if err != nil {
			nr.Status = "failed"
			nr.Error = err.Error()
			results = append(results, nr)
			continue
		}
		switch resp.GetStatus() {
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED:
			nr.Status = "applied"
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_RESTART_REQUIRED:
			nr.Status = "restart_required"
			nr.RestartRequiredPaths = resp.GetRestartRequiredPaths()
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED:
			nr.Status = "failed"
			nr.Error = resp.GetError()
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_SKIPPED:
			nr.Status = "skipped"
		default:
			nr.Status = "skipped"
		}
		results = append(results, nr)
	}
	return results
}

func (s *Server) enrichRunEvent(event *iop.RunEvent) {
	if event == nil || event.GetNodeAlias() != "" || s.registry == nil {
		return
	}
	nodeID := event.GetNodeId()
	if nodeID == "" {
		return
	}
	if entry, ok := s.registry.Get(nodeID); ok && entry.Alias != "" {
		event.NodeAlias = entry.Alias
	}
}

func (s *Server) emitNodeEvent(event *iop.EdgeNodeEvent) {
	s.logger.Debug("node event received",
		zap.String("node_id", event.GetNodeId()),
		zap.String("type", event.GetType()),
		zap.String("reason", event.GetReason()),
	)
	s.handlerMu.RLock()
	handler := s.onNodeEvent
	s.handlerMu.RUnlock()
	if handler != nil {
		handler(event)
	}
}

func safePrefix(s string) string {
	if len(s) > 8 {
		return s[:8] + "..."
	}
	return s
}

func transportDisconnectMetadata(info toki.DisconnectInfo) map[string]string {
	metadata := make(map[string]string, 2)
	if info.Reason != "" {
		metadata[events.MetadataTransportCloseReason] = info.Reason
	}
	if info.Error != "" {
		metadata[events.MetadataTransportCloseError] = info.Error
	}
	if len(metadata) == 0 {
		return nil
	}
	return metadata
}

func transportDisconnectFields(info toki.DisconnectInfo) []zap.Field {
	fields := make([]zap.Field, 0, 2)
	if info.Reason != "" {
		fields = append(fields, zap.String("transport_close_reason", info.Reason))
	}
	if info.Error != "" {
		fields = append(fields, zap.String("transport_close_error", info.Error))
	}
	return fields
}