clusterflux-public/crates/clusterflux-coordinator/src/service/nodes.rs

use std::collections::BTreeSet;
use std::time::{SystemTime, UNIX_EPOCH};

use clusterflux_core::{
    generate_opaque_token, verify_node_request_signature, Actor, ArtifactId, CredentialKind,
    Digest, NodeCapabilities, NodeDescriptor, NodeId, NodeSignedRequest, ProjectId,
    SourceProviderKind, TenantId, UserId,
};

use crate::CoordinatorError;

use super::{
    bounded_ttl, enrollment_grant_key, CoordinatorResponse, CoordinatorService,
    CoordinatorServiceError,
};

impl CoordinatorService {
    pub fn set_node_stale_after_seconds(&mut self, seconds: u64) {
        self.node_stale_after_seconds = seconds.max(1);
    }

    pub(super) fn liveness_now_epoch_seconds(&self) -> u64 {
        #[cfg(test)]
        if let Some(now) = self.server_time_override {
            return now;
        }
        unix_timestamp_seconds()
    }

    pub(super) fn node_is_live(&self, node: &NodeId) -> bool {
        self.node_last_seen_epoch_seconds
            .get(node)
            .is_some_and(|last_seen| {
                self.liveness_now_epoch_seconds().saturating_sub(*last_seen)
                    <= self.node_stale_after_seconds
            })
    }

    pub(super) fn live_node_descriptors(&self) -> Vec<NodeDescriptor> {
        self.node_descriptors
            .values()
            .cloned()
            .map(|mut descriptor| {
                descriptor.online = self.node_is_live(&descriptor.id);
                descriptor
            })
            .collect()
    }

    pub(super) fn handle_attach_node(
        &mut self,
        tenant: String,
        project: String,
        node: String,
        public_key: String,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let node = NodeId::new(node);
        self.coordinator.ensure_tenant_active(&tenant)?;
        self.coordinator.upsert_tenant(tenant.clone());
        self.coordinator.upsert_user(
            tenant.clone(),
            UserId::from("local-user"),
            CredentialKind::CliDeviceSession,
        );
        self.coordinator
            .upsert_project(tenant.clone(), project.clone(), "local");
        self.coordinator.upsert_source_provider_config(
            tenant.clone(),
            project.clone(),
            SourceProviderKind::Filesystem,
            Digest::sha256("local-filesystem"),
        );
        self.coordinator.enroll_node(
            tenant.clone(),
            project.clone(),
            node.clone(),
            public_key,
            "node:attach",
        );
        self.persist_durable_state()?;
        Ok(CoordinatorResponse::NodeAttached {
            node,
            tenant,
            project,
        })
    }

    pub(super) fn handle_create_node_enrollment_grant(
        &mut self,
        tenant: String,
        project: String,
        actor_user: String,
        ttl_seconds: u64,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let actor = UserId::new(actor_user);
        self.coordinator.ensure_tenant_active(&tenant)?;
        self.coordinator.upsert_tenant(tenant.clone());
        self.coordinator
            .upsert_user(tenant.clone(), actor, CredentialKind::CliDeviceSession);
        self.coordinator
            .upsert_project(tenant.clone(), project.clone(), "local");
        self.coordinator.upsert_source_provider_config(
            tenant.clone(),
            project.clone(),
            SourceProviderKind::Filesystem,
            Digest::sha256("local-filesystem"),
        );
        let now_epoch_seconds = self.current_epoch_seconds()?;
        self.enrollment_grants.retain(|_, grant| {
            !grant.consumed && grant.expires_at_epoch_seconds >= now_epoch_seconds
        });
        if self
            .enrollment_grants
            .values()
            .filter(|grant| grant.tenant == tenant && grant.project == project)
            .count()
            >= super::MAX_ENROLLMENT_GRANTS_PER_PROJECT
        {
            return Err(CoordinatorServiceError::Protocol(
                "node enrollment grant limit reached for this project; consume a grant or wait for one to expire"
                    .to_owned(),
            ));
        }
        let grant =
            generate_opaque_token("node_grant").map_err(CoordinatorServiceError::Protocol)?;
        let ttl_seconds = bounded_ttl(ttl_seconds, self.admission.max_node_enrollment_ttl_seconds);
        let scope = "node:attach".to_owned();
        let expires_at_epoch_seconds = now_epoch_seconds.saturating_add(ttl_seconds);
        let enrollment = self.coordinator.create_node_enrollment_grant(
            tenant.clone(),
            project.clone(),
            grant.clone(),
            scope.clone(),
            expires_at_epoch_seconds,
        );
        self.enrollment_grants
            .insert(enrollment_grant_key(&tenant, &project, &grant), enrollment);
        self.persist_durable_state()?;
        Ok(CoordinatorResponse::NodeEnrollmentGrantCreated {
            tenant,
            project,
            grant,
            scope,
            expires_at_epoch_seconds,
        })
    }

