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

package transport

import (
	toki "git.toki-labs.com/toki/proto-socket/go"
	"go.uber.org/zap"

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

// registerRunEventListener forwards node RunEvents to the current run-event
// handler after alias enrichment.
func (s *Server) registerRunEventListener(client *toki.TcpClient) {
	toki.AddListenerTyped[*iop.RunEvent](&client.Communicator, func(e *iop.RunEvent) {
		s.logger.Debug("run event received",
			zap.String("run_id", e.GetRunId()),
			zap.String("type", e.GetType()),
		)
		s.enrichRunEvent(e)
		s.handlerMu.RLock()
		lifecycle := s.onRunLifecycle
		handler := s.onRunEvent
		s.handlerMu.RUnlock()
		// Correctness first: the lifecycle hook settles run accounting
		// synchronously, then the event goes out for observation. Publishing
		// first would make a dropped fanout lose the terminal signal.
		if lifecycle != nil {
			lifecycle(e)
		}
		if handler != nil {
			handler(e)
		}
	})
}

// registerTunnelFrameListener routes raw provider tunnel frames to the current
// tunnel handler, dropping them when none is registered so they never reach the
// run event bus.
func (s *Server) registerTunnelFrameListener(client *toki.TcpClient) {
	toki.AddListenerTyped[*iop.ProviderTunnelFrame](&client.Communicator, func(f *iop.ProviderTunnelFrame) {
		s.handlerMu.RLock()
		handler := s.onTunnelFrame
		s.handlerMu.RUnlock()
		if handler == nil {
			s.logger.Warn("provider tunnel frame dropped: no tunnel handler",
				zap.String("run_id", f.GetRunId()),
				zap.String("tunnel_id", f.GetTunnelId()),
			)
			return
		}
		handler(f)
	})
}

// registerNodeRequestListener installs the token-based RegisterRequest handler.
func (s *Server) registerNodeRequestListener(client *toki.TcpClient) {
	toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		func(req *iop.RegisterRequest) (*iop.RegisterResponse, error) {
			return s.handleRegisterRequest(client, req)
		},
	)
}

// handleRegisterRequest validates the node token, binds the registered-node
// listeners, and registers the node. The event/registry ordering is preserved:
// the node-event and disconnect listeners are bound before RegisterIfAbsent, and
// the connected event is emitted only after a successful registration.
func (s *Server) handleRegisterRequest(client *toki.TcpClient, req *iop.RegisterRequest) (*iop.RegisterResponse, error) {
	nodeStore := s.nodeStoreSnapshot()
	if nodeStore == nil {
		return &iop.RegisterResponse{Accepted: false, Reason: "edge node store unavailable"}, nil
	}
	rec, ok := nodeStore.FindByToken(req.GetToken())
	if !ok {
		s.logger.Warn("unknown token", zap.String("token_prefix", safePrefix(req.GetToken())))
		s.emitNodeEvent(events.NewEdgeNodeEvent(
			events.SourceEdge,
			events.TypeNodeRegistrationFailed,
			"",
			"",
			events.ReasonUnknownToken,
			map[string]string{
				events.MetadataFailureReason: events.ReasonUnknownToken,
				events.MetadataTokenPrefix:   safePrefix(req.GetToken()),
			},
		))
		return &iop.RegisterResponse{Accepted: false, Reason: "unknown token"}, nil
	}
	if s.requirePeerID {
		conn, found := s.peerConnection(client)
		identity, identityErr := auth.PeerWorkloadIdentity(conn)
		if !found || identityErr != nil || identity.Role != "node" || identity.Name != rec.ID {
			s.logger.Warn("node registration rejected: authenticated identity mismatch",
				zap.String("node_id", rec.ID),
			)
			return &iop.RegisterResponse{Accepted: false, Reason: "authenticated node identity mismatch"}, nil
		}
	}

	cfg, err := edgenode.BuildConfigPayload(rec)
	if err != nil {
		s.logger.Error("build config payload failed",
			zap.String("node_id", rec.ID),
			zap.Error(err),
		)
		return &iop.RegisterResponse{Accepted: false, Reason: "internal config error"}, nil
	}

	entry := &edgenode.NodeEntry{
		NodeID:                       rec.ID,
		Alias:                        rec.Alias,
		LifecycleState:               edgenode.LifecycleConnected,
		Client:                       client,
		Index:                        rec.Index,
		HasIndex:                     true,
		CredentialRecipientKeyID:     req.GetCredentialRecipientKeyId(),
		CredentialRecipientPublicKey: append([]byte(nil), req.GetCredentialRecipientPublicKey()...),
	}
	s.bindNodeEventListener(client, rec)
	s.bindDisconnectListener(client, rec)
	if !s.registry.RegisterIfAbsent(entry) {
		return s.rejectDuplicateRegistration(rec), nil
	}
	// Accepted registration claims ownership and delivers config only. It does NOT
	// open dispatch eligibility, pump stranded waiters, or emit the connected
	// event yet: the node has not applied this config or installed its message
	// handler, so any run/tunnel request sent now would be dropped on the node
	// side. The node applies the config from this response, installs its handler,
	// and then sends NodeReadyRequest — the authoritative connect lifecycle and
	// the connected event run there, exactly once, once dispatch is safe.
	s.logger.Info("node registration accepted, awaiting dispatch-ready",
		zap.String("node_id", rec.ID),
		zap.String("alias", rec.Alias),
		zap.Uint64("connection_generation", entry.ConnectionGeneration),
	)

	return &iop.RegisterResponse{
		Accepted: true,
		NodeId:   rec.ID,
		Alias:    rec.Alias,
		Config:   cfg,
	}, nil
}

