Public release release-f70e47af061d

Source commit: f70e47af061d8f7ba27cd769c8cd2e2f8707dc72

Public tree identity: sha256:03a43db60d295811688bedaa5ad6460df0dbde27fbf37033997f4d8100f83ab2

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

@@ -0,0 +1,452 @@
+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)
+}