    pub(super) fn handle_exchange_node_enrollment_grant(
        &mut self,
        tenant: String,
        project: String,
        node: String,
        public_key: String,
        enrollment_grant: String,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let node = NodeId::new(node);
        let now_epoch_seconds = self.current_epoch_seconds()?;
        self.enrollment_grants.retain(|_, grant| {
            !grant.consumed && grant.expires_at_epoch_seconds >= now_epoch_seconds
        });
        self.coordinator.ensure_tenant_active(&tenant)?;
        if self.coordinator.node_identity(&node).is_none() {
            self.quota.ensure_node_admission(
                &tenant,
                self.coordinator.node_identity_count_for_tenant(&tenant),
            )?;
        }
        let grant_key = enrollment_grant_key(&tenant, &project, &enrollment_grant);
        let grant =
            self.enrollment_grants
                .get_mut(&grant_key)
                .ok_or(CoordinatorError::Enrollment(
                    clusterflux_core::EnrollmentError::Expired,
                ))?;
        let credential = self.coordinator.exchange_node_enrollment_grant(
            grant,
            node.clone(),
            &public_key,
            "node:attach",
            now_epoch_seconds,
        )?;
        self.enrollment_grants.remove(&grant_key);
        self.persist_durable_state()?;
        Ok(CoordinatorResponse::NodeEnrollmentExchanged {
            node,
            tenant,
            project,
            credential,
        })
    }

    pub(super) fn handle_node_heartbeat(
        &mut self,
        node: String,
        node_signature: Option<NodeSignedRequest>,
        payload_digest: &Digest,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let node = NodeId::new(node);
        self.authenticate_node_request(&node, node_signature, "node_heartbeat", payload_digest)?;
        Ok(CoordinatorResponse::NodeHeartbeat {
            node,
            epoch: self.coordinator.coordinator_epoch(),
        })
    }

    pub(super) fn handle_report_node_capabilities(
        &mut self,
        tenant: String,
        project: String,
        node: String,
        capabilities: NodeCapabilities,
        cached_environment_digests: Vec<Digest>,
        dependency_cache_digests: Vec<Digest>,
        source_snapshots: Vec<Digest>,
        artifact_locations: Vec<String>,
        direct_connectivity: bool,
        _online: bool,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let node = NodeId::new(node);
        let identity = self
            .coordinator
            .node_identity(&node)
            .ok_or(CoordinatorError::UnknownNode)?;
        if identity.tenant != tenant || identity.project != project {
            return Err(CoordinatorError::Unauthorized(
                "node capability report is outside the enrolled tenant/project scope".to_owned(),
            )
            .into());
        }
        capabilities.validate_public_report()?;
        for (kind, count) in [
            ("cached environments", cached_environment_digests.len()),
            ("dependency caches", dependency_cache_digests.len()),
            ("source snapshots", source_snapshots.len()),
            ("artifact locations", artifact_locations.len()),
        ] {
            if count > super::MAX_NODE_REPORTED_OBJECTS_PER_KIND {
                return Err(CoordinatorServiceError::Protocol(format!(
                    "node capability report contains {count} {kind}; limit is {}",
                    super::MAX_NODE_REPORTED_OBJECTS_PER_KIND
                )));
            }
        }
        if cached_environment_digests
            .iter()
            .chain(&dependency_cache_digests)
            .chain(&source_snapshots)
            .any(|digest| !digest.is_valid_sha256())
        {
            return Err(CoordinatorServiceError::Protocol(
                "node capability report contains an invalid digest".to_owned(),
            ));
        }
        if artifact_locations.iter().any(|artifact| {
            artifact.trim().is_empty()
                || artifact.len() > 256
                || artifact
                    .chars()
                    .any(|character| matches!(character, '/' | '\\' | '\0'))
        }) {
            return Err(CoordinatorServiceError::Protocol(
                "node capability report contains an invalid artifact id".to_owned(),
            ));
        }
        let artifact_locations = artifact_locations
            .into_iter()
            .map(ArtifactId::new)
            .collect::<BTreeSet<_>>();
        self.artifact_registry
            .reconcile_node_retention(&node, &artifact_locations);

        let online = self.node_is_live(&node);
        self.node_descriptors.insert(
            node.clone(),
            NodeDescriptor {
                id: node.clone(),
                tenant,
                project,
                capabilities,
                cached_environments: cached_environment_digests.into_iter().collect(),
                dependency_caches: dependency_cache_digests.into_iter().collect(),
                source_snapshots: source_snapshots.into_iter().collect(),
                artifact_locations,
                direct_connectivity,
                online,
            },
        );
        Ok(CoordinatorResponse::NodeCapabilitiesRecorded {
            node,
            node_descriptors: self.node_descriptors.len(),
        })
    }