// registerNodeReadyListener installs the NodeReadyRequest handler. The node sends
// this after applying its config and installing its message handler.
func (s *Server) registerNodeReadyListener(client *toki.TcpClient) {
	toki.AddRequestListenerTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&client.Communicator,
		func(req *iop.NodeReadyRequest) (*iop.NodeReadyResponse, error) {
			return s.handleReadyRequest(client, req), nil
		},
	)
}

// handleReadyRequest transitions the accepted connection identified by
// (req.NodeId, client) to dispatch-ready. Only the current owner of that id — the
// exact client that still holds the registry entry — can ready it, and only the
// first ready runs the authoritative connect lifecycle (which re-activates the
// node's provider resources and pumps its stranded waiters) and emits the
// connected event. A duplicate ready acks success without repeating either; a
// stale ready (superseded by a reconnect, or the entry already gone) is rejected
// so the node closes and reconnects instead of sitting idle believing it is live.
func (s *Server) handleReadyRequest(client *toki.TcpClient, req *iop.NodeReadyRequest) *iop.NodeReadyResponse {
	nodeID := req.GetNodeId()
	entry, transitioned, ok := s.registry.MarkDispatchReadyOwner(nodeID, client)
	if !ok {
		s.logger.Warn("stale node ready signal rejected", zap.String("node_id", nodeID))
		return &iop.NodeReadyResponse{Ready: false, Reason: "connection is not the current registry owner"}
	}
	if !transitioned {
		// Duplicate ready for an already-ready owner: idempotent success, no extra
		// pump or event.
		return &iop.NodeReadyResponse{Ready: true}
	}
	// First ready for this connection. Authoritative connected lifecycle first,
	// then the observability event.
	//
	// Both steps are fenced on this exact connection, but by two gates that share
	// the same generation, not one lock: the activation gate lives inside the
	// queue lock (HandleNodeConnect → activateNode re-checks the current owner
	// generation atomically with the resource mutation and pump, serialized against
	// a disconnect's releaseNode), and the event gate is WithCurrentOwner. They are
	// deliberately not wrapped in a single WithCurrentOwner: notifyNodeConnected
	// activates provider resources whose pump re-resolves candidates against the
	// registry (queue-lock → registry-lock order), so holding the registry lock
	// across it would invert that order and deadlock. Each gate independently makes
	// a stale (superseded or already-disconnected) generation a no-op, so a
	// reconnect that supersedes this connection between the two steps still yields
	// zero activation and zero connected event for the stale generation, and
	// exactly one of each for the current owner.
	s.notifyNodeConnected(nodeID, entry.ConnectionGeneration)
	s.registry.WithCurrentOwner(entry, func() {
		s.emitNodeReady(entry)
	})
	return &iop.NodeReadyResponse{Ready: true}
}

