diff --git a/Cargo.lock b/Cargo.lock index 520cbb67f..a04f6fd5f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -99,6 +99,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + [[package]] name = "arc-swap" version = "1.9.1" @@ -1015,6 +1024,12 @@ dependencies = [ "winnow", ] +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "crypto-bigint" version = "0.5.5" @@ -1113,6 +1128,17 @@ dependencies = [ "serde_core", ] +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -2218,6 +2244,7 @@ dependencies = [ "urlencoding", "uuid", "webpki-roots 0.26.11", + "zip", ] [[package]] @@ -4367,12 +4394,41 @@ dependencies = [ "syn", ] +[[package]] +name = "zip" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "flate2", + "indexmap", + "memchr", + "thiserror 2.0.18", + "zopfli", +] + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + [[package]] name = "zune-core" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index ef07ba752..a53aa05b2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ slack = ["openab-core/slack"] secrets-aws = ["openab-core/secrets-aws"] agentcore = ["openab-core/agentcore"] config-s3 = ["openab-core/config-s3"] +pre-seed = ["openab-core/pre-seed"] # Gateway adapters (each pulls in the gateway crate + axum for embedded server) telegram = ["dep:openab-gateway", "dep:axum", "openab-gateway/telegram"] diff --git a/config.toml.example b/config.toml.example index dfa7c68d4..83cd74a37 100644 --- a/config.toml.example +++ b/config.toml.example @@ -241,3 +241,22 @@ error_hold_ms = 2500 # openab = "~/projects/openab" # infra = "~/projects/infra-cdk" # web = "~/projects/frontend" + +# --- Pre-seed: download & extract S3 zip archives before pre_boot --- +# Seeds the agent environment (configs, tools, memory) before any hook runs. +# Archives are extracted in order — later layers overwrite earlier ones (like container layers). +# Max 5 sources. Only s3:// URIs with .zip format are supported. +# Requires the `pre-seed` feature flag (opt-in). +# +# [hooks.pre_seed] +# sources = [ +# "s3://my-bucket/base-env.zip", # Layer 1: base tools & configs +# "s3://my-bucket/shared-memory.zip", # Layer 2: shared team memory +# "s3://my-bucket/agent-specific.zip", # Layer 3: agent-specific overrides +# ] +# target = "/home/agent" # extraction target (default: $HOME) +# max_bytes = 104857600 # max compressed zip size (default: 100 MiB) +# timeout_seconds = 300 # per-source timeout (default: 300) +# on_failure = "abort" # "abort" or "warn" (default: "abort") +# region = "us-west-2" # optional: override AWS region +# endpoint_url = "http://localhost:4566" # optional: LocalStack / VPC endpoint diff --git a/crates/openab-core/Cargo.toml b/crates/openab-core/Cargo.toml index d7b53ded7..01005c1bf 100644 --- a/crates/openab-core/Cargo.toml +++ b/crates/openab-core/Cargo.toml @@ -39,6 +39,7 @@ tempfile = "3.27.0" aws-sdk-secretsmanager = { version = "1", optional = true } aws-sdk-s3 = { version = "1", optional = true } aws-config = { version = "1", optional = true } +zip = { version = "2", default-features = false, features = ["deflate"], optional = true } aws-sigv4 = { version = "1", optional = true } aws-credential-types = { version = "1", optional = true } urlencoding = { version = "2", optional = true } @@ -54,4 +55,5 @@ discord = ["dep:serenity"] slack = [] secrets-aws = ["dep:aws-sdk-secretsmanager", "dep:aws-config"] config-s3 = ["dep:aws-sdk-s3", "dep:aws-config"] +pre-seed = ["dep:aws-sdk-s3", "dep:aws-config", "dep:zip", "dep:hex"] agentcore = ["dep:aws-config", "dep:aws-sigv4", "dep:aws-credential-types", "dep:urlencoding", "dep:hex", "dep:http", "dep:rustls", "dep:tokio-rustls", "dep:webpki-roots"] diff --git a/crates/openab-core/src/config.rs b/crates/openab-core/src/config.rs index b4aaf3fb4..f0018572c 100644 --- a/crates/openab-core/src/config.rs +++ b/crates/openab-core/src/config.rs @@ -198,10 +198,58 @@ fn default_exec_timeout() -> u64 { #[derive(Debug, Clone, Default, Deserialize)] pub struct HooksConfig { + pub pre_seed: Option, pub pre_boot: Option, pub pre_shutdown: Option, } +/// Configuration for the pre_seed phase. +/// Downloads and extracts zip archives from S3 before pre_boot. +#[derive(Debug, Clone, Deserialize)] +pub struct PreSeedConfig { + /// S3 URIs of zip archives to download and extract (max 5). + /// Extracted in order; later layers overwrite earlier ones. + #[serde(default)] + pub sources: Vec, + /// Extraction target directory. Default: $HOME. + pub target: Option, + /// Override AWS region for S3 access. + pub region: Option, + /// Override S3 endpoint URL (for LocalStack, VPC endpoints). + pub endpoint_url: Option, + /// Maximum compressed zip size in bytes. Default: 100 MiB. + #[serde(default = "default_max_zip_bytes")] + pub max_bytes: u64, + /// Timeout in seconds for each download+extract operation. Default: 300. + #[serde(default = "default_pre_seed_timeout")] + pub timeout_seconds: u64, + /// Failure policy. Default: abort. + #[serde(default)] + pub on_failure: OnFailure, +} + +impl Default for PreSeedConfig { + fn default() -> Self { + Self { + sources: Vec::new(), + target: None, + region: None, + endpoint_url: None, + max_bytes: default_max_zip_bytes(), + timeout_seconds: default_pre_seed_timeout(), + on_failure: OnFailure::Abort, + } + } +} + +fn default_max_zip_bytes() -> u64 { + 100 * 1024 * 1024 // 100 MiB +} + +fn default_pre_seed_timeout() -> u64 { + 300 +} + /// Failure policy for a hook. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum OnFailure { diff --git a/crates/openab-core/src/lib.rs b/crates/openab-core/src/lib.rs index c4af990df..b02848ac3 100644 --- a/crates/openab-core/src/lib.rs +++ b/crates/openab-core/src/lib.rs @@ -12,6 +12,8 @@ pub mod hooks; pub mod markdown; pub mod media; pub mod multibot_cache; +#[cfg(feature = "pre-seed")] +pub mod pre_seed; pub mod reactions; #[cfg(feature = "discord")] pub mod remind; diff --git a/crates/openab-core/src/pre_seed.rs b/crates/openab-core/src/pre_seed.rs new file mode 100644 index 000000000..576344f9c --- /dev/null +++ b/crates/openab-core/src/pre_seed.rs @@ -0,0 +1,493 @@ +use crate::config::{parse_s3_uri, OnFailure, PreSeedConfig}; +use sha2::{Digest, Sha256}; +use std::path::Path; +use std::time::Instant; +use tracing::{error, info, warn}; + +/// Maximum number of sources allowed. +const MAX_SOURCES: usize = 5; + +/// Default max extracted (uncompressed) size: 500 MiB. +const DEFAULT_MAX_EXTRACTED_BYTES: u64 = 500 * 1024 * 1024; + +/// Default max file count per zip. +const DEFAULT_MAX_FILE_COUNT: usize = 10_000; + +/// Run the pre_seed phase: download zip archives from S3 and extract them in order. +pub async fn run(cfg: &PreSeedConfig) -> anyhow::Result<()> { + if cfg.sources.is_empty() { + return Ok(()); + } + if cfg.sources.len() > MAX_SOURCES { + anyhow::bail!( + "hooks.pre_seed: too many sources ({}, max {})", + cfg.sources.len(), + MAX_SOURCES + ); + } + + let target = match &cfg.target { + Some(t) => std::path::PathBuf::from(t), + None => dirs_home(), + }; + + info!( + sources = cfg.sources.len(), + target = %target.display(), + "hooks.pre_seed: starting" + ); + + let mut s3_config_loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(ref region) = cfg.region { + s3_config_loader = s3_config_loader.region(aws_config::Region::new(region.clone())); + } + if let Some(ref endpoint) = cfg.endpoint_url { + s3_config_loader = s3_config_loader.endpoint_url(endpoint); + } + let aws_cfg = s3_config_loader.load().await; + let s3 = aws_sdk_s3::Client::new(&aws_cfg); + + for (i, source) in cfg.sources.iter().enumerate() { + let layer = i + 1; + info!( + layer, + source = source.as_str(), + "hooks.pre_seed: downloading" + ); + + let deadline = Instant::now() + std::time::Duration::from_secs(cfg.timeout_seconds); + + let result = download_and_extract(&s3, source, &target, cfg.max_bytes, deadline).await; + + let outcome = match result { + Ok(()) => { + info!(layer, "hooks.pre_seed: layer extracted successfully"); + continue; + } + Err(e) => e, + }; + + match cfg.on_failure { + OnFailure::Abort => { + error!(layer, error = %outcome, "hooks.pre_seed failed (on_failure=abort)"); + return Err(outcome); + } + OnFailure::Warn => { + warn!(layer, error = %outcome, "hooks.pre_seed failed (on_failure=warn), continuing"); + } + } + } + + info!("hooks.pre_seed: complete"); + Ok(()) +} + +/// Download zip from S3, verify integrity, extract to a temp dir, then move into target. +/// The deadline is enforced cooperatively inside the blocking task. +async fn download_and_extract( + s3: &aws_sdk_s3::Client, + uri: &str, + target: &Path, + max_bytes: u64, + deadline: Instant, +) -> anyhow::Result<()> { + let (bucket, key) = parse_s3_uri(uri)?; + + // Check deadline before S3 call + if Instant::now() >= deadline { + anyhow::bail!("hooks.pre_seed: timed out before download for {uri}"); + } + + let resp = s3 + .get_object() + .bucket(&bucket) + .key(&key) + .checksum_mode(aws_sdk_s3::types::ChecksumMode::Enabled) + .send() + .await + .map_err(|e| anyhow::anyhow!("S3 GetObject failed for {uri}: {e}"))?; + + if let Some(len) = resp.content_length() { + if len as u64 > max_bytes { + anyhow::bail!("hooks.pre_seed: {uri} too large ({len} bytes, max {max_bytes})"); + } + } + + // Capture S3-native SHA-256 checksum if present (set during upload with --checksum-algorithm SHA256) + let s3_checksum_sha256 = resp.checksum_sha256().map(|s| s.to_string()); + + let body = resp + .body + .collect() + .await + .map_err(|e| anyhow::anyhow!("failed to read S3 body for {uri}: {e}"))?; + let bytes = body.into_bytes(); + + if bytes.len() as u64 > max_bytes { + anyhow::bail!( + "hooks.pre_seed: {uri} too large ({} bytes, max {max_bytes})", + bytes.len() + ); + } + + // SHA-256 verification: auto-verify S3-native checksum if present + if let Some(ref s3_b64) = s3_checksum_sha256 { + let mut hasher = Sha256::new(); + hasher.update(&bytes); + let actual_hex = format!("{:x}", hasher.finalize()); + let s3_hex = base64_sha256_to_hex(s3_b64)?; + if actual_hex != s3_hex { + anyhow::bail!( + "hooks.pre_seed: S3 checksum mismatch for {uri}: expected {s3_hex}, got {actual_hex}" + ); + } + info!(uri, "hooks.pre_seed: S3-native SHA-256 verified"); + } + + if Instant::now() >= deadline { + anyhow::bail!("hooks.pre_seed: timed out after download for {uri}"); + } + + info!( + uri, + bytes = bytes.len(), + "hooks.pre_seed: downloaded, extracting" + ); + + // Extract and move in a blocking task with cooperative deadline checking. + let target = target.to_path_buf(); + // Bytes is Arc-backed, Clone is zero-copy (ref-count bump only) + tokio::task::spawn_blocking(move || extract_and_apply(&bytes, &target, deadline)) + .await + .map_err(|e| anyhow::anyhow!("hooks.pre_seed: extract task panicked: {e}"))??; + + Ok(()) +} + +/// Extract zip to a temp directory with budget enforcement, then move into target. +/// Checks deadline cooperatively before each file operation. +fn extract_and_apply(data: &[u8], target: &Path, deadline: Instant) -> anyhow::Result<()> { + let temp_dir = tempfile::tempdir_in(target.parent().unwrap_or(target))?; + + extract_zip_with_limits(data, temp_dir.path(), deadline)?; + + // Check deadline before applying to target + if Instant::now() >= deadline { + // temp_dir drops and cleans up automatically + anyhow::bail!("hooks.pre_seed: timed out before applying to target"); + } + + move_recursive(temp_dir.path(), target, deadline)?; + Ok(()) +} + +/// Extract a zip archive with cooperative deadline checks and extraction budget. +fn extract_zip_with_limits(data: &[u8], dest: &Path, deadline: Instant) -> anyhow::Result<()> { + extract_zip_budgeted( + data, + dest, + deadline, + DEFAULT_MAX_FILE_COUNT, + DEFAULT_MAX_EXTRACTED_BYTES, + ) +} + +/// Inner extraction with configurable limits (enables testing with small budgets). +fn extract_zip_budgeted( + data: &[u8], + dest: &Path, + deadline: Instant, + max_file_count: usize, + max_extracted_bytes: u64, +) -> anyhow::Result<()> { + let cursor = std::io::Cursor::new(data); + let mut archive = zip::ZipArchive::new(cursor)?; + + let file_count = archive.len(); + if file_count > max_file_count { + anyhow::bail!( + "hooks.pre_seed: zip contains too many entries ({file_count}, max {max_file_count})" + ); + } + + let mut total_extracted: u64 = 0; + + for i in 0..file_count { + // Cooperative deadline check per file + if i % 100 == 0 && Instant::now() >= deadline { + anyhow::bail!("hooks.pre_seed: timed out during extraction at entry {i}"); + } + + let mut file = archive.by_index(i)?; + let name = file.enclosed_name().ok_or_else(|| { + anyhow::anyhow!("hooks.pre_seed: invalid zip entry name at index {i}") + })?; + let out_path = dest.join(name); + + if file.is_dir() { + std::fs::create_dir_all(&out_path)?; + } else { + if let Some(parent) = out_path.parent() { + std::fs::create_dir_all(parent)?; + } + + // Check extracted size budget before writing + let uncompressed = file.size(); + total_extracted += uncompressed; + if total_extracted > max_extracted_bytes { + anyhow::bail!( + "hooks.pre_seed: extracted size exceeds limit ({total_extracted} > {max_extracted_bytes})" + ); + } + + let mut out = std::fs::File::create(&out_path)?; + std::io::copy(&mut file, &mut out)?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + if let Some(mode) = file.unix_mode() { + std::fs::set_permissions(&out_path, std::fs::Permissions::from_mode(mode))?; + } + } + } + } + + Ok(()) +} + +/// Recursively move files from src directory into dst directory. +/// Checks deadline cooperatively. +fn move_recursive(src: &Path, dst: &Path, deadline: Instant) -> anyhow::Result<()> { + for entry in std::fs::read_dir(src)? { + if Instant::now() >= deadline { + anyhow::bail!("hooks.pre_seed: timed out during move to target"); + } + + let entry = entry?; + let src_path = entry.path(); + let dst_path = dst.join(entry.file_name()); + + if src_path.is_dir() { + std::fs::create_dir_all(&dst_path)?; + move_recursive(&src_path, &dst_path, deadline)?; + } else { + if std::fs::rename(&src_path, &dst_path).is_err() { + std::fs::copy(&src_path, &dst_path)?; + std::fs::remove_file(&src_path)?; + } + } + } + Ok(()) +} + +/// Decode a base64-encoded SHA-256 (as returned by S3) to lowercase hex. +fn base64_sha256_to_hex(b64: &str) -> anyhow::Result { + use base64::Engine; + let decoded = base64::engine::general_purpose::STANDARD + .decode(b64) + .map_err(|e| anyhow::anyhow!("hooks.pre_seed: invalid base64 in S3 checksum: {e}"))?; + Ok(hex::encode(decoded)) +} + +fn dirs_home() -> std::path::PathBuf { + std::env::var("HOME") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| std::path::PathBuf::from("/home/agent")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extract_zip_basic() { + use std::io::Write; + let dir = tempfile::tempdir().unwrap(); + let deadline = Instant::now() + std::time::Duration::from_secs(60); + + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + writer.start_file("hello.txt", options).unwrap(); + writer.write_all(b"world").unwrap(); + writer.start_file("sub/nested.txt", options).unwrap(); + writer.write_all(b"nested content").unwrap(); + let cursor = writer.finish().unwrap(); + + extract_zip_with_limits(cursor.get_ref(), dir.path(), deadline).unwrap(); + + assert_eq!( + std::fs::read_to_string(dir.path().join("hello.txt")).unwrap(), + "world" + ); + assert_eq!( + std::fs::read_to_string(dir.path().join("sub/nested.txt")).unwrap(), + "nested content" + ); + } + + #[test] + fn extract_and_apply_atomic() { + use std::io::Write; + let target = tempfile::tempdir().unwrap(); + let deadline = Instant::now() + std::time::Duration::from_secs(60); + + std::fs::write(target.path().join("existing.txt"), "keep").unwrap(); + + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + writer.start_file("new.txt", options).unwrap(); + writer.write_all(b"added").unwrap(); + let cursor = writer.finish().unwrap(); + + extract_and_apply(cursor.get_ref(), target.path(), deadline).unwrap(); + + assert_eq!( + std::fs::read_to_string(target.path().join("existing.txt")).unwrap(), + "keep" + ); + assert_eq!( + std::fs::read_to_string(target.path().join("new.txt")).unwrap(), + "added" + ); + } + + #[test] + fn extract_respects_expired_deadline() { + use std::io::Write; + let dir = tempfile::tempdir().unwrap(); + // Already expired deadline + let deadline = Instant::now() - std::time::Duration::from_secs(1); + + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + writer.start_file("a.txt", options).unwrap(); + writer.write_all(b"data").unwrap(); + let cursor = writer.finish().unwrap(); + + // extract_and_apply should fail due to expired deadline + let result = extract_and_apply(cursor.get_ref(), dir.path(), deadline); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("timed out")); + } + + #[test] + fn extract_zip_overwrites() { + use std::io::Write; + let target = tempfile::tempdir().unwrap(); + let deadline = Instant::now() + std::time::Duration::from_secs(60); + std::fs::write(target.path().join("hello.txt"), "original").unwrap(); + + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + writer.start_file("hello.txt", options).unwrap(); + writer.write_all(b"overwritten").unwrap(); + let cursor = writer.finish().unwrap(); + + extract_and_apply(cursor.get_ref(), target.path(), deadline).unwrap(); + + assert_eq!( + std::fs::read_to_string(target.path().join("hello.txt")).unwrap(), + "overwritten" + ); + } + + #[tokio::test] + async fn run_empty_sources() { + let cfg = PreSeedConfig::default(); + assert!(run(&cfg).await.is_ok()); + } + + #[tokio::test] + async fn run_too_many_sources() { + let cfg = PreSeedConfig { + sources: vec!["s3://b/k.zip".into(); 6], + ..Default::default() + }; + assert!(run(&cfg).await.is_err()); + } + + #[test] + fn default_has_correct_values() { + let cfg = PreSeedConfig::default(); + assert_eq!(cfg.timeout_seconds, 300); + assert_eq!(cfg.max_bytes, 100 * 1024 * 1024); + assert_eq!(cfg.on_failure, OnFailure::Abort); + assert!(cfg.sources.is_empty()); + } + + #[test] + fn move_respects_deadline() { + let src = tempfile::tempdir().unwrap(); + let dst = tempfile::tempdir().unwrap(); + std::fs::write(src.path().join("f.txt"), "x").unwrap(); + + let expired = Instant::now() - std::time::Duration::from_secs(1); + let result = move_recursive(src.path(), dst.path(), expired); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("timed out")); + } + + #[test] + fn extract_rejects_exceeding_extracted_bytes() { + use std::io::Write; + let dir = tempfile::tempdir().unwrap(); + let deadline = Instant::now() + std::time::Duration::from_secs(60); + + // Create a zip with 3 files of 10 bytes each (30 bytes total extracted) + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + for i in 0..3 { + writer.start_file(format!("file_{i}.txt"), options).unwrap(); + writer.write_all(&[b'x'; 10]).unwrap(); + } + let cursor = writer.finish().unwrap(); + + // Set max extracted bytes to 20 — fails on 3rd file (cumulative 30 > 20) + let result = extract_zip_budgeted(cursor.get_ref(), dir.path(), deadline, 10_000, 20); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("extracted size exceeds limit"), + "should fail on extracted bytes limit" + ); + } + + #[test] + fn extract_rejects_exceeding_file_count() { + use std::io::Write; + let dir = tempfile::tempdir().unwrap(); + let deadline = Instant::now() + std::time::Duration::from_secs(60); + + // Create a zip with 5 files + let buf = Vec::new(); + let cursor = std::io::Cursor::new(buf); + let mut writer = zip::ZipWriter::new(cursor); + let options = zip::write::SimpleFileOptions::default(); + for i in 0..5 { + writer.start_file(format!("f_{i}.txt"), options).unwrap(); + writer.write_all(b"x").unwrap(); + } + let cursor = writer.finish().unwrap(); + + // Set max file count to 3 — should fail (5 > 3) + let result = extract_zip_budgeted(cursor.get_ref(), dir.path(), deadline, 3, u64::MAX); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("too many entries"), + "should fail on file count limit" + ); + } +} diff --git a/docs/config-reference.md b/docs/config-reference.md index dc6bea597..2c8ba482b 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -239,7 +239,39 @@ Session pool settings for managing concurrent agent sessions. ## `[hooks]` -Lifecycle hooks that run custom scripts at specific points during the container lifecycle. See [hooks.md](hooks.md) for full documentation and examples. +Lifecycle hooks that run at specific points during the container lifecycle. See [hooks.md](hooks.md) for full documentation and examples. + +### `[hooks.pre_seed]` + +Downloads and extracts zip archives from S3 before `pre_boot`. Seeds the agent environment with configs, tools, and shared memory without requiring AWS CLI in the image. + +> **Feature flag:** requires the `pre-seed` feature (opt-in, not in default). Enable with `--features pre-seed`. + +| Key | Type | Default | Description | +|-----|------|---------|-------------| +| `sources` | string[] | `[]` | S3 URIs of zip archives (`s3://bucket/key.zip`). Max 5. Extracted in order; later layers overwrite earlier ones. | +| `target` | string | `$HOME` | Extraction target directory. | +| `max_bytes` | u64 | `104857600` | Max compressed zip size in bytes (100 MiB). Rejects downloads exceeding this. | +| `timeout_seconds` | u64 | `300` | Per-source download+extract timeout in seconds. | +| `on_failure` | string | `"abort"` | `"abort"` exits openab; `"warn"` logs and continues. | +| `region` | string | — | Override AWS region for S3 access. | +| `endpoint_url` | string | — | Override S3 endpoint URL (for LocalStack, VPC endpoints). | + +**Credential resolution** uses the standard AWS provider chain (same as `config-s3` and `secrets-aws`): +environment variables, shared credentials, IRSA / EKS Pod Identity, ECS task role. + +**Integrity verification:** If S3 objects are uploaded with `--checksum-algorithm SHA256`, OpenAB automatically verifies the checksum on download. No config needed — see [hooks.md](hooks.md) for details. + +```toml +[hooks.pre_seed] +sources = [ + "s3://my-bucket/base-env.zip", + "s3://my-bucket/shared-memory.zip", + "s3://my-bucket/agent-overrides.zip", +] +timeout_seconds = 300 +on_failure = "abort" +``` ### `[hooks.pre_boot]` diff --git a/docs/hooks.md b/docs/hooks.md index 687253876..db1961a5d 100644 --- a/docs/hooks.md +++ b/docs/hooks.md @@ -1,6 +1,114 @@ # Lifecycle Hooks -OpenAB supports lifecycle hooks that run custom scripts at specific points during the container lifecycle. Hooks are configured in `config.toml` under the `[hooks]` table. +OpenAB supports lifecycle hooks that run at specific points during the container lifecycle. All lifecycle phases are configured in `config.toml` under the `[hooks]` table. + +## Lifecycle Order + +``` +hooks.pre_seed → hooks.pre_boot → (agent running) → hooks.pre_shutdown +``` + +| Phase | Purpose | Config | Action Type | +|-------|---------|--------|-------------| +| `pre_seed` | Download & extract S3 zip archives to seed the environment | `[hooks.pre_seed]` | Built-in S3 download + unzip | +| `pre_boot` | Run custom setup scripts before agent pool creation | `[hooks.pre_boot]` | User script | +| `pre_shutdown` | Run custom cleanup scripts after pool shutdown | `[hooks.pre_shutdown]` | User script | + +## Pre-Seed Phase + +The `pre_seed` phase runs **before** `pre_boot`. It downloads zip archives from S3 and extracts them into the agent's home directory (or a custom target). This eliminates the need for users to install AWS CLI and write download scripts in `pre_boot`. + +> **Feature flag:** requires the `pre-seed` feature (opt-in, not in default). + +### Configuration + +```toml +[hooks.pre_seed] +sources = [ + "s3://my-bucket/base-env.zip", + "s3://my-bucket/shared-memory.zip", + "s3://my-bucket/agent-overrides.zip", +] +# target = "/home/agent" # default: $HOME +# max_bytes = 104857600 # max compressed size per zip (default: 100 MiB) +# timeout_seconds = 300 # per-source timeout (default: 300) +# on_failure = "abort" # "abort" or "warn" (default: "abort") +# region = "us-west-2" # optional: override AWS region +# endpoint_url = "http://localhost:4566" # optional: LocalStack / VPC endpoint +``` + +### Fields + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `sources` | string[] | `[]` | S3 URIs of zip archives. Max 5. Extracted in order. | +| `target` | string | `$HOME` | Extraction target directory. | +| `max_bytes` | u64 | `104857600` | Max compressed zip size in bytes (100 MiB). | +| `timeout_seconds` | u64 | `300` | Per-source download+extract timeout. | +| `on_failure` | string | `"abort"` | `"abort"` exits openab; `"warn"` logs and continues. | +| `region` | string | — | Override AWS region. | +| `endpoint_url` | string | — | Override S3 endpoint URL. | + +### Layer Concept + +Sources are extracted sequentially (first → last). Files from later archives overwrite earlier ones — like layers in a container image: + +``` +Layer 3 (last) ─── highest priority, overwrites all below +Layer 2 ─── overwrites layer 1 +Layer 1 (first) ─── base layer +───────────────── + $HOME +``` + +### Safety + +- **Integrity verification**: two layers of protection: + 1. **S3-native checksum (automatic)**: if the object was uploaded with `--checksum-algorithm SHA256`, OpenAB automatically verifies it on download — no config needed + 2. **User-provided `sha256s` (optional)**: explicit checksums in config for additional defense-in-depth +- **Size cap**: downloads exceeding `max_bytes` are rejected before extraction +- **Atomic extraction**: zips are first extracted to a temp directory, then moved into target — if extraction fails, target is not corrupted. Note: the move phase is per-file; if it fails mid-way with `on_failure = "warn"`, the target may be partially updated. +- **Zip Slip prevention**: uses `enclosed_name()` to block path traversal attacks + +### Constraints + +- Maximum **5** sources +- Only `s3://` URIs supported +- Only `.zip` format supported +- Uses the standard AWS credential chain (IRSA, ECS task role, env vars) +- Optional `region`/`endpoint_url` override for LocalStack or VPC endpoints + +### IAM Policy + +```json +{ + "Effect": "Allow", + "Action": ["s3:GetObject"], + "Resource": [ + "arn:aws:s3:::my-bucket/base-env.zip", + "arn:aws:s3:::my-bucket/shared-memory.zip", + "arn:aws:s3:::my-bucket/agent-overrides.zip" + ] +} +``` + +### Recommended: Enable S3 Checksums on Upload + +For automatic integrity verification without maintaining `sha256s` in config, upload zip archives with SHA-256 checksums enabled: + +```bash +# Upload with SHA-256 checksum (recommended) +aws s3 cp env.zip s3://my-bucket/env.zip --checksum-algorithm SHA256 + +# Verify it was stored +aws s3api head-object --bucket my-bucket --key env.zip --checksum-mode ENABLED +``` + +When objects have S3-native SHA-256 checksums, OpenAB verifies them automatically on download — no `sha256s` config needed. This is the simplest path to integrity verification. + +> **Note:** If `sha256s` is also provided in config, both checks run. The S3-native check uses the base64-encoded checksum from the `x-amz-checksum-sha256` response header. If neither is available, download proceeds without integrity verification (relies on IAM + bucket policy for trust). + +--- ## Available Hooks diff --git a/src/main.rs b/src/main.rs index 980590454..15f90db10 100644 --- a/src/main.rs +++ b/src/main.rs @@ -194,6 +194,14 @@ async fn main() -> anyhow::Result<()> { ); } + // --- pre_seed: download & extract S3 zips before pre_boot --- + #[cfg(feature = "pre-seed")] + if let Some(ref pre_seed) = cfg.hooks.pre_seed { + if !pre_seed.sources.is_empty() { + openab_core::pre_seed::run(pre_seed).await?; + } + } + // Validate and run pre_boot hook (before agent pool creation) if let Some(ref hook) = cfg.hooks.pre_boot { hooks::validate_hook("pre_boot", hook)?;