    pub(super) fn handle_list_node_descriptors(
        &mut self,
        tenant: String,
        project: String,
        actor_user: String,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let actor = UserId::new(actor_user);
        let descriptors = self
            .live_node_descriptors()
            .into_iter()
            .filter(|descriptor| descriptor.tenant == tenant && descriptor.project == project)
            .collect();
        Ok(CoordinatorResponse::NodeDescriptors { descriptors, actor })
    }

    pub(super) fn handle_revoke_node_credential(
        &mut self,
        tenant: String,
        project: String,
        actor_user: String,
        node: String,
    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
        let tenant = TenantId::new(tenant);
        let project = ProjectId::new(project);
        let actor = UserId::new(actor_user);
        let node = NodeId::new(node);
        let context = clusterflux_core::AuthContext {
            tenant: tenant.clone(),
            project: project.clone(),
            actor: Actor::User(actor.clone()),
        };
        self.coordinator.revoke_node_credential(&context, &node)?;
        let descriptor_removed = self.node_descriptors.remove(&node).is_some();
        self.node_last_seen_epoch_seconds.remove(&node);
        self.artifact_registry.garbage_collect_node(&node);
        let queued_assignments_removed = self
            .task_assignments
            .remove(&(tenant.clone(), project.clone(), node.clone()))
            .map_or(0, |assignments| assignments.len());
        self.active_tasks
            .retain(|(task_tenant, task_project, _, task_node, _)| {
                task_tenant != &tenant || task_project != &project || task_node != &node
            });
        self.task_cancellations
            .retain(|(task_tenant, task_project, _, task_node, _)| {
                task_tenant != &tenant || task_project != &project || task_node != &node
            });
        self.task_aborts
            .retain(|(task_tenant, task_project, _, task_node, _)| {
                task_tenant != &tenant || task_project != &project || task_node != &node
            });
        self.task_placements
            .retain(|(task_tenant, task_project, _, task_node, _), _| {
                task_tenant != &tenant || task_project != &project || task_node != &node
            });
        self.persist_durable_state()?;
        Ok(CoordinatorResponse::NodeCredentialRevoked {
            node,
            tenant,
            project,
            actor,
            descriptor_removed,
            queued_assignments_removed,
        })
    }

    pub(super) fn authenticate_node_request(
        &mut self,
        node: &NodeId,
        node_signature: Option<NodeSignedRequest>,
        request_kind: &str,
        payload_digest: &Digest,
    ) -> Result<(), CoordinatorServiceError> {
        let identity = self
            .coordinator
            .node_identity(node)
            .ok_or(CoordinatorError::UnknownNode)?;
        let signature = node_signature.ok_or_else(|| {
            CoordinatorError::Unauthorized(
                "node request requires a signed proof of enrolled private-key possession"
                    .to_owned(),
            )
        })?;
        if signature.nonce.trim().is_empty() || signature.nonce.len() > 256 {
            return Err(CoordinatorError::Unauthorized(
                "node signed request nonce is missing or invalid".to_owned(),
            )
            .into());
        }
        let now_epoch_seconds = unix_timestamp_seconds();
        if signature
            .issued_at_epoch_seconds
            .abs_diff(now_epoch_seconds)
            > super::NODE_SIGNATURE_WINDOW_SECONDS
        {
            return Err(CoordinatorError::Unauthorized(
                "node signed request is expired or outside the allowed clock skew".to_owned(),
            )
            .into());
        }
        let replay_key = (node.clone(), signature.nonce.clone());
        self.node_replay_nonces.retain(|_, accepted_at| {
            now_epoch_seconds <= accepted_at.saturating_add(super::NODE_SIGNATURE_WINDOW_SECONDS)
        });
        if self.node_replay_nonces.contains_key(&replay_key) {
            return Err(CoordinatorError::Unauthorized(
                "node signed request nonce has already been used".to_owned(),
            )
            .into());
        }
        verify_node_request_signature(
            &identity.public_key,
            node,
            request_kind,
            payload_digest,
            &signature,
        )
        .map_err(CoordinatorError::Unauthorized)?;
        if self
            .node_replay_nonces
            .keys()
            .filter(|(retained_node, _)| retained_node == node)
            .count()
            >= super::MAX_NODE_REPLAY_NONCES_PER_AUTHORITY
        {
            return Err(CoordinatorError::Unauthorized(
                "node signed request replay window is full; retry after the bounded signature window advances"
                    .to_owned(),
            )
            .into());
        }
        self.node_replay_nonces
            .insert(replay_key, now_epoch_seconds);
        let seen_at = self.liveness_now_epoch_seconds();
        self.node_last_seen_epoch_seconds
            .insert(node.clone(), seen_at);
        if let Some(descriptor) = self.node_descriptors.get_mut(node) {
            descriptor.online = true;
        }
        Ok(())
    }
}

fn unix_timestamp_seconds() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_secs())
        .unwrap_or(0)
}