1use anyhow::{Context, Result};
54use std::path::PathBuf;
55use std::time::Duration;
56use tokio::process::Command;
57use tokio::time::interval;
58use tracing::{error, info, warn};
59
60#[cfg(unix)]
61use tokio::signal::unix::{signal, SignalKind};
62
63#[derive(Debug)]
65pub struct Config {
66 pub bucket: String,
67 pub prefix: String,
68 pub queue_name: String,
77 pub pool_id: String,
83 pub supervisor: PathBuf,
84 pub poll_interval: Duration,
85 pub heartbeat_interval: Duration,
86 pub node_id: String,
87}
88
89impl Config {
90 pub fn from_env() -> Result<Self> {
91 let bucket =
92 std::env::var("CELLOS_FLEET_BUCKET").context("CELLOS_FLEET_BUCKET is required")?;
93 let prefix = std::env::var("CELLOS_FLEET_PREFIX").unwrap_or_else(|_| "fleet".to_string());
94 let queue_name = std::env::var("CELLOS_FLEET_QUEUE_NAME")
95 .map(|value| value.trim().to_string())
96 .unwrap_or_default();
97 let pool_id = std::env::var("CELLOS_FLEET_POOL_ID")
98 .map(|value| value.trim().to_string())
99 .unwrap_or_default();
100 let supervisor = PathBuf::from(
101 std::env::var("CELLOS_FLEET_SUPERVISOR")
102 .unwrap_or_else(|_| "cellos-supervisor".to_string()),
103 );
104 let poll_ms: u64 = std::env::var("CELLOS_FLEET_POLL_INTERVAL_MS")
105 .ok()
106 .and_then(|v| v.parse().ok())
107 .unwrap_or(5000);
108 let heartbeat_ms: u64 = std::env::var("CELLOS_FLEET_HEARTBEAT_INTERVAL_MS")
109 .ok()
110 .and_then(|v| v.parse().ok())
111 .unwrap_or(30_000);
112 let node_id = std::env::var("CELLOS_FLEET_NODE_ID").unwrap_or_else(|_| {
113 hostname::get()
114 .ok()
115 .and_then(|h| h.into_string().ok())
116 .unwrap_or_else(|| "unknown-node".to_string())
117 });
118 Ok(Config {
119 bucket,
120 prefix,
121 queue_name,
122 pool_id,
123 supervisor,
124 poll_interval: Duration::from_millis(poll_ms),
125 heartbeat_interval: Duration::from_millis(heartbeat_ms),
126 node_id,
127 })
128 }
129
130 pub fn queue_prefix(&self) -> String {
131 let base_prefix = self.prefix.trim_end_matches('/');
132 if self.queue_name.is_empty() {
133 base_prefix.to_string()
134 } else {
135 format!("{}/{}/", base_prefix, self.queue_name)
136 .trim_end_matches('/')
137 .to_string()
138 }
139 }
140
141 pub fn pending_prefix(&self) -> String {
142 format!("{}/pending/", self.queue_prefix())
143 }
144
145 pub fn claimed_key(&self, spec_id: &str) -> String {
146 format!("{}/claimed/{}.json", self.queue_prefix(), spec_id)
147 }
148
149 pub fn completed_key(&self, spec_id: &str) -> String {
150 format!("{}/completed/{}.json", self.queue_prefix(), spec_id)
151 }
152
153 pub fn failed_key(&self, spec_id: &str) -> String {
154 format!("{}/failed/{}.json", self.queue_prefix(), spec_id)
155 }
156
157 pub fn should_dispatch(&self, spec_pool_id: Option<&str>) -> bool {
170 if self.pool_id.is_empty() {
171 return true;
172 }
173 match spec_pool_id {
174 None => true,
175 Some(spec_pool) => spec_pool == self.pool_id,
176 }
177 }
178}
179
180#[derive(Debug, Clone, PartialEq, Eq)]
203pub struct NodeIdentity {
204 pub node_id: String,
206 pub public_key_ref: String,
209 pub attestation: Option<()>,
211}
212
213impl NodeIdentity {
214 pub fn new(node_id: impl Into<String>, public_key_ref: impl Into<String>) -> Self {
217 NodeIdentity {
218 node_id: node_id.into(),
219 public_key_ref: public_key_ref.into(),
220 attestation: None,
221 }
222 }
223}
224
225pub const NODE_ATTESTATION_EVENT_TYPE: &str = "dev.cellos.events.fleet.v1.node-attestation";
244
245pub const IDENTITY_KIND_KEY_POSSESSION: &str = "key-possession";
250
251pub fn build_node_attestation(
277 identity: &NodeIdentity,
278 signer: &dyn cellos_core::Signer,
279 event_id: impl Into<String>,
280 source: impl Into<String>,
281 time: Option<String>,
282) -> Result<cellos_core::SignedEventEnvelopeV1, cellos_core::CellosError> {
283 let event = cellos_core::CloudEventV1 {
284 specversion: "1.0".to_string(),
285 id: event_id.into(),
286 source: source.into(),
287 ty: NODE_ATTESTATION_EVENT_TYPE.to_string(),
288 datacontenttype: Some("application/json".to_string()),
289 data: Some(serde_json::json!({
290 "nodeId": identity.node_id,
291 "publicKeyRef": identity.public_key_ref,
292 "identityKind": IDENTITY_KIND_KEY_POSSESSION,
293 "attestation": serde_json::Value::Null,
296 })),
297 time,
298 traceparent: None,
299 cex: None,
300 };
301 cellos_core::sign_event_with(signer, &event)
302}
303
304pub fn verify_node_attestation(
326 envelope: &cellos_core::SignedEventEnvelopeV1,
327 verifying_keys: &std::collections::HashMap<String, cellos_core::crypto::TrustAnchorPublicKey>,
328) -> Result<NodeIdentity, cellos_core::CellosError> {
329 use cellos_core::CellosError;
330
331 let empty_hmac: std::collections::HashMap<String, Vec<u8>> = std::collections::HashMap::new();
333 let event = cellos_core::verify_signed_event_envelope(envelope, verifying_keys, &empty_hmac)?;
334
335 if event.ty != NODE_ATTESTATION_EVENT_TYPE {
338 return Err(CellosError::InvalidSpec(format!(
339 "node attestation: expected event type {NODE_ATTESTATION_EVENT_TYPE}, got {:?}",
340 event.ty
341 )));
342 }
343
344 let data = event.data.as_ref().ok_or_else(|| {
346 CellosError::InvalidSpec("node attestation: event has no data payload".into())
347 })?;
348
349 let identity_kind = data.get("identityKind").and_then(|v| v.as_str());
350 if identity_kind != Some(IDENTITY_KIND_KEY_POSSESSION) {
351 return Err(CellosError::InvalidSpec(format!(
352 "node attestation: expected identityKind {IDENTITY_KIND_KEY_POSSESSION:?}, got {:?}",
353 identity_kind
354 )));
355 }
356
357 match data.get("attestation") {
361 Some(serde_json::Value::Null) => {}
362 other => {
363 return Err(CellosError::InvalidSpec(format!(
364 "node attestation: attestation must be explicit null (NOT hardware-attested), got {:?}",
365 other
366 )));
367 }
368 }
369
370 let node_id = data.get("nodeId").and_then(|v| v.as_str()).ok_or_else(|| {
371 CellosError::InvalidSpec("node attestation: missing/invalid nodeId".into())
372 })?;
373 let public_key_ref = data
374 .get("publicKeyRef")
375 .and_then(|v| v.as_str())
376 .ok_or_else(|| {
377 CellosError::InvalidSpec("node attestation: missing/invalid publicKeyRef".into())
378 })?;
379
380 Ok(NodeIdentity::new(node_id, public_key_ref))
381}
382
383pub const MIN_CEILING_EPOCH_ENV: &str = "CELLOS_FLEET_MIN_CEILING_EPOCH";
394
395#[derive(Debug, Clone, PartialEq, Eq)]
397pub enum CeilingEpochDecision {
398 Admit,
401 Refuse { reason: &'static str },
404}
405
406impl CeilingEpochDecision {
407 pub fn admits(&self) -> bool {
409 matches!(self, CeilingEpochDecision::Admit)
410 }
411}
412
413pub const CEILING_EPOCH_STALE: &str = "ceiling_epoch_stale";
415
416pub fn ceiling_epoch_admits(observed_epoch: u64, floor: Option<u64>) -> CeilingEpochDecision {
428 match floor {
429 None => CeilingEpochDecision::Admit,
430 Some(floor) if observed_epoch >= floor => CeilingEpochDecision::Admit,
431 Some(_) => CeilingEpochDecision::Refuse {
432 reason: CEILING_EPOCH_STALE,
433 },
434 }
435}
436
437pub fn read_min_ceiling_epoch() -> Result<Option<u64>> {
444 match std::env::var(MIN_CEILING_EPOCH_ENV) {
445 Err(std::env::VarError::NotPresent) => Ok(None),
446 Err(e) => Err(anyhow::anyhow!(
447 "{MIN_CEILING_EPOCH_ENV} is not valid unicode: {e}"
448 )),
449 Ok(raw) => {
450 let parsed = raw.trim().parse::<u64>().with_context(|| {
451 format!("{MIN_CEILING_EPOCH_ENV}={raw:?} is not a valid u64 epoch")
452 })?;
453 Ok(Some(parsed))
454 }
455 }
456}
457
458pub const FLEET_PROFILE_ENV: &str = "CELLOS_FLEET_PROFILE";
470
471#[derive(Debug, Clone, Copy, PartialEq, Eq)]
473pub enum FleetProfile {
474 Default,
477 Hardened,
479}
480
481impl FleetProfile {
482 pub fn from_raw(raw: Option<&str>) -> Self {
486 match raw.map(|s| s.trim().to_ascii_lowercase()) {
487 Some(s) if s == "hardened" => FleetProfile::Hardened,
488 _ => FleetProfile::Default,
489 }
490 }
491
492 pub fn from_env() -> Self {
494 Self::from_raw(std::env::var(FLEET_PROFILE_ENV).ok().as_deref())
495 }
496
497 pub fn is_hardened(&self) -> bool {
499 matches!(self, FleetProfile::Hardened)
500 }
501}
502
503#[derive(Debug, Clone, Copy)]
508pub struct PreflightChecks {
509 pub can_sign_attestation: bool,
511 pub ceiling_epoch_fresh: bool,
513 pub audit_sink_reachable: bool,
515}
516
517impl PreflightChecks {
518 pub fn all_ok() -> Self {
520 PreflightChecks {
521 can_sign_attestation: true,
522 ceiling_epoch_fresh: true,
523 audit_sink_reachable: true,
524 }
525 }
526
527 pub fn failure_reasons(&self) -> Vec<&'static str> {
530 let mut reasons = Vec::new();
531 if !self.can_sign_attestation {
532 reasons.push("cannot_sign_attestation");
533 }
534 if !self.ceiling_epoch_fresh {
535 reasons.push(CEILING_EPOCH_STALE);
536 }
537 if !self.audit_sink_reachable {
538 reasons.push("audit_sink_unreachable");
539 }
540 reasons
541 }
542}
543
544#[derive(Debug, Clone, PartialEq, Eq)]
546pub enum ProfileAdmission {
547 Admit { advisories: Vec<&'static str> },
551 Refuse { reasons: Vec<&'static str> },
554}
555
556impl ProfileAdmission {
557 pub fn admits(&self) -> bool {
559 matches!(self, ProfileAdmission::Admit { .. })
560 }
561}
562
563pub fn profile_admits(profile: FleetProfile, checks: &PreflightChecks) -> ProfileAdmission {
574 let failures = checks.failure_reasons();
575 if failures.is_empty() {
576 return ProfileAdmission::Admit {
577 advisories: Vec::new(),
578 };
579 }
580 if profile.is_hardened() {
581 ProfileAdmission::Refuse { reasons: failures }
582 } else {
583 ProfileAdmission::Admit {
584 advisories: failures,
585 }
586 }
587}
588
589#[derive(Debug, serde::Deserialize)]
595#[serde(rename_all = "camelCase")]
596struct FleetSpecView {
597 #[serde(default)]
598 spec: FleetSpecBody,
599}
600
601#[derive(Debug, Default, serde::Deserialize)]
602#[serde(rename_all = "camelCase")]
603struct FleetSpecBody {
604 #[serde(default)]
605 placement: Option<FleetPlacementView>,
606}
607
608#[derive(Debug, Default, serde::Deserialize)]
609#[serde(rename_all = "camelCase")]
610struct FleetPlacementView {
611 #[serde(default)]
612 pool_id: Option<String>,
613}
614
615pub fn read_spec_pool_id(spec_path: &std::path::Path) -> Result<Option<String>> {
619 let bytes = std::fs::read(spec_path)
620 .with_context(|| format!("read spec file {}", spec_path.display()))?;
621 let view: FleetSpecView = serde_json::from_slice(&bytes)
622 .with_context(|| format!("parse spec file {}", spec_path.display()))?;
623 Ok(view.spec.placement.and_then(|p| p.pool_id))
624}
625
626pub async fn list_pending(cfg: &Config) -> Result<Vec<String>> {
631 let output = Command::new("aws")
632 .args([
633 "s3api",
634 "list-objects-v2",
635 "--bucket",
636 &cfg.bucket,
637 "--prefix",
638 &cfg.pending_prefix(),
639 "--query",
640 "Contents[?ends_with(Key, '.json')].Key",
641 "--output",
642 "json",
643 ])
644 .output()
645 .await
646 .context("aws s3api list-objects-v2 failed")?;
647
648 if !output.status.success() {
649 let stderr = String::from_utf8_lossy(&output.stderr);
650 warn!("list_pending stderr: {stderr}");
651 return Ok(vec![]);
652 }
653
654 let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
655 if stdout == "null" || stdout.is_empty() {
656 return Ok(vec![]);
657 }
658
659 let keys: Vec<String> =
660 serde_json::from_str(&stdout).context("failed to parse list-objects-v2 output")?;
661 Ok(keys)
662}
663
664pub async fn try_claim(cfg: &Config, pending_key: &str) -> Result<bool> {
669 let spec_id = pending_key
671 .trim_start_matches(&cfg.pending_prefix())
672 .trim_end_matches(".json");
673
674 let src = format!("s3://{}/{}", cfg.bucket, pending_key);
675 let dst = format!("s3://{}/{}", cfg.bucket, cfg.claimed_key(spec_id));
676
677 let status = Command::new("aws")
678 .args(["s3", "mv", &src, &dst])
679 .status()
680 .await
681 .context("aws s3 mv (claim) failed")?;
682
683 Ok(status.success())
684}
685
686pub async fn download_spec(cfg: &Config, spec_id: &str) -> Result<tempfile::NamedTempFile> {
688 let tmp = tempfile::Builder::new()
689 .prefix("cellos-fleet-spec-")
690 .suffix(".json")
691 .tempfile()
692 .context("failed to create temp file for spec")?;
693
694 let s3_key = format!("s3://{}/{}", cfg.bucket, cfg.claimed_key(spec_id));
695
696 let status = Command::new("aws")
697 .args(["s3", "cp", &s3_key, tmp.path().to_str().unwrap()])
698 .status()
699 .await
700 .context("aws s3 cp (download spec) failed")?;
701
702 anyhow::ensure!(status.success(), "aws s3 cp exited non-zero");
703 Ok(tmp)
704}
705
706pub async fn peek_pending_spec(cfg: &Config, pending_key: &str) -> Result<tempfile::NamedTempFile> {
709 let tmp = tempfile::Builder::new()
710 .prefix("cellos-fleet-peek-")
711 .suffix(".json")
712 .tempfile()
713 .context("failed to create temp file for peek")?;
714
715 let s3_key = format!("s3://{}/{}", cfg.bucket, pending_key);
716
717 let status = Command::new("aws")
718 .args(["s3", "cp", &s3_key, tmp.path().to_str().unwrap()])
719 .status()
720 .await
721 .context("aws s3 cp (peek pending spec) failed")?;
722
723 anyhow::ensure!(status.success(), "aws s3 cp (peek) exited non-zero");
724 Ok(tmp)
725}
726
727pub async fn run_cell(cfg: &Config, spec_path: &std::path::Path) -> Result<i32> {
731 let status = Command::new(&cfg.supervisor)
732 .arg(spec_path)
733 .status()
734 .await
735 .context("cellos-supervisor failed to launch")?;
736
737 Ok(status.code().unwrap_or(1))
738}
739
740pub async fn finalize(cfg: &Config, spec_id: &str, exit_code: i32) -> Result<()> {
742 let src = format!("s3://{}/{}", cfg.bucket, cfg.claimed_key(spec_id));
743 let dst = if exit_code == 0 {
744 format!("s3://{}/{}", cfg.bucket, cfg.completed_key(spec_id))
745 } else {
746 format!("s3://{}/{}", cfg.bucket, cfg.failed_key(spec_id))
747 };
748
749 Command::new("aws")
750 .args(["s3", "mv", &src, &dst])
751 .status()
752 .await
753 .context("aws s3 mv (finalize) failed")?;
754
755 Ok(())
756}
757
758pub async fn process_spec(cfg: &Config, pending_key: &str) -> Result<()> {
760 let spec_id = pending_key
761 .trim_start_matches(&cfg.pending_prefix())
762 .trim_end_matches(".json");
763
764 if !cfg.pool_id.is_empty() {
771 let peek_tmp = peek_pending_spec(cfg, pending_key).await?;
772 let spec_pool = read_spec_pool_id(peek_tmp.path()).unwrap_or_else(|e| {
773 warn!(spec_id, error = %e, "failed to read placement.poolId — treating as no constraint");
774 None
775 });
776 if !cfg.should_dispatch(spec_pool.as_deref()) {
777 info!(
778 node = %cfg.node_id,
779 spec_id,
780 "skipping spec {}: placement.poolId={} != runner poolId={}",
781 spec_id,
782 spec_pool.as_deref().unwrap_or("<none>"),
783 cfg.pool_id,
784 );
785 return Ok(());
786 }
787 }
788
789 info!(node = %cfg.node_id, spec_id, "claiming spec");
790
791 if !try_claim(cfg, pending_key).await? {
792 info!(spec_id, "spec already claimed by another node, skipping");
793 return Ok(());
794 }
795
796 info!(node = %cfg.node_id, spec_id, "claimed — downloading");
797 let tmp = download_spec(cfg, spec_id).await?;
798
799 info!(node = %cfg.node_id, spec_id, path = %tmp.path().display(), "running cell");
800 let exit_code = run_cell(cfg, tmp.path()).await?;
801
802 info!(node = %cfg.node_id, spec_id, exit_code, "cell completed — finalizing");
803 finalize(cfg, spec_id, exit_code).await?;
804
805 if exit_code == 0 {
806 info!(node = %cfg.node_id, spec_id, "spec completed successfully");
807 } else {
808 warn!(node = %cfg.node_id, spec_id, exit_code, "spec completed with failure");
809 }
810
811 Ok(())
812}
813
814#[cfg(unix)]
815async fn wait_for_shutdown_signal() -> Result<()> {
816 let mut sigterm =
817 signal(SignalKind::terminate()).context("failed to install SIGTERM handler")?;
818 sigterm.recv().await;
819 Ok(())
820}
821
822#[cfg(not(unix))]
823async fn wait_for_shutdown_signal() -> Result<()> {
824 tokio::signal::ctrl_c()
825 .await
826 .context("failed to install Ctrl+C handler")?;
827 Ok(())
828}
829
830pub async fn run(cfg: Config) -> Result<()> {
839 info!(
840 node = %cfg.node_id,
841 bucket = %cfg.bucket,
842 prefix = %cfg.prefix,
843 queue_name = %cfg.queue_name,
844 pool_id = %cfg.pool_id,
845 supervisor = %cfg.supervisor.display(),
846 poll_interval_ms = cfg.poll_interval.as_millis(),
847 heartbeat_interval_ms = cfg.heartbeat_interval.as_millis(),
848 "cellos-fleet agent starting"
849 );
850
851 let mut poll_tick = interval(cfg.poll_interval);
852 let mut heartbeat_tick = interval(cfg.heartbeat_interval);
853 let shutdown = wait_for_shutdown_signal();
854 tokio::pin!(shutdown);
855
856 poll_tick.tick().await;
858 heartbeat_tick.tick().await;
859
860 loop {
861 tokio::select! {
862 _ = poll_tick.tick() => {
864 match list_pending(&cfg).await {
865 Err(e) => error!("list_pending error: {e:#}"),
866 Ok(keys) if keys.is_empty() => {}
867 Ok(keys) => {
868 for key in &keys {
869 if let Err(e) = process_spec(&cfg, key).await {
870 error!(key, "process_spec error: {e:#}");
871 }
872 }
873 }
874 }
875 }
876
877 _ = heartbeat_tick.tick() => {
879 info!(
880 event_type = "dev.cellos.events.fleet.v1.heartbeat",
881 node = %cfg.node_id,
882 bucket = %cfg.bucket,
883 prefix = %cfg.prefix,
884 queue_name = %cfg.queue_name,
885 pool_id = %cfg.pool_id,
886 "heartbeat"
887 );
888 }
889
890 _ = &mut shutdown => {
892 info!(
893 node = %cfg.node_id,
894 "SIGTERM received — draining (no new work accepted)"
895 );
896 break;
897 }
898 }
899 }
900
901 info!(node = %cfg.node_id, "cellos-fleet agent stopped (drain complete)");
902 Ok(())
903}
904
905#[cfg(test)]
906mod tests {
907 use super::{Config, NodeIdentity};
908 use std::path::PathBuf;
909 use std::time::Duration;
910
911 fn config(prefix: &str, queue_name: &str) -> Config {
912 config_with_pool(prefix, queue_name, "")
913 }
914
915 fn config_with_pool(prefix: &str, queue_name: &str, pool_id: &str) -> Config {
916 Config {
917 bucket: "bucket".into(),
918 prefix: prefix.into(),
919 queue_name: queue_name.into(),
920 pool_id: pool_id.into(),
921 supervisor: PathBuf::from("cellos-supervisor"),
922 poll_interval: Duration::from_secs(5),
923 heartbeat_interval: Duration::from_secs(30),
924 node_id: "node-a".into(),
925 }
926 }
927
928 #[test]
929 fn uses_legacy_layout_when_queue_name_is_empty() {
930 let cfg = config("fleet", "");
931
932 assert_eq!(cfg.pending_prefix(), "fleet/pending/");
933 assert_eq!(cfg.claimed_key("spec-1"), "fleet/claimed/spec-1.json");
934 assert_eq!(cfg.completed_key("spec-1"), "fleet/completed/spec-1.json");
935 assert_eq!(cfg.failed_key("spec-1"), "fleet/failed/spec-1.json");
936 }
937
938 #[test]
939 fn uses_queue_qualified_layout_when_queue_name_is_set() {
940 let cfg = config("fleet", "gpu-runners");
941
942 assert_eq!(cfg.pending_prefix(), "fleet/gpu-runners/pending/");
943 assert_eq!(
944 cfg.claimed_key("spec-1"),
945 "fleet/gpu-runners/claimed/spec-1.json"
946 );
947 assert_eq!(
948 cfg.completed_key("spec-1"),
949 "fleet/gpu-runners/completed/spec-1.json"
950 );
951 assert_eq!(
952 cfg.failed_key("spec-1"),
953 "fleet/gpu-runners/failed/spec-1.json"
954 );
955 }
956
957 #[test]
958 fn trims_trailing_slash_from_prefix() {
959 let cfg = config("fleet/", "gpu-runners");
960
961 assert_eq!(cfg.pending_prefix(), "fleet/gpu-runners/pending/");
962 assert_eq!(
963 cfg.claimed_key("spec-1"),
964 "fleet/gpu-runners/claimed/spec-1.json"
965 );
966 }
967
968 #[test]
973 fn dispatch_matrix_for_pool_id_placement_gate() {
974 let unbounded = config_with_pool("fleet", "", "");
976 assert!(
977 unbounded.should_dispatch(None),
978 "no-pool runner must accept specs without a poolId constraint"
979 );
980 assert!(
981 unbounded.should_dispatch(Some("runner-pool-amd64")),
982 "no-pool runner must accept specs with any poolId constraint"
983 );
984
985 let amd64 = config_with_pool("fleet", "", "runner-pool-amd64");
988 assert!(
989 amd64.should_dispatch(None),
990 "pool-bound runner must accept specs with no poolId constraint"
991 );
992 assert!(
993 amd64.should_dispatch(Some("runner-pool-amd64")),
994 "pool-bound runner must accept matching poolId"
995 );
996 assert!(
997 !amd64.should_dispatch(Some("runner-pool-arm64")),
998 "pool-bound runner must skip mismatching poolId"
999 );
1000 }
1001
1002 #[test]
1003 fn read_spec_pool_id_parses_placement_and_handles_absence() {
1004 use std::io::Write;
1005
1006 let mut with_pool = tempfile::NamedTempFile::new().unwrap();
1008 write!(
1009 with_pool,
1010 r#"{{
1011 "apiVersion": "cellos.io/v1",
1012 "kind": "ExecutionCell",
1013 "spec": {{
1014 "id": "test",
1015 "placement": {{ "poolId": "runner-pool-amd64" }}
1016 }}
1017 }}"#
1018 )
1019 .unwrap();
1020 let pool = super::read_spec_pool_id(with_pool.path()).unwrap();
1021 assert_eq!(pool.as_deref(), Some("runner-pool-amd64"));
1022
1023 let mut without = tempfile::NamedTempFile::new().unwrap();
1025 write!(
1026 without,
1027 r#"{{
1028 "apiVersion": "cellos.io/v1",
1029 "kind": "ExecutionCell",
1030 "spec": {{ "id": "test" }}
1031 }}"#
1032 )
1033 .unwrap();
1034 assert_eq!(super::read_spec_pool_id(without.path()).unwrap(), None);
1035 }
1036
1037 #[test]
1041 fn node_identity_new_defaults_attestation_to_none() {
1042 let id = NodeIdentity::new("node-a", "arn:aws:kms:us-east-1:0:key/abc");
1043 assert_eq!(id.node_id, "node-a");
1044 assert_eq!(id.public_key_ref, "arn:aws:kms:us-east-1:0:key/abc");
1045 assert_eq!(id.attestation, None);
1046 }
1047
1048 #[test]
1049 fn node_identity_is_constructible_with_explicit_fields() {
1050 let id = NodeIdentity {
1051 node_id: "node-b".into(),
1052 public_key_ref: "fingerprint:deadbeef".into(),
1053 attestation: None,
1054 };
1055 assert_eq!(id, NodeIdentity::new("node-b", "fingerprint:deadbeef"));
1056 }
1057
1058 use super::{
1060 build_node_attestation, ceiling_epoch_admits, profile_admits, verify_node_attestation,
1061 CeilingEpochDecision, FleetProfile, PreflightChecks, ProfileAdmission,
1062 IDENTITY_KIND_KEY_POSSESSION, NODE_ATTESTATION_EVENT_TYPE,
1063 };
1064 use cellos_core::crypto::{dalek::public_key_from_seed, TrustAnchorPublicKey};
1065 use cellos_core::SoftwareSigner;
1066 use std::collections::HashMap;
1067
1068 fn org_root_keys(kid: &str, seed: [u8; 32]) -> HashMap<String, TrustAnchorPublicKey> {
1070 let mut keys = HashMap::new();
1071 keys.insert(
1072 kid.to_string(),
1073 TrustAnchorPublicKey::from_bytes_unchecked(public_key_from_seed(&seed).unwrap()),
1074 );
1075 keys
1076 }
1077
1078 #[test]
1079 fn node_attestation_verifies_offline_against_org_root_only() {
1080 let seed = [9u8; 32];
1083 let kid = "org-root-2026";
1084 let signer = SoftwareSigner::from_seed(kid, seed).expect("signer");
1085 let identity = NodeIdentity::new("node-alpha", "arn:aws:kms:us-east-1:0:key/abc");
1086
1087 let envelope = build_node_attestation(
1088 &identity,
1089 &signer,
1090 "ev-attest-1",
1091 "/cellos-fleet/node-alpha",
1092 Some("2026-06-25T00:00:00Z".into()),
1093 )
1094 .expect("build attestation");
1095
1096 assert_eq!(envelope.event.ty, NODE_ATTESTATION_EVENT_TYPE);
1098 let data = envelope.event.data.as_ref().unwrap();
1099 assert_eq!(
1100 data.get("identityKind").and_then(|v| v.as_str()),
1101 Some(IDENTITY_KIND_KEY_POSSESSION)
1102 );
1103 assert!(data.get("attestation").unwrap().is_null());
1105
1106 let keys = org_root_keys(kid, seed);
1107 let asserted = verify_node_attestation(&envelope, &keys).expect("verify offline");
1108 assert_eq!(asserted, identity);
1109 }
1110
1111 #[test]
1112 fn node_attestation_tamper_fails_closed() {
1113 let seed = [9u8; 32];
1115 let kid = "org-root-2026";
1116 let signer = SoftwareSigner::from_seed(kid, seed).expect("signer");
1117 let identity = NodeIdentity::new("node-alpha", "ref-1");
1118
1119 let mut envelope =
1120 build_node_attestation(&identity, &signer, "ev-1", "/src", None).expect("build");
1121 envelope.event.data = Some(serde_json::json!({
1123 "nodeId": "node-impostor",
1124 "publicKeyRef": "ref-1",
1125 "identityKind": IDENTITY_KIND_KEY_POSSESSION,
1126 "attestation": serde_json::Value::Null,
1127 }));
1128
1129 let keys = org_root_keys(kid, seed);
1130 let err = verify_node_attestation(&envelope, &keys)
1131 .expect_err("tampered attestation must fail verify");
1132 assert!(
1133 format!("{err}").contains("ed25519 verify failed"),
1134 "expected signature failure, got: {err}"
1135 );
1136 }
1137
1138 #[test]
1139 fn node_attestation_wrong_org_root_fails() {
1140 let signer = SoftwareSigner::from_seed("org-root-2026", [9u8; 32]).expect("signer");
1142 let identity = NodeIdentity::new("node-alpha", "ref-1");
1143 let envelope =
1144 build_node_attestation(&identity, &signer, "ev-1", "/src", None).expect("build");
1145
1146 let wrong_keys = org_root_keys("org-root-2026", [1u8; 32]);
1148 let err = verify_node_attestation(&envelope, &wrong_keys)
1149 .expect_err("wrong org-root key must fail");
1150 assert!(format!("{err}").contains("ed25519 verify failed"));
1151 }
1152
1153 #[test]
1155 fn ceiling_epoch_floor_refuses_when_observed_below_floor() {
1156 let decision = ceiling_epoch_admits(4, Some(5));
1158 assert_eq!(
1159 decision,
1160 CeilingEpochDecision::Refuse {
1161 reason: super::CEILING_EPOCH_STALE
1162 }
1163 );
1164 assert!(!decision.admits());
1165
1166 assert!(ceiling_epoch_admits(5, Some(5)).admits());
1168 assert!(ceiling_epoch_admits(6, Some(5)).admits());
1170
1171 assert!(ceiling_epoch_admits(0, None).admits());
1173 assert!(ceiling_epoch_admits(u64::MAX, None).admits());
1174 }
1175
1176 #[test]
1178 fn profile_from_raw_only_hardened_selects_hardened() {
1179 assert_eq!(
1180 FleetProfile::from_raw(Some("hardened")),
1181 FleetProfile::Hardened
1182 );
1183 assert_eq!(
1184 FleetProfile::from_raw(Some(" Hardened ")),
1185 FleetProfile::Hardened
1186 );
1187 assert_eq!(
1188 FleetProfile::from_raw(Some("HARDENED")),
1189 FleetProfile::Hardened
1190 );
1191 assert_eq!(
1192 FleetProfile::from_raw(Some("default")),
1193 FleetProfile::Default
1194 );
1195 assert_eq!(FleetProfile::from_raw(Some("")), FleetProfile::Default);
1196 assert_eq!(FleetProfile::from_raw(None), FleetProfile::Default);
1197 }
1198
1199 #[test]
1200 fn profile_admits_all_ok_regardless_of_profile() {
1201 let ok = PreflightChecks::all_ok();
1202 assert_eq!(
1203 profile_admits(FleetProfile::Hardened, &ok),
1204 ProfileAdmission::Admit { advisories: vec![] }
1205 );
1206 assert_eq!(
1207 profile_admits(FleetProfile::Default, &ok),
1208 ProfileAdmission::Admit { advisories: vec![] }
1209 );
1210 }
1211
1212 #[test]
1213 fn hardened_fails_closed_on_each_precondition() {
1214 let cannot_sign = PreflightChecks {
1216 can_sign_attestation: false,
1217 ..PreflightChecks::all_ok()
1218 };
1219 let d = profile_admits(FleetProfile::Hardened, &cannot_sign);
1220 assert!(!d.admits());
1221 assert_eq!(
1222 d,
1223 ProfileAdmission::Refuse {
1224 reasons: vec!["cannot_sign_attestation"]
1225 }
1226 );
1227
1228 let stale = PreflightChecks {
1230 ceiling_epoch_fresh: false,
1231 ..PreflightChecks::all_ok()
1232 };
1233 assert_eq!(
1234 profile_admits(FleetProfile::Hardened, &stale),
1235 ProfileAdmission::Refuse {
1236 reasons: vec![super::CEILING_EPOCH_STALE]
1237 }
1238 );
1239
1240 let sink_down = PreflightChecks {
1242 audit_sink_reachable: false,
1243 ..PreflightChecks::all_ok()
1244 };
1245 assert_eq!(
1246 profile_admits(FleetProfile::Hardened, &sink_down),
1247 ProfileAdmission::Refuse {
1248 reasons: vec!["audit_sink_unreachable"]
1249 }
1250 );
1251 }
1252
1253 #[test]
1254 fn default_profile_fails_open_with_named_advisory() {
1255 let cannot_sign = PreflightChecks {
1257 can_sign_attestation: false,
1258 ..PreflightChecks::all_ok()
1259 };
1260 let d = profile_admits(FleetProfile::Default, &cannot_sign);
1261 assert!(d.admits());
1262 assert_eq!(
1263 d,
1264 ProfileAdmission::Admit {
1265 advisories: vec!["cannot_sign_attestation"]
1266 }
1267 );
1268
1269 let multi = PreflightChecks {
1271 can_sign_attestation: false,
1272 ceiling_epoch_fresh: false,
1273 audit_sink_reachable: false,
1274 };
1275 assert_eq!(
1276 profile_admits(FleetProfile::Default, &multi),
1277 ProfileAdmission::Admit {
1278 advisories: vec![
1279 "cannot_sign_attestation",
1280 super::CEILING_EPOCH_STALE,
1281 "audit_sink_unreachable"
1282 ]
1283 }
1284 );
1285 }
1286}