// rejectDuplicateRegistration logs and emits the duplicate-connection failure
// event and returns the rejection response.
func (s *Server) rejectDuplicateRegistration(rec *edgenode.NodeRecord) *iop.RegisterResponse {
	s.logger.Warn("duplicate registration rejected",
		zap.String("node_id", rec.ID),
	)
	s.emitNodeEvent(events.NewEdgeNodeEvent(
		events.SourceEdge,
		events.TypeNodeRegistrationFailed,
		rec.ID,
		rec.Alias,
		events.ReasonDuplicateConnection,
		map[string]string{events.MetadataFailureReason: events.ReasonDuplicateConnection},
	))
	return &iop.RegisterResponse{Accepted: false, Reason: "node already connected"}
}

// notifyNodeConnected invokes the authoritative connect handler for an accepted
// registration, ahead of the observability fanout. The handler read mirrors the
// disconnect path: it runs synchronously so a reconnect's resource activation
// and queue pump cannot be lost to a dropped bus delivery.
func (s *Server) notifyNodeConnected(nodeID string, generation uint64) {
	s.handlerMu.RLock()
	handler := s.onNodeConnect
	s.handlerMu.RUnlock()
	if handler != nil {
		handler(nodeID, generation)
	}
}

// emitNodeReady logs and emits the node-connected event once the node has
// signalled dispatch readiness. It is driven off the live registry entry so the
// announced identity is exactly the connection that became ready.
func (s *Server) emitNodeReady(entry *edgenode.NodeEntry) {
	s.logger.Info("node ready",
		zap.String("node_id", entry.NodeID),
		zap.String("alias", entry.Alias),
	)
	s.emitNodeEvent(events.NewEdgeNodeEvent(
		events.SourceEdge,
		events.TypeNodeConnected,
		entry.NodeID,
		entry.Alias,
		events.ReasonRegistered,
		map[string]string{events.MetadataLifecycleState: edgenode.LifecycleConnected},
	))
}

// bindNodeEventListener forwards node-originated EdgeNodeEvents, backfilling the
// node id/alias from the registered record when the node omitted them.
func (s *Server) bindNodeEventListener(client *toki.TcpClient, rec *edgenode.NodeRecord) {
	toki.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(e *iop.EdgeNodeEvent) {
		if e.NodeId == "" {
			e.NodeId = rec.ID
		}
		if e.Alias == "" {
			e.Alias = rec.Alias
		}
		s.emitNodeEvent(e)
	})
}

// bindDisconnectListener classifies the disconnect reason, removes the node from
// the registry, and — only when this client still owned the entry — settles
// lifecycle cleanup before emitting the disconnected event.
//
// The ownership check gates everything downstream: a stale duplicate connection
// closing must not release the live connection's leases, nor announce that the
// still-connected node went away.
func (s *Server) bindDisconnectListener(client *toki.TcpClient, rec *edgenode.NodeRecord) {
	client.AddDisconnectListener(func(_ *toki.TcpClient) {
		transportInfo := client.DisconnectInfo()
		fields := []zap.Field{zap.String("node_id", rec.ID)}
		fields = append(fields, transportDisconnectFields(transportInfo)...)

		// UnregisterIfClient is the ownership gate: it succeeds only for the client
		// that still owns the entry and returns that connection's generation. A
		// rejected duplicate or a connection already superseded by a reconnect gets
		// ok=false here, so nothing downstream — lifecycle callback, lease release,
		// or disconnected event — runs for a stale close.
		generation, ok := s.registry.UnregisterIfClient(rec.ID, client)
		if !ok {
			s.logger.Debug("stale node connection closed", fields...)
			return
		}
		fields = append(fields, zap.Uint64("connection_generation", generation))
		s.logger.Info("node unregistered", fields...)

		reason := events.ReasonTransportClosed
		if transportInfo.Reason == "heartbeat_timeout" {
			reason = events.ReasonHeartbeatTimeout
		} else if s.stopping.Load() {
			reason = events.ReasonEdgeShutdown
		}

		s.handlerMu.RLock()
		disconnectHandler := s.onNodeDisconnect
		s.handlerMu.RUnlock()
		if disconnectHandler != nil {
			disconnectHandler(rec.ID, generation, reason)
		}

		meta := transportDisconnectMetadata(transportInfo)
		if meta == nil {
			meta = make(map[string]string)
		}
		if reason == events.ReasonHeartbeatTimeout {
			meta[events.MetadataFailureReason] = events.ReasonHeartbeatTimeout
			meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
		} else {
			meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
		}
		s.emitNodeEvent(events.NewEdgeNodeEvent(
			events.SourceEdge,
			events.TypeNodeDisconnected,
			rec.ID,
			rec.Alias,
			reason,
			meta,
		))
	})
}