clawft_kernel/workload_runtime/
cog_spec.rs

1//! The parts of a cog's `cog.toml` the runtime enforces (ADR-100 sections
2//! 2-4): `[console]` limits, `[resources]` and the `[config]` CLI surface,
3//! plus building a [`VerifiedWorkload`] from a verified package.
4
5use std::collections::{BTreeMap, BTreeSet};
6use std::sync::Arc;
7
8use serde::Deserialize;
9
10use super::types::{Binary, RuntimeError, SignedPayload, VerifiedWorkload, WorkloadSource};
11use crate::workload_pkg::manifest::valid_cog_id;
12use crate::workload_pkg::verify::check_file;
13use crate::workload_pkg::{FileSource, VerifiedPackage};
14
15/// Default `[console].max_runtime_secs` when a cog does not set one.
16pub const DEFAULT_MAX_RUNTIME_SECS: u64 = 15;
17/// Largest `max_runtime_secs` honored.
18pub const MAX_RUNTIME_SECS: u64 = 3600;
19/// Default `[console].output_limit_bytes`.
20pub const DEFAULT_OUTPUT_LIMIT: usize = 64 * 1024;
21/// Largest output limit honored.
22pub const MAX_OUTPUT_LIMIT: usize = 16 * 1024 * 1024;
23/// Memory cap when `[resources].ram_mb` is absent.
24pub const DEFAULT_RAM_MB: u32 = 256;
25/// Largest `ram_mb` honored.
26pub const MAX_RAM_MB: u32 = 64 * 1024;
27/// Longest single argument value accepted.
28pub const MAX_ARG_LEN: usize = 128;
29/// Most extra arguments accepted.
30pub const MAX_ARGS: usize = 32;
31
32/// Enforced `[console]` limits.
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct ConsoleLimits {
35    /// Exact argument strings a console run may use.
36    pub allowed_commands: Vec<String>,
37    /// Wall-clock cap for a console run.
38    pub max_runtime_secs: u64,
39    /// Captured bytes per stream.
40    pub output_limit_bytes: usize,
41}
42
43/// Enforced `[resources]`.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub struct Resources {
46    /// Memory cap in MiB.
47    pub ram_mb: u32,
48    /// CPU share in percent of one core (1..=400).
49    pub cpu_pct: u32,
50}
51
52/// Parsed, validated cog manifest subset.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct CogSpec {
55    /// `[cog].id`.
56    pub id: String,
57    /// `[cog].version`.
58    pub version: String,
59    /// Console limits.
60    pub console: ConsoleLimits,
61    /// Resource limits.
62    pub resources: Resources,
63    /// `cli_arg` flags that take a value.
64    pub value_args: BTreeSet<String>,
65    /// `cli_arg` flags of boolean options (no value).
66    pub flag_args: BTreeSet<String>,
67}
68
69#[derive(Deserialize)]
70struct RawToml {
71    cog: RawCog,
72    #[serde(default)]
73    console: Option<RawConsole>,
74    #[serde(default)]
75    resources: Option<RawResources>,
76    #[serde(default)]
77    config: BTreeMap<String, toml::Value>,
78}
79
80#[derive(Deserialize)]
81struct RawCog {
82    id: String,
83    version: String,
84}
85
86#[derive(Deserialize)]
87struct RawConsole {
88    #[serde(default)]
89    allowed_commands: Vec<String>,
90    max_runtime_secs: Option<u64>,
91    output_limit_bytes: Option<usize>,
92}
93
94#[derive(Deserialize)]
95struct RawResources {
96    ram_mb: Option<u32>,
97    cpu_pct: Option<u32>,
98}
99
100fn bad(msg: impl Into<String>) -> RuntimeError {
101    RuntimeError::InvalidConfig(msg.into())
102}
103
104fn valid_flag(s: &str) -> bool {
105    s.len() > 2
106        && s.len() <= 64
107        && s.starts_with("--")
108        && s[2..]
109            .bytes()
110            .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-')
111}
112
113/// A plain argument value: printable, no whitespace or shell metacharacters.
114pub fn valid_value(s: &str) -> bool {
115    !s.is_empty()
116        && s.len() <= MAX_ARG_LEN
117        && (!s.starts_with('-') || s.parse::<f64>().is_ok())
118        && s.bytes()
119            .all(|b| b.is_ascii_alphanumeric() || b"._:-+/".contains(&b))
120}
121
122impl CogSpec {
123    /// Parse `cog.toml` bytes.
124    pub fn parse(bytes: &[u8]) -> Result<Self, RuntimeError> {
125        let text = std::str::from_utf8(bytes).map_err(|_| bad("cog.toml is not UTF-8"))?;
126        let raw: RawToml = toml::from_str(text).map_err(|e| bad(format!("cog.toml: {e}")))?;
127        if !valid_cog_id(&raw.cog.id) {
128            return Err(bad(format!("bad cog id {:?}", raw.cog.id)));
129        }
130        let console = raw.console.unwrap_or(RawConsole {
131            allowed_commands: Vec::new(),
132            max_runtime_secs: None,
133            output_limit_bytes: None,
134        });
135        let max_runtime_secs = console.max_runtime_secs.unwrap_or(DEFAULT_MAX_RUNTIME_SECS);
136        if !(1..=MAX_RUNTIME_SECS).contains(&max_runtime_secs) {
137            return Err(bad(format!(
138                "max_runtime_secs must be 1..={MAX_RUNTIME_SECS}"
139            )));
140        }
141        let output_limit_bytes = console.output_limit_bytes.unwrap_or(DEFAULT_OUTPUT_LIMIT);
142        if !(1..=MAX_OUTPUT_LIMIT).contains(&output_limit_bytes) {
143            return Err(bad(format!(
144                "output_limit_bytes must be 1..={MAX_OUTPUT_LIMIT}"
145            )));
146        }
147        let res = raw.resources.unwrap_or(RawResources {
148            ram_mb: None,
149            cpu_pct: None,
150        });
151        let ram_mb = res.ram_mb.unwrap_or(DEFAULT_RAM_MB);
152        let cpu_pct = res.cpu_pct.unwrap_or(100);
153        if !(1..=MAX_RAM_MB).contains(&ram_mb) || !(1..=400).contains(&cpu_pct) {
154            return Err(bad(
155                "resources out of range (ram_mb 1..=65536, cpu_pct 1..=400)",
156            ));
157        }
158        let (mut value_args, mut flag_args) = (BTreeSet::new(), BTreeSet::new());
159        for (key, v) in &raw.config {
160            let Some(arg) = v.get("cli_arg").and_then(|a| a.as_str()) else {
161                continue;
162            };
163            if !valid_flag(arg) {
164                return Err(bad(format!("config.{key}: bad cli_arg {arg:?}")));
165            }
166            if v.get("type").and_then(|t| t.as_str()) == Some("boolean") {
167                flag_args.insert(arg.to_string());
168            } else {
169                value_args.insert(arg.to_string());
170            }
171        }
172        let spec = Self {
173            id: raw.cog.id,
174            version: raw.cog.version,
175            console: ConsoleLimits {
176                allowed_commands: console.allowed_commands,
177                max_runtime_secs,
178                output_limit_bytes,
179            },
180            resources: Resources { ram_mb, cpu_pct },
181            value_args,
182            flag_args,
183        };
184        for c in &spec.console.allowed_commands {
185            spec.split_command(c)?;
186        }
187        Ok(spec)
188    }
189
190    /// Validate place-time arguments against the `[config]` CLI surface.
191    pub fn validate_args(&self, args: &[String]) -> Result<(), RuntimeError> {
192        if args.len() > MAX_ARGS {
193            return Err(bad(format!("more than {MAX_ARGS} arguments")));
194        }
195        let mut it = args.iter();
196        while let Some(a) = it.next() {
197            if self.flag_args.contains(a) {
198                continue;
199            }
200            if a == "--once" || a == "--interval" {
201                return Err(bad(format!("{a} is set by the run mode, not by args")));
202            }
203            if !self.value_args.contains(a) {
204                return Err(bad(format!(
205                    "argument {a:?} is not a config cli_arg of this cog"
206                )));
207            }
208            match it.next() {
209                Some(v) if valid_value(v) => {}
210                _ => return Err(bad(format!("{a} needs a plain value"))),
211            }
212        }
213        Ok(())
214    }
215
216    /// Argv tail for a console command, which must be exactly one of
217    /// `[console].allowed_commands`.
218    pub fn console_args(&self, command: &str) -> Result<Vec<String>, RuntimeError> {
219        if !self.console.allowed_commands.iter().any(|c| c == command) {
220            return Err(bad(format!(
221                "console command {command:?} is not in [console].allowed_commands"
222            )));
223        }
224        self.split_command(command)
225    }
226
227    fn split_command(&self, command: &str) -> Result<Vec<String>, RuntimeError> {
228        let parts: Vec<String> = command.split_whitespace().map(str::to_string).collect();
229        if parts.len() > MAX_ARGS
230            || parts
231                .iter()
232                .any(|p| p.len() > MAX_ARG_LEN || p.bytes().any(|b| b.is_ascii_control()))
233        {
234            return Err(bad(format!("unsafe console command {command:?}")));
235        }
236        Ok(parts)
237    }
238}
239
240impl VerifiedWorkload {
241    /// Build a cog workload from a verified package. Every binary and the
242    /// `cog.toml` are read again through `source` and re-checked against
243    /// their pinned hashes, so the bytes handed to an adapter are the bytes
244    /// that were signed.
245    pub fn from_package(
246        pkg: &VerifiedPackage,
247        source: &dyn FileSource,
248    ) -> Result<Self, RuntimeError> {
249        let body = &pkg.body;
250        let read = |f| {
251            let bytes = source
252                .read(f)
253                .map_err(|e| RuntimeError::AdmissionRefused(e.to_string()))?;
254            check_file(f, &bytes).map_err(|e| RuntimeError::AdmissionRefused(e.to_string()))?;
255            Ok::<_, RuntimeError>(bytes)
256        };
257        let spec = CogSpec::parse(&read(&body.cog_toml)?)?;
258        if spec.id != body.id {
259            return Err(RuntimeError::AdmissionRefused(format!(
260                "cog.toml id {:?} does not match package id {:?}",
261                spec.id, body.id
262            )));
263        }
264        let mut binaries = BTreeMap::new();
265        for (arch, f) in &body.binaries {
266            binaries.insert(
267                arch.clone(),
268                Binary {
269                    blake3: f.blake3.clone(),
270                    bytes: Arc::new(read(f)?),
271                },
272            );
273        }
274        let signer_keys = pkg
275            .envelope
276            .signatures
277            .iter()
278            .filter(|s| pkg.signers.iter().any(|a| a.key_id == s.key_id))
279            .map(|s| s.public_key.clone())
280            .collect();
281        let mut artifact_hashes: Vec<String> = body.files().map(|f| f.blake3.clone()).collect();
282        artifact_hashes.sort();
283        artifact_hashes.dedup();
284        Ok(Self {
285            kind: pkg.envelope.kind.clone(),
286            id: body.id.clone(),
287            version: body.version.clone(),
288            source: WorkloadSource::SignedPackage(SignedPayload {
289                package_id: pkg.package_id.clone(),
290                signer_keys,
291                spec,
292                binaries,
293                artifact_hashes,
294            }),
295        })
296    }
297
298    /// A device-store cog reference. Only the Seed adapter accepts it, and
299    /// only when the operator pinned `(id, version)`.
300    pub fn store_pin(
301        registry: &str,
302        id: &str,
303        version: &str,
304        sha256: Option<&str>,
305    ) -> Result<Self, RuntimeError> {
306        if !valid_cog_id(id) || !valid_value(version) || !valid_cog_id(registry) {
307            return Err(bad("store pin needs a valid registry, cog id and version"));
308        }
309        if let Some(h) = sha256
310            && !crate::workload_pkg::codec::is_lower_hex(h, 64)
311        {
312            return Err(bad("store pin sha256 must be 64 lower-case hex"));
313        }
314        Ok(Self {
315            kind: crate::workload_pkg::KIND_COG.to_string(),
316            id: id.to_string(),
317            version: version.to_string(),
318            source: WorkloadSource::StorePin {
319                registry: registry.to_string(),
320                sha256: sha256.map(str::to_string),
321            },
322        })
323    }
324
325    /// What the revocation list is checked against for this workload:
326    /// `(package id, signer keys, artifact hashes)`. Never all empty: a
327    /// store pin or project workload gets a synthetic package id, so a
328    /// request built from any workload names something.
329    pub fn revocation_refs(&self) -> (String, Vec<String>, Vec<String>) {
330        match &self.source {
331            WorkloadSource::SignedPackage(p) => (
332                p.package_id.clone(),
333                p.signer_keys.clone(),
334                p.artifact_hashes.clone(),
335            ),
336            WorkloadSource::StorePin { .. } => (
337                format!("store.{}.{}", self.id, self.version),
338                Vec::new(),
339                Vec::new(),
340            ),
341            WorkloadSource::Project(p) => (
342                format!("project.{}.{}", p.project_id, p.cert_serial),
343                Vec::new(),
344                Vec::new(),
345            ),
346            WorkloadSource::Inference(_) => (
347                format!("inference.{}.{}", self.id, self.version),
348                Vec::new(),
349                Vec::new(),
350            ),
351        }
352    }
353
354    /// The signed payload, or an admission refusal naming the adapter.
355    pub fn signed(&self, runtime: &str) -> Result<&SignedPayload, RuntimeError> {
356        match &self.source {
357            WorkloadSource::SignedPackage(p) => Ok(p),
358            WorkloadSource::StorePin { .. } => Err(RuntimeError::AdmissionRefused(format!(
359                "{runtime} runs only signed packages; store pins go through the Seed adapter"
360            ))),
361            WorkloadSource::Project(_) => Err(RuntimeError::AdmissionRefused(format!(
362                "{runtime} runs only signed packages; project workloads are not supported yet"
363            ))),
364            WorkloadSource::Inference(_) => Err(RuntimeError::AdmissionRefused(format!(
365                "{runtime} runs only signed packages; inference workloads go through the infer.* adapters"
366            ))),
367        }
368    }
369}
clawft_kernel/workload_runtime/cog_spec.rs

clawft_kernel/workload_runtime/
cog_spec.rs
