From 43ec1b6742c77cc9116fb425d286ce90d3a0c440 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Wed, 14 Oct 2026 18:47:51 +0000 Subject: [PATCH 01/13] trident-acl-agent: persist Nebraska affinity cookie across requests Nebraska issues a load-balancer affinity cookie (e.g. an App-Service-style ASLBSA cookie) via Set-Cookie on any request that does not already carry one, and expects it echoed back on every later request/response. - Transport::post_xml now takes an optional outgoing cookie and returns any captured Set-Cookie value (raw name=value pairs only, attributes like Path/Secure dropped), so ReqwestTransport sends/receives it over the wire without relying on reqwest's short-lived built-in cookie jar. - Client carries the cookie across its own calls via interior mutability (RefCell), with with_cookie()/cookie() to seed/read it across process and reboot boundaries, since callers (check_for_update, report_progress, complete_after_reboot, report_failure) build a fresh Client per call. - New CookieStore (core::nebraska::cookie_store) persists the cookie to its own file (/var/lib/trident-acl-agent/nebraska-cookie by default), deliberately separate from annotations::state's state.json: that file is purpose-built to bridge one update operation's pre-/post-reboot halves, while the cookie is a host-level attribute unrelated to any operation. - Orchestrator loads/saves the cookie around both Nebraska call sites (handle_stage's check_for_update, report_nebraska_event's progress/terminal reports), best-effort: a load/save failure only logs a warning and proceeds cookie-less, never fails the underlying update. - New OrchestrationConfig::cookie_path (TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH), defaulting to /var/lib/trident-acl-agent/nebraska-cookie. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6c26d0a1-b43f-4ccb-a684-0d92c41f1321 --- .../src/annotations/orchestrator.rs | 67 +++++- crates/trident-acl-agent/src/core/config.rs | 31 +++ .../src/core/nebraska/client.rs | 131 +++++++++++- .../src/core/nebraska/cookie_store.rs | 193 ++++++++++++++++++ .../src/core/nebraska/mod.rs | 2 + .../src/core/nebraska/transport.rs | 183 ++++++++++++++++- 6 files changed, 577 insertions(+), 30 deletions(-) create mode 100644 crates/trident-acl-agent/src/core/nebraska/cookie_store.rs diff --git a/crates/trident-acl-agent/src/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index cfd5fc3090..7ab36f7758 100644 --- a/crates/trident-acl-agent/src/annotations/orchestrator.rs +++ b/crates/trident-acl-agent/src/annotations/orchestrator.rs @@ -35,7 +35,7 @@ use crate::{ core::{ config::AgentConfig, cosi_sha384, - nebraska::{CheckOutcome, Client as NebraskaClient, ProgressEvent}, + nebraska::{CheckOutcome, Client as NebraskaClient, CookieStore, ProgressEvent}, trident::{CompletedResponse, TridentClient, TridentClientError}, version::{current_active_version, FALLBACK_ALWAYS_VERSION}, }, @@ -109,6 +109,10 @@ pub struct Orchestrator { config: AgentConfig, k8s: NodeClient, state: StateStore, + /// Persists Nebraska's load-balancer affinity cookie across requests + /// and reboots. Deliberately separate from `state`: see + /// `core::nebraska::CookieStore`'s docs for why. + nebraska_cookie: CookieStore, annotation_keys: AnnotationKeys, } @@ -118,6 +122,7 @@ impl Orchestrator { let annotation_keys = AnnotationKeys::new(&config.kubernetes.annotation_prefix); Ok(Self { state: StateStore::new(config.orchestration.state_path.clone()), + nebraska_cookie: CookieStore::new(config.orchestration.cookie_path.clone()), config, k8s, annotation_keys, @@ -701,16 +706,22 @@ impl Orchestrator { Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); - let outcome = task::spawn_blocking({ + let initial_cookie = self.load_nebraska_cookie(); + let (check_result, cookie) = task::spawn_blocking({ let current_version = current_version.clone(); move || { - let client = NebraskaClient::new(endpoint, app_id, track, machine_id); - client.check_for_update(¤t_version) + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + if let Some(cookie) = initial_cookie { + client = client.with_cookie(cookie); + } + let result = client.check_for_update(¤t_version); + (result, client.cookie()) } }) .await - .context("Nebraska query task panicked")? - .context("Nebraska query failed")?; + .context("Nebraska query task panicked")?; + self.save_nebraska_cookie(cookie); + let outcome = check_result.context("Nebraska query failed")?; let offered = match outcome { CheckOutcome::UpToDate => { let status = UpdateStatus::new( @@ -1569,6 +1580,30 @@ impl Orchestrator { } } + /// Loads the persisted Nebraska affinity cookie, warning (not failing) + /// if the store can't be read. The whole mechanism is a best-effort + /// session-affinity hint, never a requirement for talking to Nebraska - + /// a load failure just means this round-trip proceeds cookie-less, the + /// same as this host's very first-ever contact with Nebraska. + fn load_nebraska_cookie(&self) -> Option { + match self.nebraska_cookie.load() { + Ok(cookie) => cookie, + Err(err) => { + warn!("failed to load Nebraska cookie, proceeding without one: {err}"); + None + } + } + } + + /// Persists a cookie a Nebraska round-trip ended up holding, if any. + fn save_nebraska_cookie(&self, cookie: Option) { + if let Some(cookie) = cookie { + if let Err(err) = self.nebraska_cookie.save(&cookie) { + warn!("failed to persist Nebraska cookie: {err}"); + } + } + } + /// Sends a single Nebraska event report, logging (but never /// propagating) failure. /// @@ -1628,9 +1663,13 @@ impl Orchestrator { } }; let label = report.label(); + let initial_cookie = self.load_nebraska_cookie(); let result = task::spawn_blocking(move || { - let client = NebraskaClient::new(endpoint, app_id, track, machine_id); - match report { + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + if let Some(cookie) = initial_cookie { + client = client.with_cookie(cookie); + } + let outcome = match report { NebraskaReport::Progress { version, event } => { client.report_progress(&version, event) } @@ -1640,12 +1679,18 @@ impl Orchestrator { NebraskaReport::Failed { previous, current } => { client.report_failure(&previous, ¤t) } - } + }; + (outcome, client.cookie()) }) .await; match result { - Ok(Ok(())) => debug!("reported Nebraska '{label}' event"), - Ok(Err(err)) => warn!("Nebraska '{label}' event report failed: {err}"), + Ok((outcome, cookie)) => { + self.save_nebraska_cookie(cookie); + match outcome { + Ok(()) => debug!("reported Nebraska '{label}' event"), + Err(err) => warn!("Nebraska '{label}' event report failed: {err}"), + } + } Err(err) => warn!("Nebraska '{label}' event report task panicked: {err}"), } } diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index df03ebf728..8cefc49617 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -61,6 +61,14 @@ const DEFAULT_NODE_GONE_MAX_WAIT: Option = None; /// File name for the persisted agent state (see `annotations::state`). pub const STATE_FILE_NAME: &str = "state.json"; pub const DEFAULT_STATE_PATH: &str = formatcp!("/var/lib/trident-acl-agent/{STATE_FILE_NAME}"); +/// File name for the persisted Nebraska affinity cookie (see +/// `core::nebraska::CookieStore`). Deliberately a separate file from +/// `STATE_FILE_NAME`: the cookie is a host/connection-level attribute with +/// no relationship to any update operation, unlike everything else in +/// `state.json`. +pub const NEBRASKA_COOKIE_FILE_NAME: &str = "nebraska-cookie"; +pub const DEFAULT_NEBRASKA_COOKIE_PATH: &str = + formatcp!("/var/lib/trident-acl-agent/{NEBRASKA_COOKIE_FILE_NAME}"); pub const DEFAULT_KUBELET_KUBECONFIG: &str = "/var/lib/kubelet/kubeconfig"; /// Default annotation-key prefix. /// Override with `TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX`. @@ -135,6 +143,9 @@ impl AgentConfig { state_path: orchestration .state_path .unwrap_or_else(|| PathBuf::from(DEFAULT_STATE_PATH)), + cookie_path: orchestration + .cookie_path + .unwrap_or_else(|| PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH)), stage_timeout: orchestration.stage_timeout.unwrap_or(DEFAULT_STAGE_TIMEOUT), finalize_timeout: orchestration .finalize_timeout @@ -213,6 +224,8 @@ struct RawOrchestrationConfig { mode: Option, #[serde(deserialize_with = "empty_path_as_none")] state_path: Option, + #[serde(deserialize_with = "empty_path_as_none")] + cookie_path: Option, #[serde(deserialize_with = "empty_duration_as_none")] stage_timeout: Option, #[serde(deserialize_with = "empty_duration_as_none")] @@ -397,6 +410,11 @@ impl FromStr for Mode { pub struct OrchestrationConfig { pub mode: Mode, pub state_path: PathBuf, + /// Path to the persisted Nebraska affinity cookie (see + /// `core::nebraska::CookieStore`). Deliberately a separate file/field + /// from `state_path`: the cookie has no relationship to any particular + /// update operation, unlike everything `state_path` persists. + pub cookie_path: PathBuf, /// Placeholder default pending real data from storm aclagent scenario runs. pub stage_timeout: Duration, /// Placeholder default pending real data from storm aclagent scenario runs. @@ -424,6 +442,7 @@ impl Default for OrchestrationConfig { Self { mode: Mode::Annotations, state_path: PathBuf::from(DEFAULT_STATE_PATH), + cookie_path: PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH), stage_timeout: DEFAULT_STAGE_TIMEOUT, finalize_timeout: DEFAULT_FINALIZE_TIMEOUT, heartbeat_interval: DEFAULT_HEARTBEAT_INTERVAL, @@ -486,6 +505,10 @@ mod tests { config.orchestration.state_path, PathBuf::from(DEFAULT_STATE_PATH) ); + assert_eq!( + config.orchestration.cookie_path, + PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH) + ); assert_eq!(config.orchestration.stage_timeout, DEFAULT_STAGE_TIMEOUT); assert_eq!( config.orchestration.finalize_timeout, @@ -533,6 +556,10 @@ mod tests { "TRIDENT_ACL_AGENT_ORCHESTRATION_STATE_PATH", "/var/lib/trident-acl-agent/custom-state.json", ), + ( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH", + "/var/lib/trident-acl-agent/custom-nebraska-cookie", + ), ("TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT", "21m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT", "11m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_HEARTBEAT_INTERVAL", "45s"), @@ -566,6 +593,10 @@ mod tests { config.orchestration.state_path, PathBuf::from("/var/lib/trident-acl-agent/custom-state.json") ); + assert_eq!( + config.orchestration.cookie_path, + PathBuf::from("/var/lib/trident-acl-agent/custom-nebraska-cookie") + ); assert_eq!( config.orchestration.stage_timeout, Duration::from_secs(21 * 60) diff --git a/crates/trident-acl-agent/src/core/nebraska/client.rs b/crates/trident-acl-agent/src/core/nebraska/client.rs index e11435e21e..422fbd05f7 100644 --- a/crates/trident-acl-agent/src/core/nebraska/client.rs +++ b/crates/trident-acl-agent/src/core/nebraska/client.rs @@ -1,6 +1,6 @@ //! The high-level [`Client`] for talking to a Nebraska server. -use std::{collections::BTreeMap, thread, time::Duration}; +use std::{cell::RefCell, collections::BTreeMap, thread, time::Duration}; use log::{debug, trace, warn}; use semver::Version; @@ -309,6 +309,16 @@ pub struct Client { track: String, machine_id: MachineId, client_version: Option, + /// Nebraska's load-balancer affinity cookie (e.g. an App-Service-style + /// `ASLBSA` cookie), carried on every request once known and updated + /// from every response that sets one. `RefCell` because every method + /// here takes `&self` (see the type docs), but a single `Client` still + /// needs to pick up a cookie mid-sequence (e.g. during + /// [`check_for_update`](Client::check_for_update)) and carry it into its + /// own later calls. Seeding it across *separate* `Client`s - including + /// across a reboot - is the caller's job: see [`with_cookie`](Client::with_cookie) + /// and [`cookie`](Client::cookie). + cookie: RefCell>, transport: T, } @@ -342,10 +352,30 @@ impl Client { track: track.into(), machine_id, client_version: None, + cookie: RefCell::new(None), transport, } } + /// Seeds this client with a previously-persisted Nebraska affinity + /// cookie (see the [`cookie`](Client::cookie) getter), so it is sent on + /// this client's very first request rather than only picked up from a + /// response. + pub fn with_cookie(self, cookie: impl Into) -> Self { + *self.cookie.borrow_mut() = Some(cookie.into()); + self + } + + /// The most recent Nebraska affinity cookie this client has seen, if + /// any - either seeded via [`with_cookie`](Client::with_cookie) or + /// picked up from a response. Callers that need the cookie to outlive + /// this `Client` (across a reboot, or just a later process invocation) + /// must read this after use and persist it themselves; nothing here + /// does that automatically. + pub fn cookie(&self) -> Option { + self.cookie.borrow().clone() + } + /// Names the **updater** — the program driving the update — in /// ``. /// @@ -599,7 +629,15 @@ impl Client { redacted(&self.endpoint), String::from_utf8_lossy(&body) ); - let text = self.transport.post_xml(&self.endpoint, &body)?; + // Deliberately never logged: this is a session-affinity credential, + // not protocol content. + let cookie = self.cookie.borrow().clone(); + let (text, set_cookie) = + self.transport + .post_xml(&self.endpoint, &body, cookie.as_deref())?; + if let Some(set_cookie) = set_cookie { + *self.cookie.borrow_mut() = Some(set_cookie); + } trace!("Nebraska response: {} bytes", text.len()); let response = wire::parse_response(&text).map_err(NebraskaError::Parse)?; trace!("Nebraska response: {}", summarize(&response)); @@ -820,11 +858,14 @@ mod tests { /// Stands in for the updater version a real caller would supply. const TEST_CLIENT_VERSION: &str = "test-updater-1.2.3"; - /// A canned transport: records the last request body and returns a fixed - /// response, so client logic is exercised without a network. + /// A canned transport: records the last request body/cookie and returns + /// a fixed response (optionally with a canned `Set-Cookie`), so client + /// logic is exercised without a network. struct MockTransport { response: String, last_body: RefCell>, + last_cookie_sent: RefCell>, + set_cookie: RefCell>, } impl MockTransport { @@ -832,14 +873,29 @@ mod tests { Self { response: response.into(), last_body: RefCell::new(None), + last_cookie_sent: RefCell::new(None), + set_cookie: RefCell::new(None), } } + + /// Makes every subsequent `post_xml` call return `cookie` as the + /// captured `Set-Cookie` value, as if the server had just issued one. + fn with_set_cookie(self, cookie: impl Into) -> Self { + *self.set_cookie.borrow_mut() = Some(cookie.into()); + self + } } impl Transport for MockTransport { - fn post_xml(&self, _endpoint: &Url, body: &[u8]) -> Result { + fn post_xml( + &self, + _endpoint: &Url, + body: &[u8], + cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError> { *self.last_body.borrow_mut() = Some(String::from_utf8_lossy(body).into_owned()); - Ok(self.response.clone()) + *self.last_cookie_sent.borrow_mut() = cookie.map(str::to_string); + Ok((self.response.clone(), self.set_cookie.borrow().clone())) } } @@ -862,10 +918,15 @@ mod tests { } impl Transport for ScriptedTransport { - fn post_xml(&self, _endpoint: &Url, _body: &[u8]) -> Result { + fn post_xml( + &self, + _endpoint: &Url, + _body: &[u8], + _cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError> { let index = (self.calls.get() as usize).min(self.responses.len() - 1); self.calls.set(self.calls.get() + 1); - Ok(self.responses[index].clone()) + Ok((self.responses[index].clone(), None)) } } @@ -1489,6 +1550,53 @@ mod tests { ); } + #[test] + fn no_cookie_is_sent_until_one_is_known() { + let client = client_with(NO_UPDATE); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert_eq!(client.cookie(), None); + } + + #[test] + fn with_cookie_seeds_the_first_request() { + let client = client_with(NO_UPDATE).with_cookie("ASLBSA=seeded"); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=seeded") + ); + } + + #[test] + fn cookie_from_a_response_is_echoed_on_the_next_request() { + // Two sequential calls on the *same* Client, mirroring what happens + // mid-operation (e.g. check_for_update picking up a cookie that a + // later report_progress on the same Client must then carry). + let client = Client::with_transport( + Url::parse("https://nebraska.example/v1/update/").unwrap(), + "app-1", + "stable", + MachineId::new("mid-1").unwrap(), + MockTransport::new(NO_UPDATE).with_set_cookie("ASLBSA=fresh"), + ); + + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().clone(), + None, + "first request has no cookie to send yet" + ); + assert_eq!(client.cookie().as_deref(), Some("ASLBSA=fresh")); + + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=fresh"), + "second request must echo back the cookie the first response set" + ); + } + #[test] fn redacted_endpoint_drops_path_and_query() { let url = Url::parse("https://nebraska.example:8443/v1/update/?secret=hunter2").unwrap(); @@ -1585,7 +1693,12 @@ mod tests { calls: Cell, } impl Transport for AlwaysFails { - fn post_xml(&self, _endpoint: &Url, _body: &[u8]) -> Result { + fn post_xml( + &self, + _endpoint: &Url, + _body: &[u8], + _cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError> { self.calls.set(self.calls.get() + 1); Err(NebraskaError::Transport("down".into())) } diff --git a/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs new file mode 100644 index 0000000000..ed66963ec3 --- /dev/null +++ b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs @@ -0,0 +1,193 @@ +//! Persists Nebraska's load-balancer affinity cookie +//! (`/var/lib/trident-acl-agent/nebraska-cookie` by default) across process +//! invocations and reboots. +//! +//! This is deliberately **not** part of `annotations::state`'s +//! `state.json`: that file exists to bridge one update *operation's* +//! pre-/post-reboot halves (every field there is keyed by `operation_id` or +//! is the single in-flight `pendingCommit`), and is read with +//! `deny_unknown_fields`. Nebraska's affinity cookie is a host/connection +//! -level attribute with no relationship to any particular operation - it +//! must outlive operations completing, failing, or being reset, and +//! vice versa - so it gets its own small file instead. + +use std::{ + fs, + io::{ErrorKind, Write}, + os::unix::fs::OpenOptionsExt, + path::{Path, PathBuf}, + process, +}; + +use anyhow::{Context, Error}; + +/// The cookie is a bearer-style session-affinity credential, so store it +/// with owner-only permissions - the same posture `annotations::state` +/// already applies to `state.json`. +const COOKIE_FILE_MODE: u32 = 0o600; + +/// Durable store for the single Nebraska affinity cookie this host has most +/// recently seen (e.g. an App-Service-style `ASLBSA`/`ASLBSACORS` pair, +/// already reduced to a `Cookie:`-header-ready string by +/// [`transport::ReqwestTransport`](super::transport::ReqwestTransport)). +#[derive(Debug, Clone)] +pub struct CookieStore { + path: PathBuf, +} + +impl CookieStore { + pub fn new(path: impl Into) -> Self { + Self { path: path.into() } + } + + pub fn path(&self) -> &Path { + &self.path + } + + /// Loads the last-persisted cookie, or `None` if this host has never + /// seen one yet (including a brand-new install with no file at all). + pub fn load(&self) -> Result, Error> { + match fs::read_to_string(&self.path) { + Ok(raw) => { + let trimmed = raw.trim(); + Ok(if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + }) + } + Err(err) if err.kind() == ErrorKind::NotFound => Ok(None), + Err(err) => Err(err).with_context(|| format!("failed to read {}", self.path.display())), + } + } + + /// Atomically persists `cookie`, replacing any previous value. + /// + /// Mirrors `annotations::state::StateStore::save`'s crash safety: write + /// to a sibling temp file at owner-only (0600) permissions, fsync it, + /// then atomically rename over the real path. + pub fn save(&self, cookie: &str) -> Result<(), Error> { + let parent = match self.path.parent() { + Some(parent) => { + fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + parent + } + None => Path::new("."), + }; + + let temp_path = parent.join(format!( + "{}.tmp-{}", + self.path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("nebraska-cookie"), + process::id() + )); + + { + // create_new (not create+truncate) so the file is always newly + // created with COOKIE_FILE_MODE - see StateStore::save for why + // truncating a stale temp file left behind by a prior crash + // would silently keep that file's existing permissions. A + // leftover temp file (same process::id() reused across + // reboots) can only be this process's own abandoned write, + // never another process's live file, so it is removed and + // retried once. + let mut open_opts = fs::OpenOptions::new(); + open_opts + .write(true) + .create_new(true) + .mode(COOKIE_FILE_MODE); + let mut file = match open_opts.open(&temp_path) { + Ok(file) => file, + Err(err) if err.kind() == ErrorKind::AlreadyExists => { + fs::remove_file(&temp_path).with_context(|| { + format!("failed to remove stale {}", temp_path.display()) + })?; + open_opts + .open(&temp_path) + .with_context(|| format!("failed to create {}", temp_path.display()))? + } + Err(err) => { + return Err(err) + .with_context(|| format!("failed to create {}", temp_path.display())) + } + }; + file.write_all(cookie.as_bytes()) + .with_context(|| format!("failed to write {}", temp_path.display()))?; + file.sync_all() + .with_context(|| format!("failed to fsync {}", temp_path.display()))?; + } + + fs::rename(&temp_path, &self.path).with_context(|| { + format!( + "failed to atomically replace {} with {}", + self.path.display(), + temp_path.display() + ) + })?; + + // Best-effort: POSIX doesn't guarantee a rename is durable until the + // containing directory's metadata is also synced, so fsync it too. + // Not fatal if this fails (e.g. unsupported on some filesystems) - + // the rename itself has already succeeded. + if let Ok(dir) = fs::File::open(parent) { + let _ = dir.sync_all(); + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn store() -> (tempfile::TempDir, CookieStore) { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let path = dir.path().join("nebraska-cookie"); + let store = CookieStore::new(path); + (dir, store) + } + + #[test] + fn load_returns_none_when_file_is_missing() { + let (_dir, store) = store(); + assert_eq!(store.load().unwrap(), None); + } + + #[test] + fn save_then_load_round_trips() { + let (_dir, store) = store(); + store.save("ASLBSA=abc123").unwrap(); + assert_eq!(store.load().unwrap(), Some("ASLBSA=abc123".to_string())); + } + + #[test] + fn save_overwrites_a_previous_value() { + let (_dir, store) = store(); + store.save("ASLBSA=old").unwrap(); + store.save("ASLBSA=new").unwrap(); + assert_eq!(store.load().unwrap(), Some("ASLBSA=new".to_string())); + } + + #[test] + fn saved_file_is_owner_only() { + use std::os::unix::fs::PermissionsExt; + + let (_dir, store) = store(); + store.save("ASLBSA=abc123").unwrap(); + let mode = fs::metadata(store.path()).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, COOKIE_FILE_MODE); + } + + #[test] + fn save_creates_missing_parent_directories() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let path = dir.path().join("nested").join("nebraska-cookie"); + let store = CookieStore::new(path); + store.save("ASLBSA=abc123").unwrap(); + assert_eq!(store.load().unwrap(), Some("ASLBSA=abc123".to_string())); + } +} diff --git a/crates/trident-acl-agent/src/core/nebraska/mod.rs b/crates/trident-acl-agent/src/core/nebraska/mod.rs index 484efc3328..68b43a6efe 100644 --- a/crates/trident-acl-agent/src/core/nebraska/mod.rs +++ b/crates/trident-acl-agent/src/core/nebraska/mod.rs @@ -76,6 +76,7 @@ //! ``` mod client; +mod cookie_store; mod error; mod event; mod id; @@ -84,6 +85,7 @@ mod transport; mod wire; pub use client::{redacted, CheckOutcome, Client, PackageFile, PackageHash, UpdateOffer}; +pub use cookie_store::CookieStore; pub use error::NebraskaError; pub use event::ProgressEvent; pub use id::MachineId; diff --git a/crates/trident-acl-agent/src/core/nebraska/transport.rs b/crates/trident-acl-agent/src/core/nebraska/transport.rs index d0b5b985a1..7c0e4c3167 100644 --- a/crates/trident-acl-agent/src/core/nebraska/transport.rs +++ b/crates/trident-acl-agent/src/core/nebraska/transport.rs @@ -40,8 +40,22 @@ const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); /// return the response text, mapping failures to [`NebraskaError::Transport`] /// (connection failures) or [`NebraskaError::Http`] (non-success status). pub trait Transport { - /// POSTs `body` to `endpoint` and returns the response body as a string. - fn post_xml(&self, endpoint: &Url, body: &[u8]) -> Result; + /// POSTs `body` to `endpoint`, optionally carrying `cookie` as a + /// `Cookie:` header, and returns the response body as a string together + /// with any `Cookie:`-header-ready value derived from the response's + /// `Set-Cookie` header(s) (`None` if the response set no cookie at all). + /// + /// This exists for Nebraska's load-balancer affinity cookie (e.g. an + /// App-Service-style `ASLBSA` cookie): Nebraska issues one on any request + /// that doesn't already carry it, and expects it echoed back on every + /// later request. See [`Client`](crate::nebraska::Client), which owns + /// the cookie across calls. + fn post_xml( + &self, + endpoint: &Url, + body: &[u8], + cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError>; } /// The default [`Transport`], backed by a blocking `reqwest` client. @@ -80,17 +94,36 @@ impl ReqwestTransport { } impl Transport for ReqwestTransport { - fn post_xml(&self, endpoint: &Url, body: &[u8]) -> Result { + fn post_xml( + &self, + endpoint: &Url, + body: &[u8], + cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError> { + let mut request = self + .client + .post(endpoint.as_str()) + .header("Content-Type", "application/xml") + .timeout(self.timeout); + if let Some(cookie) = cookie { + request = request.header(reqwest::header::COOKIE, cookie); + } + // `reqwest::Error`'s `Display` embeds the request URL, which may carry an // Omaha secret; `without_url` strips it so the secret cannot reach a log // or a propagated error. - self.client - .post(endpoint.as_str()) - .header("Content-Type", "application/xml") - .timeout(self.timeout) + let response = request .body(body.to_vec()) .send() - .map_err(|e| NebraskaError::Transport(e.without_url().to_string()))? + .map_err(|e| NebraskaError::Transport(e.without_url().to_string()))?; + + // Captured before the status check: a cookie set alongside a + // non-2xx response is not propagated today (the error path below + // returns no cookie), but the success path - by far the common + // case - must not miss one. + let set_cookie = extract_cookie(&response); + + let text = response .error_for_status() .map_err(|e| NebraskaError::Http { status: e.status().map(|s| s.as_u16()), @@ -100,7 +133,36 @@ impl Transport for ReqwestTransport { .map_err(|e| NebraskaError::Http { status: None, message: e.without_url().to_string(), - }) + })?; + + Ok((text, set_cookie)) + } +} + +/// Builds a `Cookie:`-header-ready value from a response's `Set-Cookie` +/// header(s), or `None` if it carried none. +/// +/// Nebraska's affinity cookie can come as more than one `Set-Cookie` header +/// (e.g. an App-Service-style `ASLBSA`/`ASLBSACORS` pair); only the +/// `name=value` portion before each header's first `;` belongs in an +/// outgoing `Cookie:` header, so attributes like `Path`/`Secure`/`Max-Age` +/// are deliberately dropped here. Multiple cookies are joined with `; `. +fn extract_cookie(response: &reqwest::blocking::Response) -> Option { + let pairs: Vec = response + .headers() + .get_all(reqwest::header::SET_COOKIE) + .iter() + .filter_map(|value| value.to_str().ok()) + .filter_map(|raw| raw.split(';').next()) + .map(str::trim) + .filter(|pair| !pair.is_empty()) + .map(str::to_string) + .collect(); + + if pairs.is_empty() { + None + } else { + Some(pairs.join("; ")) } } @@ -129,7 +191,7 @@ mod tests { let started = Instant::now(); let err = ReqwestTransport::with_timeout(TEST_TIMEOUT) - .post_xml(&endpoint, b"") + .post_xml(&endpoint, b"", None) .unwrap_err(); assert!(matches!(err, NebraskaError::Transport(_)), "got {err:?}"); @@ -143,4 +205,105 @@ mod tests { // into the error text. assert!(!err.to_string().contains("/v1/update"), "{err}"); } + + #[test] + fn post_xml_sends_cookie_header_when_provided() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/") + .match_header("cookie", "ASLBSA=abc123") + .with_status(200) + .with_body("") + .expect(1) + .create(); + + let (text, set_cookie) = ReqwestTransport::new() + .post_xml( + &Url::parse(&server.url()).unwrap(), + b"", + Some("ASLBSA=abc123"), + ) + .unwrap(); + + mock.assert(); + assert_eq!(text, ""); + assert_eq!(set_cookie, None); + } + + #[test] + fn post_xml_omits_cookie_header_when_none_stored_yet() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_body("") + .expect(1) + .create(); + + ReqwestTransport::new() + .post_xml(&Url::parse(&server.url()).unwrap(), b"", None) + .unwrap(); + + mock.assert(); + } + + #[test] + fn post_xml_captures_set_cookie_from_response() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/") + .with_status(200) + // Real affinity cookies carry attributes (Path/Secure/...); only + // the name=value part should survive into the returned cookie. + .with_header("set-cookie", "ASLBSA=abc123; path=/; secure; httponly") + .with_body("") + .create(); + + let (_, set_cookie) = ReqwestTransport::new() + .post_xml(&Url::parse(&server.url()).unwrap(), b"", None) + .unwrap(); + + mock.assert(); + assert_eq!(set_cookie.as_deref(), Some("ASLBSA=abc123")); + } + + #[test] + fn post_xml_joins_multiple_set_cookie_headers() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/") + .with_status(200) + .with_header("set-cookie", "ASLBSA=abc123; path=/") + .with_header("set-cookie", "ASLBSACORS=def456; path=/") + .with_body("") + .create(); + + let (_, set_cookie) = ReqwestTransport::new() + .post_xml(&Url::parse(&server.url()).unwrap(), b"", None) + .unwrap(); + + mock.assert(); + assert_eq!( + set_cookie.as_deref(), + Some("ASLBSA=abc123; ASLBSACORS=def456") + ); + } + + #[test] + fn post_xml_returns_none_when_response_sets_no_cookie() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/") + .with_status(200) + .with_body("") + .create(); + + let (_, set_cookie) = ReqwestTransport::new() + .post_xml(&Url::parse(&server.url()).unwrap(), b"", None) + .unwrap(); + + mock.assert(); + assert_eq!(set_cookie, None); + } } From 79bd465d28f3d381cfbdd769486b3f0fc57f1fdd Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Wed, 14 Oct 2026 20:18:55 +0000 Subject: [PATCH 02/13] storm/aclagent: simulate and assert Nebraska's affinity cookie Extends the earlier trident-acl-agent cookie-persistence change with storm coverage. Real Nebraska's own Omaha handler never sets the Azure-App-Service-style affinity cookie production is fronted by - it's an infrastructure-level artifact, not something Nebraska's application code emits - so storm's fake Nebraska (which wraps the real upstream Omaha handler) never exercised this path before. - NebraskaProxy.Handler() now simulates that cookie: issues a stable, per-instance value on any request lacking one, and classifies every request as missing/correct/wrong via new RequestsMissingCookie(), RequestsWithCorrectCookie(), RequestsWithWrongCookie(), and AffinityCookie() accessors. Purely additive to the test harness; the real Omaha protocol handling is untouched. - New nebraska_test.go subtest drives the simulated cookie directly over HTTP (no VM needed): asserts it's issued once, not re-issued once echoed back, and that a stale cookie is classified as wrong. - RunABUpdate (tools/storm/aclagent/tests/update.go) now asserts, after its existing stage/finalize/reboot/commit run: exactly one request arrived with no cookie (the agent's first-ever contact), none arrived with a wrong cookie, at least one echoed the correct cookie, and - the strongest check - the cookie persisted on disk (/var/lib/trident-acl-agent/nebraska-cookie, the same persistent partition the Nebraska CA cert is pinned to) matches exactly across the real "systemctl reboot" finalize triggers, proving cross-process, cross-reboot persistence with the real compiled binary under systemd. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6c26d0a1-b43f-4ccb-a684-0d92c41f1321 --- tools/storm/aclagent/proxies/nebraska.go | 93 +++++++++++++++++++ tools/storm/aclagent/proxies/nebraska_test.go | 88 ++++++++++++++++++ tools/storm/aclagent/tests/update.go | 34 +++++++ 3 files changed, 215 insertions(+) diff --git a/tools/storm/aclagent/proxies/nebraska.go b/tools/storm/aclagent/proxies/nebraska.go index ced67b5257..085ab32f2f 100644 --- a/tools/storm/aclagent/proxies/nebraska.go +++ b/tools/storm/aclagent/proxies/nebraska.go @@ -95,6 +95,28 @@ type NebraskaProxy struct { api *api.API handler *omaha.Handler appID string + + // affinityCookieValue simulates the Azure App Service-style + // load-balancer affinity cookie (e.g. "ASLBSA") that fronts the real + // production Nebraska deployment. Real Nebraska's own Omaha handler + // never sets this itself - it's purely an infrastructure-level artifact + // - so this proxy injects it the same way App Service does (issued on + // any request that doesn't already carry it, stable thereafter), purely + // so storm can exercise trident-acl-agent's cookie persistence/echo + // behavior end-to-end. See + // crates/trident-acl-agent/src/core/nebraska/cookie_store.rs. Set once + // in ListenAndServe before the server starts accepting requests, never + // mutated afterward, so it's safe to read from Handler() without a lock. + affinityCookieValue string + + // cookieMu protects the request-classification counters below, which + // Handler() updates on every request (concurrently, in principle) and + // storm tests read after a scenario completes via + // RequestsMissingCookie/RequestsWithCorrectCookie/RequestsWithWrongCookie. + cookieMu sync.Mutex + requestsMissingCookie int + requestsWithCorrectCookie int + requestsWithWrongCookie int } // ExpectedUpdateStatusSequence is the ordered sequence of Nebraska instance @@ -136,6 +158,17 @@ func (p *NebraskaProxy) Handler() http.Handler { return } defer r.Body.Close() + + // Simulate the Azure App Service-style load-balancer affinity + // cookie described in affinityCookieValue's doc comment: issue it + // (with the attributes a real one would carry) on any request that + // doesn't already have it, and classify every request by whether it + // carried the right cookie, no cookie, or a stale/wrong one. + p.recordAffinityCookie(r.Header.Get("Cookie")) + if r.Header.Get("Cookie") == "" { + w.Header().Set("Set-Cookie", p.affinityCookieValue+"; Path=/; HttpOnly") + } + ip := r.RemoteAddr if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { ip = host @@ -150,6 +183,61 @@ func (p *NebraskaProxy) Handler() http.Handler { }) } +// recordAffinityCookie classifies one incoming request's raw Cookie header +// against the simulated affinity cookie this proxy issues (see +// affinityCookieValue), updating the counters RequestsMissingCookie, +// RequestsWithCorrectCookie, and RequestsWithWrongCookie read from. +func (p *NebraskaProxy) recordAffinityCookie(cookieHeader string) { + p.cookieMu.Lock() + defer p.cookieMu.Unlock() + switch { + case cookieHeader == "": + p.requestsMissingCookie++ + case cookieHeader == p.affinityCookieValue: + p.requestsWithCorrectCookie++ + default: + p.requestsWithWrongCookie++ + } +} + +// AffinityCookie returns the exact `name=value` this proxy's simulated +// load-balancer affinity cookie carries (see affinityCookieValue), for +// storm tests that want to assert the value persisted on a VM +// (/var/lib/trident-acl-agent/nebraska-cookie) matches what this proxy +// actually issued. +func (p *NebraskaProxy) AffinityCookie() string { + return p.affinityCookieValue +} + +// RequestsMissingCookie returns how many requests this proxy has served +// with no Cookie header at all - i.e. before trident-acl-agent has ever +// picked up the simulated affinity cookie. A healthy end-to-end run expects +// exactly 1: the agent's very first-ever contact with this Nebraska +// instance. +func (p *NebraskaProxy) RequestsMissingCookie() int { + p.cookieMu.Lock() + defer p.cookieMu.Unlock() + return p.requestsMissingCookie +} + +// RequestsWithCorrectCookie returns how many requests arrived already +// carrying exactly the affinity cookie this proxy issued. +func (p *NebraskaProxy) RequestsWithCorrectCookie() int { + p.cookieMu.Lock() + defer p.cookieMu.Unlock() + return p.requestsWithCorrectCookie +} + +// RequestsWithWrongCookie returns how many requests arrived carrying a +// Cookie header that does not match the affinity cookie this proxy issued. +// Always 0 for a correct trident-acl-agent; nonzero means the agent sent a +// stale or corrupted cookie value. +func (p *NebraskaProxy) RequestsWithWrongCookie() int { + p.cookieMu.Lock() + defer p.cookieMu.Unlock() + return p.requestsWithWrongCookie +} + // ListenAndServe starts an ephemeral Postgres container, applies Nebraska's // db migrations against it, seeds an application/package/channel/group // matching p.Scenario, and starts serving the real Nebraska Omaha handler on @@ -214,6 +302,11 @@ func (p *NebraskaProxy) ListenAndServe(ctx context.Context, listenAddr string) ( return nil, nil, fmt.Errorf("failed to seed Nebraska scenario: %w", err) } + // Stable for this proxy's whole lifetime (see affinityCookieValue's doc + // comment); must be set before the server below starts accepting + // requests, since Handler() reads it on every request. + p.affinityCookieValue = fmt.Sprintf("ASLBSA=storm%d", time.Now().UnixNano()) + listener, err = net.Listen("tcp", listenAddr) if err != nil { stopEphemeralPostgres(containerID) diff --git a/tools/storm/aclagent/proxies/nebraska_test.go b/tools/storm/aclagent/proxies/nebraska_test.go index 537cca26cd..4da023f97a 100644 --- a/tools/storm/aclagent/proxies/nebraska_test.go +++ b/tools/storm/aclagent/proxies/nebraska_test.go @@ -101,6 +101,41 @@ func postEvent(t *testing.T, client *http.Client, addr, appID, machineID string, return string(body) } +// postUpdateCheckWithCookie is like postUpdateCheck but lets the caller set +// an explicit Cookie header (empty means none), and returns the response +// headers rather than the body - for asserting Set-Cookie/affinity-cookie +// behavior without needing a cookie-jar-enabled client (trident-acl-agent's +// own Client decides per request, not via a jar, whether it has a cookie to +// send yet - see crates/trident-acl-agent/src/core/nebraska/client.rs). +func postUpdateCheckWithCookie(t *testing.T, client *http.Client, addr, appID, machineID, cookie string) http.Header { + t.Helper() + body := fmt.Sprintf(` + + + + +`, appID, machineID) + + req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("https://%s/", addr), strings.NewReader(body)) + if err != nil { + t.Fatalf("building request: %v", err) + } + req.Header.Set("Content-Type", "text/xml") + if cookie != "" { + req.Header.Set("Cookie", cookie) + } + resp, err := client.Do(req) + if err != nil { + t.Fatalf("POST updatecheck: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", resp.StatusCode, respBody) + } + return resp.Header +} + // TestNebraskaProxy exercises NebraskaProxy against the real // github.com/flatcar/nebraska/backend Omaha handler and an ephemeral // Postgres container (started/torn down per subtest), rather than a @@ -199,4 +234,57 @@ func TestNebraskaProxy(t *testing.T) { t.Fatal("expected ValidateStatusHistory to reject a truncated status sequence") } }) + + // This simulated affinity cookie exists purely for storm's own + // assertions (see RunABUpdate in tools/storm/aclagent/tests/update.go) - + // real Nebraska's own Omaha handler, exercised by every other subtest + // above, never sets one itself. + t.Run("affinity-cookie-issued-and-echoed", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + p, client := newTestNebraskaProxy(t, &NebraskaScenario{Available: false}) + listener, stop, err := p.ListenAndServe(ctx, "127.0.0.1:0") + if err != nil { + t.Fatalf("ListenAndServe: %v", err) + } + t.Cleanup(func() { _ = stop() }) + addr := listener.Addr().String() + + // First contact: no Cookie header sent, so the proxy must issue one. + firstHeaders := postUpdateCheckWithCookie(t, client, addr, p.AppID(), "cookie-test-machine", "") + issued := firstHeaders.Get("Set-Cookie") + if issued == "" { + t.Fatal("expected the first cookie-less request to receive a Set-Cookie") + } + nameValue := strings.SplitN(issued, ";", 2)[0] + if nameValue != p.AffinityCookie() { + t.Fatalf("Set-Cookie name=value %q does not match AffinityCookie() %q", nameValue, p.AffinityCookie()) + } + + // Echoing the issued cookie back must not trigger a second Set-Cookie + // (real Nebraska's affinity cookie is only (re)issued when the + // request doesn't already carry one). + secondHeaders := postUpdateCheckWithCookie(t, client, addr, p.AppID(), "cookie-test-machine", nameValue) + if got := secondHeaders.Get("Set-Cookie"); got != "" { + t.Fatalf("expected no Set-Cookie once the cookie is echoed back, got %q", got) + } + + if got := p.RequestsMissingCookie(); got != 1 { + t.Fatalf("RequestsMissingCookie() = %d, want 1", got) + } + if got := p.RequestsWithCorrectCookie(); got != 1 { + t.Fatalf("RequestsWithCorrectCookie() = %d, want 1", got) + } + if got := p.RequestsWithWrongCookie(); got != 0 { + t.Fatalf("RequestsWithWrongCookie() = %d, want 0 before any stale cookie is sent", got) + } + + // A stale/wrong cookie must be classified as such, not silently + // treated as correct. + postUpdateCheckWithCookie(t, client, addr, p.AppID(), "cookie-test-machine", "ASLBSA=not-the-real-one") + if got := p.RequestsWithWrongCookie(); got != 1 { + t.Fatalf("RequestsWithWrongCookie() = %d, want 1 after sending a stale cookie", got) + } + }) } diff --git a/tools/storm/aclagent/tests/update.go b/tools/storm/aclagent/tests/update.go index 8735425fc1..4fd7c7fb89 100644 --- a/tools/storm/aclagent/tests/update.go +++ b/tools/storm/aclagent/tests/update.go @@ -460,6 +460,40 @@ func RunABUpdate(testConfig stormaclconfig.TestConfig, vmConfig stormvmconfig.Al return fmt.Errorf("ACL agent scenario failed Nebraska status validation: %w", err) } + // Nebraska's own Omaha handler doesn't set the cookie - production's + // load-balancer affinity cookie is simulated by this proxy purely to + // exercise trident-acl-agent's own pickup/persist/echo behavior (see + // crates/trident-acl-agent/src/core/nebraska/cookie_store.rs). It's a + // best-effort hint, not something Nebraska requires to drive the update + // correctly, so these checks run after - not instead of - the real + // pass/fail assertions above. + if got := nebraska.RequestsMissingCookie(); got != 1 { + collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) + return fmt.Errorf("expected exactly 1 Nebraska request with no affinity cookie (the agent's first-ever contact), got %d", got) + } + if got := nebraska.RequestsWithWrongCookie(); got != 0 { + collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) + return fmt.Errorf("trident-acl-agent sent %d request(s) carrying a stale/wrong Nebraska affinity cookie", got) + } + if got := nebraska.RequestsWithCorrectCookie(); got == 0 { + collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) + return fmt.Errorf("trident-acl-agent never echoed back the Nebraska affinity cookie after its first contact") + } + // The strongest assertion: the cookie must still be the right value on + // disk after the real "systemctl reboot" above, on the same persistent + // /var/lib/trident-acl-agent partition the Nebraska CA cert is pinned + // to (see prepareVmForAclAgent) - proving cross-reboot persistence, not + // just in-process carry-over within a single Client. + cookieOnDisk, err := stormssh.SshCommandCombinedOutput(vmConfig.VMConfig, vmIP, "sudo cat /var/lib/trident-acl-agent/nebraska-cookie") + if err != nil { + collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) + return fmt.Errorf("failed to read persisted Nebraska cookie from VM: %w", err) + } + if got := strings.TrimSpace(cookieOnDisk); got != nebraska.AffinityCookie() { + collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) + return fmt.Errorf("persisted Nebraska cookie on VM is %q, want %q (must survive the finalize reboot above)", got, nebraska.AffinityCookie()) + } + return collectAclArtifacts(vmConfig.VMConfig, vmIP, testConfig.OutputPath) } From 212d7c1878a103b9d88e2b459848e4fd0fb8afa9 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Wed, 14 Oct 2026 21:20:42 +0000 Subject: [PATCH 03/13] trident-acl-agent: wire the Nebraska connectivity probe into cookie_store Fixes a real CI failure on this branch (ADO build 1225069, [GITHUB]-trident-pr-e2e, storm aclagent run-ab-update): storm-trident::scenario::aclagent::run-ab-update -> FAIL (expected exactly 1 Nebraska request with no affinity cookie (the agent's first-ever contact), got 2) Root cause: RunABUpdate's storm test runs a successful "--validate-connection nebraska" check (via check_nebraska_reachable()) against the real Nebraska endpoint before the stage/finalize scenario even starts, to prove the env-var config override path works. That probe built a bare Client with no cookie_store wiring at all - it never loaded an existing cookie nor persisted the one Nebraska issued - so it always counted as its own permanent "first contact". The real scenario's own first check_for_update then saw no persisted cookie either, making 2 first-contacts instead of 1. Fix: check_nebraska_reachable() now takes a &CookieStore, loads/seeds the probe's Client with any persisted cookie, and persists back whatever the probe's response sets - mirroring the orchestrator's load_nebraska_cookie()/save_nebraska_cookie() pattern. Best-effort, like every other cookie_store use: a load/save failure only warns and the probe proceeds/returns normally either way. connection_check.rs's validate_connection() now builds the same CookieStore (from config.orchestration.cookie_path) the orchestrator uses and passes it through. Updated the 3 existing check_nebraska_reachable tests to pass a temp CookieStore, and added a new test proving the probe persists a cookie on one call and sends it on the next. This also resolves the "should the connectivity probe participate in the cookie contract" open question noted in wiki/projects/trident/trident-acl-agent.md - answer, confirmed by this real pipeline failure: yes, it must share the same store. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6c26d0a1-b43f-4ccb-a684-0d92c41f1321 --- .../trident-acl-agent/src/connection_check.rs | 9 ++ crates/trident-acl-agent/src/lib.rs | 116 +++++++++++++++++- 2 files changed, 121 insertions(+), 4 deletions(-) diff --git a/crates/trident-acl-agent/src/connection_check.rs b/crates/trident-acl-agent/src/connection_check.rs index 62f04f00f5..485016a7e1 100644 --- a/crates/trident-acl-agent/src/connection_check.rs +++ b/crates/trident-acl-agent/src/connection_check.rs @@ -73,14 +73,23 @@ pub async fn validate_connection( // which would make this a "can we get a valid update check" test // rather than the pure reachability check documented on // ConnectionTarget::Nebraska above. + // + // Shares the same persisted CookieStore (cookie_path) the + // orchestrator uses: this probe talks to the same Nebraska + // instance real traffic does, so it must participate in the + // affinity-cookie dance too - otherwise every + // `--validate-connection nebraska` run is its own permanent + // "first contact" that never adopts or persists a cookie. let endpoint_for_task = endpoint.clone(); let track = config.nebraska.track.clone(); + let cookie_store = nebraska::CookieStore::new(config.orchestration.cookie_path.clone()); task::spawn_blocking(move || { check_nebraska_reachable( &endpoint_for_task, &app_id, &track, IdSource::MachineIdHashed, + &cookie_store, ) }) .await diff --git a/crates/trident-acl-agent/src/lib.rs b/crates/trident-acl-agent/src/lib.rs index ab48b31565..0b3edd9517 100644 --- a/crates/trident-acl-agent/src/lib.rs +++ b/crates/trident-acl-agent/src/lib.rs @@ -14,12 +14,13 @@ //! both by this crate's agent binary and by a future Trident ACL Agent that //! orchestrates updates differently. +use log::warn; use semver::Version; use url::Url; use crate::core::{ error::AgentError, - nebraska::{Client, MachineId, NebraskaError}, + nebraska::{Client, CookieStore, MachineId, NebraskaError}, version::FALLBACK_ALWAYS_VERSION, }; @@ -76,18 +77,48 @@ fn build_machine_id(source: IdSource) -> Result { /// `` is *not* a side-effect-free no-op against Nebraska either /// (hence this being named `probe`, not `ping`). See [`Client::probe`]'s doc /// comment for the full explanation. +/// +/// Also loads/saves `cookie_store`: this probe talks to the very same +/// Nebraska instance (and, in production, the same load-balancer-fronted +/// deployment) that `check_for_update`/`report_progress`/etc. do, so it +/// must participate in the shared affinity-cookie dance too -- otherwise +/// every invocation of `--validate-connection nebraska` is itself a fresh +/// "first contact" that never adopts or persists a cookie, silently +/// defeating the whole mechanism's point for one more request each time +/// it's run. Best-effort, like every other cookie_store use: a load/save +/// failure only logs a warning and the probe proceeds/returns normally. pub fn check_nebraska_reachable( url: &Url, app_id: &str, track: &str, machine_id_source: IdSource, + cookie_store: &CookieStore, ) -> Result<(), AgentError> { let machine_id = build_machine_id(machine_id_source)?; - let client = Client::new(url.clone(), app_id, track, machine_id); - match client.probe( + let mut client = Client::new(url.clone(), app_id, track, machine_id); + let initial_cookie = match cookie_store.load() { + Ok(cookie) => cookie, + Err(err) => { + warn!("failed to load Nebraska cookie, proceeding without one: {err}"); + None + } + }; + if let Some(cookie) = initial_cookie { + client = client.with_cookie(cookie); + } + + let result = client.probe( &Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver"), - ) { + ); + + if let Some(cookie) = client.cookie() { + if let Err(err) = cookie_store.save(&cookie) { + warn!("failed to persist Nebraska cookie: {err}"); + } + } + + match result { Ok(()) => Ok(()), // A well-formed response reporting a non-OK app status still proves // the server is reachable and speaking Omaha; only a @@ -105,6 +136,15 @@ mod tests { use mockito::{Matcher, Server}; use url::Url; + /// A `CookieStore` backed by a fresh temp file, for tests that don't + /// care about cookie persistence themselves (most of them) but still + /// need to satisfy `check_nebraska_reachable`'s now-required parameter. + fn temp_cookie_store() -> (tempfile::TempDir, CookieStore) { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let path = dir.path().join("nebraska-cookie"); + (dir, CookieStore::new(path)) + } + #[test] fn test_check_nebraska_reachable_never_registers_or_updates() { // Regression test: check_nebraska_reachable() must use Client::probe(), @@ -122,6 +162,7 @@ mod tests { // this mock at all, so mockito would return no matching mock and the // call would fail instead of silently passing. let mut server = Server::new(); + let (_dir, cookie_store) = temp_cookie_store(); let bare_app_mock = server .mock("POST", "/") @@ -143,6 +184,7 @@ mod tests { "test", "track", IdSource::MachineIdHashed, + &cookie_store, ) .unwrap(); @@ -155,6 +197,7 @@ mod tests { // server and does it speak Omaha" check: a well-formed response with // a non-OK app status should still count as "reachable" here. let mut server = Server::new(); + let (_dir, cookie_store) = temp_cookie_store(); let omaha_mock = server .mock("POST", "/") @@ -173,6 +216,7 @@ mod tests { "test", "track", IdSource::MachineIdHashed, + &cookie_store, ) .unwrap(); @@ -181,14 +225,78 @@ mod tests { #[test] fn test_check_nebraska_reachable_fails_on_transport_error() { + let (_dir, cookie_store) = temp_cookie_store(); let err = check_nebraska_reachable( // Port 0 never accepts a connection. &Url::parse("http://127.0.0.1:0/").unwrap(), "test", "track", IdSource::MachineIdHashed, + &cookie_store, ) .unwrap_err(); assert!(matches!(err, AgentError::Nebraska(_))); } + + #[test] + fn test_check_nebraska_reachable_persists_and_reuses_cookie() { + // Regression test for the real pipeline failure this fixes: a probe + // that never participates in cookie_store is itself a permanent + // "first contact," silently defeating the whole mechanism for every + // `--validate-connection nebraska` invocation. + let mut server = Server::new(); + let (_dir, cookie_store) = temp_cookie_store(); + + let first_mock = server + .mock("POST", "/") + .match_header("cookie", Matcher::Missing) + .with_status(200) + .with_header("set-cookie", "ASLBSA=probe123; Path=/") + .with_body(indoc! {r#" + + + + + "#}) + .expect(1) + .create(); + + check_nebraska_reachable( + &Url::parse(&server.url()).unwrap(), + "test", + "track", + IdSource::MachineIdHashed, + &cookie_store, + ) + .unwrap(); + first_mock.assert(); + assert_eq!( + cookie_store.load().unwrap().as_deref(), + Some("ASLBSA=probe123"), + "the cookie Nebraska issued on this probe must be persisted" + ); + + let second_mock = server + .mock("POST", "/") + .match_header("cookie", "ASLBSA=probe123") + .with_status(200) + .with_body(indoc! {r#" + + + + + "#}) + .expect(1) + .create(); + + check_nebraska_reachable( + &Url::parse(&server.url()).unwrap(), + "test", + "track", + IdSource::MachineIdHashed, + &cookie_store, + ) + .unwrap(); + second_mock.assert(); + } } From 9c7db5f73611175e5cd399a7761fd762b955ed36 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 16 Oct 2026 16:02:51 +0000 Subject: [PATCH 04/13] trident-acl-agent: add cookie-enabled and cookie-lifetime config knobs Two new envy-driven OrchestrationConfig fields, alongside the existing cookie_path: - TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED (bool, default true): blanket kill switch. When false, no Cookie: header is ever sent and no Set-Cookie is ever adopted - the persisted file (if one exists from before the toggle was flipped) is never read or written either, so toggling it back on later just picks the stale file back up as-is. - TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME (enum, default "forever"): "forever" reuses a persisted cookie across every update until Nebraska replaces it (the originally-shipped behavior); "per-update" forces a fresh cookie at the start of every update by skipping the load on handle_stage's initial check_for_update call - definitionally the first Nebraska call of any update, since stage always precedes finalize/commit. Every other call site (report_nebraska_event's progress/terminal reports, the --validate-connection nebraska probe) is always a continuation and is unaffected by this choice. Implementation: - core/config.rs: new CookieLifetime enum (FromStr-parsed, mirroring Mode), cookie_enabled/cookie_lifetime fields on OrchestrationConfig, new empty_bool_as_none/empty_cookie_lifetime_as_none deserializers, defaults/override/malformed-value tests. - annotations/orchestrator.rs: load_nebraska_cookie gains a force_fresh: bool parameter, short-circuiting (without touching disk) when cookies are disabled or force_fresh is set; save_nebraska_cookie short-circuits when disabled. handle_stage's check_for_update is the only call site that computes force_fresh = cookie_lifetime == PerUpdate; report_nebraska_event always passes force_fresh: false. - lib.rs: check_nebraska_reachable gains a cookie_enabled: bool parameter (no lifetime concept - a diagnostic probe isn't part of any update); new test proving a disabled probe never touches an existing persisted cookie. connection_check.rs passes config.orchestration.cookie_enabled through. All 218 trident-acl-agent lib tests pass; clippy and fmt clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6c26d0a1-b43f-4ccb-a684-0d92c41f1321 --- .../src/annotations/orchestrator.rs | 46 +++++-- .../trident-acl-agent/src/connection_check.rs | 2 + crates/trident-acl-agent/src/core/config.rs | 125 ++++++++++++++++++ crates/trident-acl-agent/src/lib.rs | 79 +++++++++-- 4 files changed, 233 insertions(+), 19 deletions(-) diff --git a/crates/trident-acl-agent/src/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index 7ab36f7758..f98d11e3a4 100644 --- a/crates/trident-acl-agent/src/annotations/orchestrator.rs +++ b/crates/trident-acl-agent/src/annotations/orchestrator.rs @@ -33,7 +33,7 @@ use crate::{ SCHEMA_VERSION, }, core::{ - config::AgentConfig, + config::{AgentConfig, CookieLifetime}, cosi_sha384, nebraska::{CheckOutcome, Client as NebraskaClient, CookieStore, ProgressEvent}, trident::{CompletedResponse, TridentClient, TridentClientError}, @@ -706,7 +706,14 @@ impl Orchestrator { Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); - let initial_cookie = self.load_nebraska_cookie(); + // This is `handle_stage`'s own initial Nebraska call - definitionally + // the first one of a brand new update (stage always precedes + // finalize/commit) - so it's the one place `CookieLifetime::PerUpdate` + // forces a fresh cookie rather than inheriting whatever a prior + // update left behind. Every other Nebraska call in this module is a + // continuation and always passes `force_fresh: false`. + let force_fresh = self.config.orchestration.cookie_lifetime == CookieLifetime::PerUpdate; + let initial_cookie = self.load_nebraska_cookie(force_fresh); let (check_result, cookie) = task::spawn_blocking({ let current_version = current_version.clone(); move || { @@ -1580,12 +1587,22 @@ impl Orchestrator { } } - /// Loads the persisted Nebraska affinity cookie, warning (not failing) - /// if the store can't be read. The whole mechanism is a best-effort + /// Loads the persisted Nebraska affinity cookie, honoring + /// `cookie_enabled`/`force_fresh`, and warning (not failing) if the + /// store can't be read. The whole mechanism is a best-effort /// session-affinity hint, never a requirement for talking to Nebraska - - /// a load failure just means this round-trip proceeds cookie-less, the - /// same as this host's very first-ever contact with Nebraska. - fn load_nebraska_cookie(&self) -> Option { + /// a load failure (or a disabled/force-fresh skip) just means this + /// round-trip proceeds cookie-less, the same as this host's very + /// first-ever contact with Nebraska. + /// + /// `force_fresh` is set by callers that are the first Nebraska call of + /// a brand new update (see `CookieLifetime::PerUpdate`'s docs) - when + /// true, this returns `None` without even touching disk, deliberately + /// not inheriting whatever a prior update left behind. + fn load_nebraska_cookie(&self, force_fresh: bool) -> Option { + if !self.config.orchestration.cookie_enabled || force_fresh { + return None; + } match self.nebraska_cookie.load() { Ok(cookie) => cookie, Err(err) => { @@ -1595,8 +1612,14 @@ impl Orchestrator { } } - /// Persists a cookie a Nebraska round-trip ended up holding, if any. + /// Persists a cookie a Nebraska round-trip ended up holding, if any - + /// unless `cookie_enabled` is off, in which case this is a no-op (so a + /// mid-flight toggle from enabled to disabled doesn't keep writing a + /// file a later re-enable would then pick back up unexpectedly). fn save_nebraska_cookie(&self, cookie: Option) { + if !self.config.orchestration.cookie_enabled { + return; + } if let Some(cookie) = cookie { if let Err(err) = self.nebraska_cookie.save(&cookie) { warn!("failed to persist Nebraska cookie: {err}"); @@ -1663,7 +1686,12 @@ impl Orchestrator { } }; let label = report.label(); - let initial_cookie = self.load_nebraska_cookie(); + // Always a continuation of an already-started update (a progress or + // terminal event report), never the first Nebraska call of a new + // one - so `force_fresh` is always false here regardless of + // `cookie_lifetime`. See `handle_stage`'s own call for the one place + // that does force a fresh cookie. + let initial_cookie = self.load_nebraska_cookie(false); let result = task::spawn_blocking(move || { let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); if let Some(cookie) = initial_cookie { diff --git a/crates/trident-acl-agent/src/connection_check.rs b/crates/trident-acl-agent/src/connection_check.rs index 485016a7e1..a0a7ce4b03 100644 --- a/crates/trident-acl-agent/src/connection_check.rs +++ b/crates/trident-acl-agent/src/connection_check.rs @@ -83,6 +83,7 @@ pub async fn validate_connection( let endpoint_for_task = endpoint.clone(); let track = config.nebraska.track.clone(); let cookie_store = nebraska::CookieStore::new(config.orchestration.cookie_path.clone()); + let cookie_enabled = config.orchestration.cookie_enabled; task::spawn_blocking(move || { check_nebraska_reachable( &endpoint_for_task, @@ -90,6 +91,7 @@ pub async fn validate_connection( &track, IdSource::MachineIdHashed, &cookie_store, + cookie_enabled, ) }) .await diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index 8cefc49617..50df127bbe 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -73,6 +73,9 @@ pub const DEFAULT_KUBELET_KUBECONFIG: &str = "/var/lib/kubelet/kubeconfig"; /// Default annotation-key prefix. /// Override with `TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX`. pub const DEFAULT_ANNOTATION_PREFIX: &str = "acl.microsoft.com"; +/// Default for `cookie_enabled`: on, matching the originally-shipped +/// behavior. Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. +const DEFAULT_COOKIE_ENABLED: bool = true; #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct AgentConfig { @@ -146,6 +149,10 @@ impl AgentConfig { cookie_path: orchestration .cookie_path .unwrap_or_else(|| PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH)), + cookie_enabled: orchestration + .cookie_enabled + .unwrap_or(DEFAULT_COOKIE_ENABLED), + cookie_lifetime: orchestration.cookie_lifetime.unwrap_or_default(), stage_timeout: orchestration.stage_timeout.unwrap_or(DEFAULT_STAGE_TIMEOUT), finalize_timeout: orchestration .finalize_timeout @@ -226,6 +233,10 @@ struct RawOrchestrationConfig { state_path: Option, #[serde(deserialize_with = "empty_path_as_none")] cookie_path: Option, + #[serde(deserialize_with = "empty_bool_as_none")] + cookie_enabled: Option, + #[serde(deserialize_with = "empty_cookie_lifetime_as_none")] + cookie_lifetime: Option, #[serde(deserialize_with = "empty_duration_as_none")] stage_timeout: Option, #[serde(deserialize_with = "empty_duration_as_none")] @@ -294,6 +305,33 @@ where .transpose() } +/// Parses `"true"`/`"false"` (Rust's `bool::from_str`, case-sensitive +/// lowercase only), treating the empty string as unset like every other +/// `empty_*_as_none` helper here. +fn empty_bool_as_none<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + empty_as_none(String::deserialize(deserializer)?) + .map(|value| { + value + .parse::() + .map_err(|err| D::Error::custom(format!("invalid boolean {value:?}: {err}"))) + }) + .transpose() +} + +fn empty_cookie_lifetime_as_none<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + empty_as_none(String::deserialize(deserializer)?) + .map(|value| value.parse::().map_err(D::Error::custom)) + .transpose() +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct NebraskaConfig { pub endpoint: Option, @@ -406,6 +444,41 @@ impl FromStr for Mode { } } +/// Controls whether a persisted Nebraska affinity cookie is reused across +/// updates, or deliberately dropped at the start of each new one. See +/// [`OrchestrationConfig::cookie_lifetime`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum CookieLifetime { + /// Reuse a persisted cookie across every update until Nebraska itself + /// replaces it. This is the originally-shipped behavior. + #[default] + Forever, + /// Force a fresh cookie at the start of every update: the very first + /// Nebraska call of a new update (`handle_stage`'s `check_for_update` - + /// definitionally the first call of any update, since stage always + /// precedes finalize/commit) never sends whatever cookie is currently + /// persisted, so Nebraska issues/re-confirms one for this update + /// specifically. Every later call within that same update (progress + /// events, finalize, the post-reboot commit) still reuses whatever that + /// first call established - this is "don't inherit a prior update's + /// session," not "never persist at all." + PerUpdate, +} + +impl FromStr for CookieLifetime { + type Err = Error; + + fn from_str(s: &str) -> Result { + match s { + "forever" => Ok(CookieLifetime::Forever), + "per-update" => Ok(CookieLifetime::PerUpdate), + other => Err(anyhow!( + "unknown cookie lifetime {other:?} (expected \"forever\" or \"per-update\")" + )), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct OrchestrationConfig { pub mode: Mode, @@ -415,6 +488,18 @@ pub struct OrchestrationConfig { /// from `state_path`: the cookie has no relationship to any particular /// update operation, unlike everything `state_path` persists. pub cookie_path: PathBuf, + /// Blanket kill switch for the whole cookie mechanism. When `false`, no + /// `Cookie:` header is ever sent and no `Set-Cookie` is ever adopted - + /// the persisted file (if one exists from before the toggle was + /// flipped) is never read *or* written, so toggling this back to `true` + /// later just picks the stale file back up as-is. Override with + /// `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. + pub cookie_enabled: bool, + /// Whether a persisted cookie is reused forever, or deliberately + /// dropped at the start of every new update. See [`CookieLifetime`]. + /// Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME` + /// (`"forever"` or `"per-update"`). + pub cookie_lifetime: CookieLifetime, /// Placeholder default pending real data from storm aclagent scenario runs. pub stage_timeout: Duration, /// Placeholder default pending real data from storm aclagent scenario runs. @@ -443,6 +528,8 @@ impl Default for OrchestrationConfig { mode: Mode::Annotations, state_path: PathBuf::from(DEFAULT_STATE_PATH), cookie_path: PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH), + cookie_enabled: DEFAULT_COOKIE_ENABLED, + cookie_lifetime: CookieLifetime::Forever, stage_timeout: DEFAULT_STAGE_TIMEOUT, finalize_timeout: DEFAULT_FINALIZE_TIMEOUT, heartbeat_interval: DEFAULT_HEARTBEAT_INTERVAL, @@ -509,6 +596,14 @@ mod tests { config.orchestration.cookie_path, PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH) ); + assert!( + config.orchestration.cookie_enabled, + "cookie_enabled should default to on, matching the originally-shipped behavior" + ); + assert_eq!( + config.orchestration.cookie_lifetime, + CookieLifetime::Forever + ); assert_eq!(config.orchestration.stage_timeout, DEFAULT_STAGE_TIMEOUT); assert_eq!( config.orchestration.finalize_timeout, @@ -560,6 +655,11 @@ mod tests { "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH", "/var/lib/trident-acl-agent/custom-nebraska-cookie", ), + ("TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", "false"), + ( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + "per-update", + ), ("TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT", "21m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT", "11m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_HEARTBEAT_INTERVAL", "45s"), @@ -597,6 +697,11 @@ mod tests { config.orchestration.cookie_path, PathBuf::from("/var/lib/trident-acl-agent/custom-nebraska-cookie") ); + assert!(!config.orchestration.cookie_enabled); + assert_eq!( + config.orchestration.cookie_lifetime, + CookieLifetime::PerUpdate + ); assert_eq!( config.orchestration.stage_timeout, Duration::from_secs(21 * 60) @@ -661,6 +766,26 @@ mod tests { assert!(format!("{err:#}").contains("bogus"), "{err:#}"); } + #[test] + fn malformed_cookie_enabled_is_a_parse_error() { + let err = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", + "yes please", + )])) + .unwrap_err(); + assert!(format!("{err:#}").contains("yes please"), "{err:#}"); + } + + #[test] + fn malformed_cookie_lifetime_is_a_parse_error() { + let err = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + "bogus", + )])) + .unwrap_err(); + assert!(format!("{err:#}").contains("bogus"), "{err:#}"); + } + #[test] fn malformed_duration_is_a_parse_error() { let err = AgentConfig::from_vars(vars(&[( diff --git a/crates/trident-acl-agent/src/lib.rs b/crates/trident-acl-agent/src/lib.rs index 0b3edd9517..b746402a69 100644 --- a/crates/trident-acl-agent/src/lib.rs +++ b/crates/trident-acl-agent/src/lib.rs @@ -87,22 +87,32 @@ fn build_machine_id(source: IdSource) -> Result { /// defeating the whole mechanism's point for one more request each time /// it's run. Best-effort, like every other cookie_store use: a load/save /// failure only logs a warning and the probe proceeds/returns normally. +/// +/// `cookie_enabled` is the same blanket kill switch +/// (`OrchestrationConfig::cookie_enabled`) the orchestrator honors: when +/// `false`, this probe never touches `cookie_store` at all, in either +/// direction. There's no `CookieLifetime` concept here -- a diagnostic +/// probe isn't part of any update, so "first call of a new update" never +/// applies to it. pub fn check_nebraska_reachable( url: &Url, app_id: &str, track: &str, machine_id_source: IdSource, cookie_store: &CookieStore, + cookie_enabled: bool, ) -> Result<(), AgentError> { let machine_id = build_machine_id(machine_id_source)?; let mut client = Client::new(url.clone(), app_id, track, machine_id); - let initial_cookie = match cookie_store.load() { - Ok(cookie) => cookie, - Err(err) => { - warn!("failed to load Nebraska cookie, proceeding without one: {err}"); - None - } - }; + let initial_cookie = cookie_enabled + .then(|| match cookie_store.load() { + Ok(cookie) => cookie, + Err(err) => { + warn!("failed to load Nebraska cookie, proceeding without one: {err}"); + None + } + }) + .flatten(); if let Some(cookie) = initial_cookie { client = client.with_cookie(cookie); } @@ -112,9 +122,11 @@ pub fn check_nebraska_reachable( .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver"), ); - if let Some(cookie) = client.cookie() { - if let Err(err) = cookie_store.save(&cookie) { - warn!("failed to persist Nebraska cookie: {err}"); + if cookie_enabled { + if let Some(cookie) = client.cookie() { + if let Err(err) = cookie_store.save(&cookie) { + warn!("failed to persist Nebraska cookie: {err}"); + } } } @@ -185,6 +197,7 @@ mod tests { "track", IdSource::MachineIdHashed, &cookie_store, + true, ) .unwrap(); @@ -217,6 +230,7 @@ mod tests { "track", IdSource::MachineIdHashed, &cookie_store, + true, ) .unwrap(); @@ -233,6 +247,7 @@ mod tests { "track", IdSource::MachineIdHashed, &cookie_store, + true, ) .unwrap_err(); assert!(matches!(err, AgentError::Nebraska(_))); @@ -267,6 +282,7 @@ mod tests { "track", IdSource::MachineIdHashed, &cookie_store, + true, ) .unwrap(); first_mock.assert(); @@ -295,8 +311,51 @@ mod tests { "track", IdSource::MachineIdHashed, &cookie_store, + true, ) .unwrap(); second_mock.assert(); } + + #[test] + fn test_check_nebraska_reachable_cookie_enabled_false_never_touches_store() { + // The blanket kill switch: with cookie_enabled=false, no Cookie + // header is ever sent and no Set-Cookie is ever adopted/persisted, + // even when the response sets one and even when a cookie is already + // on disk from before the toggle was flipped. + let mut server = Server::new(); + let (_dir, cookie_store) = temp_cookie_store(); + cookie_store.save("ASLBSA=preexisting").unwrap(); + + let mock = server + .mock("POST", "/") + .match_header("cookie", Matcher::Missing) + .with_status(200) + .with_header("set-cookie", "ASLBSA=fresh; Path=/") + .with_body(indoc! {r#" + + + + + "#}) + .expect(1) + .create(); + + check_nebraska_reachable( + &Url::parse(&server.url()).unwrap(), + "test", + "track", + IdSource::MachineIdHashed, + &cookie_store, + false, + ) + .unwrap(); + + mock.assert(); + assert_eq!( + cookie_store.load().unwrap().as_deref(), + Some("ASLBSA=preexisting"), + "a disabled probe must not overwrite a pre-existing cookie with a fresh one" + ); + } } From 6ae4e700e2891164b9d213c5d554625312f6c714 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 16 Oct 2026 16:31:59 +0000 Subject: [PATCH 05/13] acl: honor Nebraska cookie lifetime and disable settings Clear persisted cookies at per-update boundaries and disable cookie adoption and transmission throughout client retries and probes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6400e928-51a9-4a92-aead-de8b03854676 --- .../src/annotations/orchestrator.rs | 188 +++++++++++++++++- .../src/core/nebraska/client.rs | 85 +++++++- crates/trident-acl-agent/src/lib.rs | 3 +- 3 files changed, 262 insertions(+), 14 deletions(-) diff --git a/crates/trident-acl-agent/src/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index f98d11e3a4..5adc45c7de 100644 --- a/crates/trident-acl-agent/src/annotations/orchestrator.rs +++ b/crates/trident-acl-agent/src/annotations/orchestrator.rs @@ -714,10 +714,12 @@ impl Orchestrator { // continuation and always passes `force_fresh: false`. let force_fresh = self.config.orchestration.cookie_lifetime == CookieLifetime::PerUpdate; let initial_cookie = self.load_nebraska_cookie(force_fresh); + let cookie_enabled = self.config.orchestration.cookie_enabled; let (check_result, cookie) = task::spawn_blocking({ let current_version = current_version.clone(); move || { - let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id) + .with_cookie_enabled(cookie_enabled); if let Some(cookie) = initial_cookie { client = client.with_cookie(cookie); } @@ -1591,16 +1593,22 @@ impl Orchestrator { /// `cookie_enabled`/`force_fresh`, and warning (not failing) if the /// store can't be read. The whole mechanism is a best-effort /// session-affinity hint, never a requirement for talking to Nebraska - - /// a load failure (or a disabled/force-fresh skip) just means this + /// a load failure (or a disabled/fresh session) just means this /// round-trip proceeds cookie-less, the same as this host's very /// first-ever contact with Nebraska. /// /// `force_fresh` is set by callers that are the first Nebraska call of /// a brand new update (see `CookieLifetime::PerUpdate`'s docs) - when - /// true, this returns `None` without even touching disk, deliberately - /// not inheriting whatever a prior update left behind. + /// enabled, it invalidates the persisted cookie too, so later reports + /// cannot reload a prior update's cookie if Nebraska issues no new one. fn load_nebraska_cookie(&self, force_fresh: bool) -> Option { - if !self.config.orchestration.cookie_enabled || force_fresh { + if !self.config.orchestration.cookie_enabled { + return None; + } + if force_fresh { + if let Err(err) = self.nebraska_cookie.save("") { + warn!("failed to clear Nebraska cookie for new update: {err}"); + } return None; } match self.nebraska_cookie.load() { @@ -1692,8 +1700,10 @@ impl Orchestrator { // `cookie_lifetime`. See `handle_stage`'s own call for the one place // that does force a fresh cookie. let initial_cookie = self.load_nebraska_cookie(false); + let cookie_enabled = self.config.orchestration.cookie_enabled; let result = task::spawn_blocking(move || { - let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id) + .with_cookie_enabled(cookie_enabled); if let Some(cookie) = initial_cookie { client = client.with_cookie(cookie); } @@ -2337,6 +2347,172 @@ mod tests { core::trident::mock::{connect_mock_client, MockTridentdConfig, Outcome}, }; + async fn cookie_orchestrator( + dir: &std::path::Path, + enabled: bool, + lifetime: CookieLifetime, + ) -> Orchestrator { + use crate::core::config::{KubernetesConfig, OrchestrationConfig}; + + let kubeconfig = dir.join("kubeconfig"); + std::fs::write( + &kubeconfig, + "apiVersion: v1\nkind: Config\ncurrent-context: test\ncontexts:\n- name: test\n context:\n cluster: test\n user: test\nclusters:\n- name: test\n cluster:\n server: http://127.0.0.1:1\nusers:\n- name: test\n user: {}\n", + ) + .unwrap(); + Orchestrator::from_config(AgentConfig { + kubernetes: KubernetesConfig { + kubeconfig, + ..Default::default() + }, + orchestration: OrchestrationConfig { + state_path: dir.join("state.json"), + cookie_path: dir.join("nebraska-cookie"), + cookie_enabled: enabled, + cookie_lifetime: lifetime, + ..Default::default() + }, + ..Default::default() + }) + .await + .unwrap() + } + + #[tokio::test] + async fn per_update_cookie_is_not_reloaded_when_initial_check_sets_none() { + use mockito::Matcher; + + let mut server = mockito::Server::new_async().await; + for (enabled, lifetime) in [ + (true, CookieLifetime::PerUpdate), + (true, CookieLifetime::Forever), + (false, CookieLifetime::PerUpdate), + ] { + let dir = tempfile::tempdir().unwrap(); + let orchestrator = cookie_orchestrator(dir.path(), enabled, lifetime).await; + orchestrator + .nebraska_cookie + .save("ASLBSA=previous") + .unwrap(); + let expected_cookie = if enabled && lifetime == CookieLifetime::Forever { + Matcher::Exact("ASLBSA=previous".to_string()) + } else { + Matcher::Missing + }; + let response = r#""#; + let check = server + .mock("POST", "/v1/update/") + .match_header("cookie", expected_cookie.clone()) + .match_body(Matcher::Regex(""#) + .expect(1) + .create_async() + .await; + let retry = server + .mock("POST", "/v1/update/") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_body(r#""#) + .expect(1) + .create_async() + .await; + let mut request = request(RequestedOperation::Finalize); + request.server = Some(Url::parse(&format!("{}/v1/update/", server.url())).unwrap()); + request.app_id = Some("app-1".to_string()); + request.track = Some("stable".to_string()); + orchestrator + .report_nebraska_event( + &request, + NebraskaReport::Completed { + previous: Version::new(1, 0, 0), + current: Version::new(2, 0, 0), + }, + ) + .await; + first.assert_async().await; + retry.assert_async().await; + assert_eq!( + orchestrator.nebraska_cookie.load().unwrap().as_deref(), + Some("ASLBSA=previous"), + ); + } + fn request(operation: RequestedOperation) -> UpdateRequest { UpdateRequest { schema_version: SCHEMA_VERSION.to_string(), diff --git a/crates/trident-acl-agent/src/core/nebraska/client.rs b/crates/trident-acl-agent/src/core/nebraska/client.rs index 422fbd05f7..67c3c65996 100644 --- a/crates/trident-acl-agent/src/core/nebraska/client.rs +++ b/crates/trident-acl-agent/src/core/nebraska/client.rs @@ -319,6 +319,7 @@ pub struct Client { /// across a reboot - is the caller's job: see [`with_cookie`](Client::with_cookie) /// and [`cookie`](Client::cookie). cookie: RefCell>, + cookie_enabled: bool, transport: T, } @@ -353,16 +354,29 @@ impl Client { machine_id, client_version: None, cookie: RefCell::new(None), + cookie_enabled: true, transport, } } + /// Enables or disables affinity cookies, including response adoption + /// and retries. Disabling also discards any previously seeded cookie. + pub fn with_cookie_enabled(mut self, enabled: bool) -> Self { + self.cookie_enabled = enabled; + if !enabled { + *self.cookie.borrow_mut() = None; + } + self + } + /// Seeds this client with a previously-persisted Nebraska affinity /// cookie (see the [`cookie`](Client::cookie) getter), so it is sent on /// this client's very first request rather than only picked up from a - /// response. + /// response. Ignored when cookies are disabled. pub fn with_cookie(self, cookie: impl Into) -> Self { - *self.cookie.borrow_mut() = Some(cookie.into()); + if self.cookie_enabled { + *self.cookie.borrow_mut() = Some(cookie.into()); + } self } @@ -631,12 +645,18 @@ impl Client { ); // Deliberately never logged: this is a session-affinity credential, // not protocol content. - let cookie = self.cookie.borrow().clone(); + let cookie = if self.cookie_enabled { + self.cookie.borrow().clone() + } else { + None + }; let (text, set_cookie) = self.transport .post_xml(&self.endpoint, &body, cookie.as_deref())?; - if let Some(set_cookie) = set_cookie { - *self.cookie.borrow_mut() = Some(set_cookie); + if self.cookie_enabled { + if let Some(set_cookie) = set_cookie { + *self.cookie.borrow_mut() = Some(set_cookie); + } } trace!("Nebraska response: {} bytes", text.len()); let response = wire::parse_response(&text).map_err(NebraskaError::Parse)?; @@ -915,6 +935,8 @@ mod tests { struct ScriptedTransport { responses: Vec, calls: Cell, + set_cookie: Option, + cookies_sent: RefCell>>, } impl Transport for ScriptedTransport { @@ -922,11 +944,14 @@ mod tests { &self, _endpoint: &Url, _body: &[u8], - _cookie: Option<&str>, + cookie: Option<&str>, ) -> Result<(String, Option), NebraskaError> { let index = (self.calls.get() as usize).min(self.responses.len() - 1); self.calls.set(self.calls.get() + 1); - Ok((self.responses[index].clone(), None)) + self.cookies_sent + .borrow_mut() + .push(cookie.map(str::to_string)); + Ok((self.responses[index].clone(), self.set_cookie.clone())) } } @@ -939,6 +964,8 @@ mod tests { ScriptedTransport { responses: responses.into_iter().map(String::from).collect(), calls: Cell::new(0), + set_cookie: None, + cookies_sent: RefCell::new(Vec::new()), }, ) } @@ -1392,6 +1419,50 @@ mod tests { assert_eq!(client.transport.calls.get(), 2); } + #[test] + fn completion_retries_honor_cookie_enabled() { + for enabled in [true, false] { + let mut client = client_with_responses(vec![IN_PROGRESS, NO_UPDATE]) + .with_cookie("ASLBSA=seeded") + .with_cookie_enabled(enabled); + client.transport.set_cookie = Some("ASLBSA=fresh".to_string()); + + let outcome = client + .complete_after_reboot_with_policy( + fast_policy(), + &Version::new(1, 0, 0), + &Version::new(2, 0, 0), + ) + .unwrap(); + + assert_eq!(outcome, CheckOutcome::UpToDate); + assert_eq!(client.transport.calls.get(), 2); + let expected = if enabled { + vec![ + Some("ASLBSA=seeded".to_string()), + Some("ASLBSA=fresh".to_string()), + ] + } else { + vec![None, None] + }; + assert_eq!(*client.transport.cookies_sent.borrow(), expected); + assert_eq!( + client.cookie().as_deref(), + enabled.then_some("ASLBSA=fresh") + ); + } + } + + #[test] + fn disabled_cookies_cannot_be_seeded() { + let client = client_with(NO_UPDATE) + .with_cookie_enabled(false) + .with_cookie("ASLBSA=seeded"); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert_eq!(client.cookie(), None); + } + #[test] fn complete_after_reboot_fails_when_completion_never_lands() { let client = client_with_responses(vec![IN_PROGRESS]); diff --git a/crates/trident-acl-agent/src/lib.rs b/crates/trident-acl-agent/src/lib.rs index b746402a69..f08cd78960 100644 --- a/crates/trident-acl-agent/src/lib.rs +++ b/crates/trident-acl-agent/src/lib.rs @@ -103,7 +103,8 @@ pub fn check_nebraska_reachable( cookie_enabled: bool, ) -> Result<(), AgentError> { let machine_id = build_machine_id(machine_id_source)?; - let mut client = Client::new(url.clone(), app_id, track, machine_id); + let mut client = + Client::new(url.clone(), app_id, track, machine_id).with_cookie_enabled(cookie_enabled); let initial_cookie = cookie_enabled .then(|| match cookie_store.load() { Ok(cookie) => cookie, From 7035f793fdd7d8407edd96e552feb0df2efe0238 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 16 Oct 2026 17:03:55 +0000 Subject: [PATCH 06/13] acl: define per-update cookies per staging attempt Clarify that restarting an unfinished stage resets affinity even with unchanged request IDs. Cover reset and continuation behavior across an orchestrator restart. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6400e928-51a9-4a92-aead-de8b03854676 --- .../src/annotations/orchestrator.rs | 52 +++++++++++++------ crates/trident-acl-agent/src/core/config.rs | 22 ++++---- 2 files changed, 47 insertions(+), 27 deletions(-) diff --git a/crates/trident-acl-agent/src/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index 5adc45c7de..e8424b3efb 100644 --- a/crates/trident-acl-agent/src/annotations/orchestrator.rs +++ b/crates/trident-acl-agent/src/annotations/orchestrator.rs @@ -706,12 +706,9 @@ impl Orchestrator { Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); - // This is `handle_stage`'s own initial Nebraska call - definitionally - // the first one of a brand new update (stage always precedes - // finalize/commit) - so it's the one place `CookieLifetime::PerUpdate` - // forces a fresh cookie rather than inheriting whatever a prior - // update left behind. Every other Nebraska call in this module is a - // continuation and always passes `force_fresh: false`. + // Every staging attempt starts fresh in PerUpdate, including a + // restart with unchanged request IDs. Progress/finalize/commit + // reuse this attempt's cookie rather than resetting it again. let force_fresh = self.config.orchestration.cookie_lifetime == CookieLifetime::PerUpdate; let initial_cookie = self.load_nebraska_cookie(force_fresh); let cookie_enabled = self.config.orchestration.cookie_enabled; @@ -1597,17 +1594,17 @@ impl Orchestrator { /// round-trip proceeds cookie-less, the same as this host's very /// first-ever contact with Nebraska. /// - /// `force_fresh` is set by callers that are the first Nebraska call of - /// a brand new update (see `CookieLifetime::PerUpdate`'s docs) - when - /// enabled, it invalidates the persisted cookie too, so later reports - /// cannot reload a prior update's cookie if Nebraska issues no new one. + /// `force_fresh` marks the initial check of a staging attempt, including + /// a restarted stage (see `CookieLifetime::PerUpdate`). When enabled, + /// it invalidates the persisted cookie too, so later reports cannot + /// reload a previous attempt's cookie if Nebraska issues no new one. fn load_nebraska_cookie(&self, force_fresh: bool) -> Option { if !self.config.orchestration.cookie_enabled { return None; } if force_fresh { if let Err(err) = self.nebraska_cookie.save("") { - warn!("failed to clear Nebraska cookie for new update: {err}"); + warn!("failed to clear Nebraska cookie for staging attempt: {err}"); } return None; } @@ -1694,11 +1691,8 @@ impl Orchestrator { } }; let label = report.label(); - // Always a continuation of an already-started update (a progress or - // terminal event report), never the first Nebraska call of a new - // one - so `force_fresh` is always false here regardless of - // `cookie_lifetime`. See `handle_stage`'s own call for the one place - // that does force a fresh cookie. + // Reports continue the current attempt, including across reboot, + // so only handle_stage's initial check resets affinity. let initial_cookie = self.load_nebraska_cookie(false); let cookie_enabled = self.config.orchestration.cookie_enabled; let result = task::spawn_blocking(move || { @@ -2466,6 +2460,32 @@ mod tests { } } + #[tokio::test] + async fn per_update_cookie_resets_when_an_unfinished_stage_restarts() { + let dir = tempfile::tempdir().unwrap(); + let orchestrator = cookie_orchestrator(dir.path(), true, CookieLifetime::PerUpdate).await; + assert_eq!(orchestrator.load_nebraska_cookie(true), None); + orchestrator.save_nebraska_cookie(Some("ASLBSA=first-attempt".to_string())); + assert_eq!( + orchestrator.load_nebraska_cookie(false).as_deref(), + Some("ASLBSA=first-attempt"), + ); + drop(orchestrator); + + let restarted = cookie_orchestrator(dir.path(), true, CookieLifetime::PerUpdate).await; + assert_eq!( + restarted.nebraska_cookie.load().unwrap().as_deref(), + Some("ASLBSA=first-attempt"), + ); + assert_eq!(restarted.load_nebraska_cookie(true), None); + assert_eq!(restarted.nebraska_cookie.load().unwrap(), None); + restarted.save_nebraska_cookie(Some("ASLBSA=second-attempt".to_string())); + assert_eq!( + restarted.load_nebraska_cookie(false).as_deref(), + Some("ASLBSA=second-attempt"), + ); + } + #[tokio::test] async fn disabled_cookie_is_not_adopted_by_completion_report_retries() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index 50df127bbe..c53883c078 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -445,7 +445,7 @@ impl FromStr for Mode { } /// Controls whether a persisted Nebraska affinity cookie is reused across -/// updates, or deliberately dropped at the start of each new one. See +/// staging attempts, or reset before each attempt's initial check. See /// [`OrchestrationConfig::cookie_lifetime`]. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum CookieLifetime { @@ -453,15 +453,15 @@ pub enum CookieLifetime { /// replaces it. This is the originally-shipped behavior. #[default] Forever, - /// Force a fresh cookie at the start of every update: the very first - /// Nebraska call of a new update (`handle_stage`'s `check_for_update` - - /// definitionally the first call of any update, since stage always - /// precedes finalize/commit) never sends whatever cookie is currently - /// persisted, so Nebraska issues/re-confirms one for this update - /// specifically. Every later call within that same update (progress - /// events, finalize, the post-reboot commit) still reuses whatever that - /// first call established - this is "don't inherit a prior update's - /// session," not "never persist at all." + /// Start each staging attempt without a cookie: clear persisted + /// affinity before `handle_stage`'s initial Nebraska check, then persist + /// and reuse any cookie it establishes for progress, finalize, and + /// post-reboot completion. + /// + /// Restarting an unfinished stage starts a new attempt, even when the + /// target version, `nodeUpdateId`, and `operationId` are unchanged. + /// Its initial check also starts cookie-less; affinity is not preserved + /// across staging attempts. Resuming finalize/commit does not reset it. PerUpdate, } @@ -496,7 +496,7 @@ pub struct OrchestrationConfig { /// `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. pub cookie_enabled: bool, /// Whether a persisted cookie is reused forever, or deliberately - /// dropped at the start of every new update. See [`CookieLifetime`]. + /// reset before every staging attempt. See [`CookieLifetime`]. /// Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME` /// (`"forever"` or `"per-update"`). pub cookie_lifetime: CookieLifetime, From c9e1ee5ff7265567324a8d08d2d29a3ef7adf04e Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Thu, 22 Oct 2026 15:16:04 +0000 Subject: [PATCH 07/13] acl: make Nebraska affinity cookies opt-in Enable cookies explicitly in the ACL test image and Storm probes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 966c15e1-1a64-4a8a-b529-e97d48d63667 --- crates/trident-acl-agent/src/core/config.rs | 30 ++++++++++++++----- .../files/trident-acl-agent-override.conf | 1 + tools/storm/aclagent/tests/update.go | 14 +++++---- 3 files changed, 32 insertions(+), 13 deletions(-) diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index 46c53184be..2e27ed3158 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -73,9 +73,9 @@ pub const DEFAULT_KUBELET_KUBECONFIG: &str = "/var/lib/kubelet/kubeconfig"; /// Default annotation-key prefix. /// Override with `TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX`. pub const DEFAULT_ANNOTATION_PREFIX: &str = "acl.microsoft.com"; -/// Default for `cookie_enabled`: on, matching the originally-shipped -/// behavior. Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. -const DEFAULT_COOKIE_ENABLED: bool = true; +/// Cookies are opt-in. +/// Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. +const DEFAULT_COOKIE_ENABLED: bool = false; #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct AgentConfig { @@ -738,8 +738,8 @@ mod tests { PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH) ); assert!( - config.orchestration.cookie_enabled, - "cookie_enabled should default to on, matching the originally-shipped behavior" + !config.orchestration.cookie_enabled, + "cookie_enabled should default to off" ); assert_eq!( config.orchestration.cookie_lifetime, @@ -796,7 +796,7 @@ mod tests { "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH", "/var/lib/trident-acl-agent/custom-nebraska-cookie", ), - ("TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", "false"), + ("TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", "true"), ( "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", "per-update", @@ -838,7 +838,7 @@ mod tests { config.orchestration.cookie_path, PathBuf::from("/var/lib/trident-acl-agent/custom-nebraska-cookie") ); - assert!(!config.orchestration.cookie_enabled); + assert!(config.orchestration.cookie_enabled); assert_eq!( config.orchestration.cookie_lifetime, CookieLifetime::PerUpdate @@ -907,6 +907,22 @@ mod tests { assert!(format!("{err:#}").contains("bogus"), "{err:#}"); } + #[test] + fn cookie_enabled_is_opt_in() { + assert!(!AgentConfig::default().orchestration.cookie_enabled); + for (value, expected) in [("", false), ("false", false), ("true", true)] { + let config = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", + value, + )])) + .unwrap(); + assert_eq!( + config.orchestration.cookie_enabled, expected, + "cookie_enabled with value {value:?}" + ); + } + } + #[test] fn malformed_cookie_enabled_is_a_parse_error() { let err = AgentConfig::from_vars(vars(&[( diff --git a/tests/images/trident-vm-testimage/base/files/trident-acl-agent-override.conf b/tests/images/trident-vm-testimage/base/files/trident-acl-agent-override.conf index 2143d67f9c..d14bfe4352 100644 --- a/tests/images/trident-vm-testimage/base/files/trident-acl-agent-override.conf +++ b/tests/images/trident-vm-testimage/base/files/trident-acl-agent-override.conf @@ -8,5 +8,6 @@ ExecStart=/usr/bin/trident-acl-agent --verbosity trace --network-verbosity warn Environment=TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX=acl.azure.com Environment=TRIDENT_ACL_AGENT_CURRENT_VERSION_KEY=IMAGE_VERSION Environment=TRIDENT_ACL_AGENT_ORCHESTRATION_MODE=annotations +Environment=TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED=true Environment=TRIDENT_ACL_AGENT_CURRENT_VERSION_PATH=/etc/os-release Environment=TRIDENT_ACL_AGENT_CURRENT_VERSION_FALLBACK=always diff --git a/tools/storm/aclagent/tests/update.go b/tools/storm/aclagent/tests/update.go index 4fd7c7fb89..5e756aa9dd 100644 --- a/tools/storm/aclagent/tests/update.go +++ b/tools/storm/aclagent/tests/update.go @@ -368,17 +368,19 @@ func RunABUpdate(testConfig stormaclconfig.TestConfig, vmConfig stormvmconfig.Al // prepareVmForAclAgent's update-ca-trust + rehash install, exactly as // a real deployment would trust Nebraska's cert via the OS trust store. nebraskaEnv := map[string]string{ - "TRIDENT_ACL_AGENT_NEBRASKA_ENDPOINT": nebraskaServer, - "TRIDENT_ACL_AGENT_NEBRASKA_APP_ID": nebraskaAppID, - "TRIDENT_ACL_AGENT_NEBRASKA_TRACK": nebraskaTrack, + "TRIDENT_ACL_AGENT_NEBRASKA_ENDPOINT": nebraskaServer, + "TRIDENT_ACL_AGENT_NEBRASKA_APP_ID": nebraskaAppID, + "TRIDENT_ACL_AGENT_NEBRASKA_TRACK": nebraskaTrack, + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED": "true", } if err := expectValidateConnection(vmConfig.VMConfig, vmIP, "nebraska", true, nebraskaEnv); err != nil { return fmt.Errorf("env-var-override validate-connection check failed (expected success): %w", err) } wrongNebraskaEnv := map[string]string{ - "TRIDENT_ACL_AGENT_NEBRASKA_ENDPOINT": fmt.Sprintf("https://%s:%d", testConfig.HostEndpointIP, testConfig.NebraskaPort+1), - "TRIDENT_ACL_AGENT_NEBRASKA_APP_ID": nebraskaAppID, - "TRIDENT_ACL_AGENT_NEBRASKA_TRACK": nebraskaTrack, + "TRIDENT_ACL_AGENT_NEBRASKA_ENDPOINT": fmt.Sprintf("https://%s:%d", testConfig.HostEndpointIP, testConfig.NebraskaPort+1), + "TRIDENT_ACL_AGENT_NEBRASKA_APP_ID": nebraskaAppID, + "TRIDENT_ACL_AGENT_NEBRASKA_TRACK": nebraskaTrack, + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED": "true", } if err := expectValidateConnection(vmConfig.VMConfig, vmIP, "nebraska", false, wrongNebraskaEnv); err != nil { return fmt.Errorf("env-var-override validate-connection check failed (expected failure against wrong port): %w", err) From 7fd70acb06ca666594635b9ca6c1dbc7bb0bd57c Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Thu, 22 Oct 2026 19:29:04 +0000 Subject: [PATCH 08/13] acl: add per-update and per-message cookie policies Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5638630c-63c7-440d-b89f-d5add6423071 --- .../src/annotations/orchestrator.rs | 215 +++++------ .../trident-acl-agent/src/connection_check.rs | 9 +- crates/trident-acl-agent/src/core/config.rs | 90 ++--- .../src/core/nebraska/README.md | 33 ++ .../src/core/nebraska/client.rs | 358 +++++++++++++----- .../src/core/nebraska/cookie_store.rs | 180 ++++++--- .../src/core/nebraska/mod.rs | 2 +- crates/trident-acl-agent/src/lib.rs | 116 ++++-- crates/trident-acl-agent/src/omahaonly/mod.rs | 10 +- tools/storm/aclagent/proxies/nebraska.go | 6 +- tools/storm/aclagent/tests/update.go | 6 +- 11 files changed, 666 insertions(+), 359 deletions(-) diff --git a/crates/trident-acl-agent/src/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index e8424b3efb..3956a5ab42 100644 --- a/crates/trident-acl-agent/src/annotations/orchestrator.rs +++ b/crates/trident-acl-agent/src/annotations/orchestrator.rs @@ -33,7 +33,7 @@ use crate::{ SCHEMA_VERSION, }, core::{ - config::{AgentConfig, CookieLifetime}, + config::AgentConfig, cosi_sha384, nebraska::{CheckOutcome, Client as NebraskaClient, CookieStore, ProgressEvent}, trident::{CompletedResponse, TridentClient, TridentClientError}, @@ -706,28 +706,22 @@ impl Orchestrator { Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); - // Every staging attempt starts fresh in PerUpdate, including a - // restart with unchanged request IDs. Progress/finalize/commit - // reuse this attempt's cookie rather than resetting it again. - let force_fresh = self.config.orchestration.cookie_lifetime == CookieLifetime::PerUpdate; - let initial_cookie = self.load_nebraska_cookie(force_fresh); let cookie_enabled = self.config.orchestration.cookie_enabled; - let (check_result, cookie) = task::spawn_blocking({ + let cookie_store = self.nebraska_cookie.clone(); + let cookie_lifetime = self.config.orchestration.cookie_lifetime; + let outcome = task::spawn_blocking({ let current_version = current_version.clone(); move || { - let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id) - .with_cookie_enabled(cookie_enabled); - if let Some(cookie) = initial_cookie { - client = client.with_cookie(cookie); + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + if cookie_enabled { + client = client.with_cookie_store(cookie_store, cookie_lifetime); } - let result = client.check_for_update(¤t_version); - (result, client.cookie()) + client.check_for_update(¤t_version) } }) .await - .context("Nebraska query task panicked")?; - self.save_nebraska_cookie(cookie); - let outcome = check_result.context("Nebraska query failed")?; + .context("Nebraska query task panicked")? + .context("Nebraska query failed")?; let offered = match outcome { CheckOutcome::UpToDate => { let status = UpdateStatus::new( @@ -1586,52 +1580,6 @@ impl Orchestrator { } } - /// Loads the persisted Nebraska affinity cookie, honoring - /// `cookie_enabled`/`force_fresh`, and warning (not failing) if the - /// store can't be read. The whole mechanism is a best-effort - /// session-affinity hint, never a requirement for talking to Nebraska - - /// a load failure (or a disabled/fresh session) just means this - /// round-trip proceeds cookie-less, the same as this host's very - /// first-ever contact with Nebraska. - /// - /// `force_fresh` marks the initial check of a staging attempt, including - /// a restarted stage (see `CookieLifetime::PerUpdate`). When enabled, - /// it invalidates the persisted cookie too, so later reports cannot - /// reload a previous attempt's cookie if Nebraska issues no new one. - fn load_nebraska_cookie(&self, force_fresh: bool) -> Option { - if !self.config.orchestration.cookie_enabled { - return None; - } - if force_fresh { - if let Err(err) = self.nebraska_cookie.save("") { - warn!("failed to clear Nebraska cookie for staging attempt: {err}"); - } - return None; - } - match self.nebraska_cookie.load() { - Ok(cookie) => cookie, - Err(err) => { - warn!("failed to load Nebraska cookie, proceeding without one: {err}"); - None - } - } - } - - /// Persists a cookie a Nebraska round-trip ended up holding, if any - - /// unless `cookie_enabled` is off, in which case this is a no-op (so a - /// mid-flight toggle from enabled to disabled doesn't keep writing a - /// file a later re-enable would then pick back up unexpectedly). - fn save_nebraska_cookie(&self, cookie: Option) { - if !self.config.orchestration.cookie_enabled { - return; - } - if let Some(cookie) = cookie { - if let Err(err) = self.nebraska_cookie.save(&cookie) { - warn!("failed to persist Nebraska cookie: {err}"); - } - } - } - /// Sends a single Nebraska event report, logging (but never /// propagating) failure. /// @@ -1691,17 +1639,15 @@ impl Orchestrator { } }; let label = report.label(); - // Reports continue the current attempt, including across reboot, - // so only handle_stage's initial check resets affinity. - let initial_cookie = self.load_nebraska_cookie(false); let cookie_enabled = self.config.orchestration.cookie_enabled; + let cookie_store = self.nebraska_cookie.clone(); + let cookie_lifetime = self.config.orchestration.cookie_lifetime; let result = task::spawn_blocking(move || { - let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id) - .with_cookie_enabled(cookie_enabled); - if let Some(cookie) = initial_cookie { - client = client.with_cookie(cookie); + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + if cookie_enabled { + client = client.with_cookie_store(cookie_store, cookie_lifetime); } - let outcome = match report { + match report { NebraskaReport::Progress { version, event } => { client.report_progress(&version, event) } @@ -1711,18 +1657,12 @@ impl Orchestrator { NebraskaReport::Failed { previous, current } => { client.report_failure(&previous, ¤t) } - }; - (outcome, client.cookie()) + } }) .await; match result { - Ok((outcome, cookie)) => { - self.save_nebraska_cookie(cookie); - match outcome { - Ok(()) => debug!("reported Nebraska '{label}' event"), - Err(err) => warn!("Nebraska '{label}' event report failed: {err}"), - } - } + Ok(Ok(())) => debug!("reported Nebraska '{label}' event"), + Ok(Err(err)) => warn!("Nebraska '{label}' event report failed: {err}"), Err(err) => warn!("Nebraska '{label}' event report task panicked: {err}"), } } @@ -2338,7 +2278,10 @@ mod tests { const MOCK_RPC_TIMEOUT: Duration = Duration::from_secs(5); use crate::{ annotations::{RequestedOperation, SCHEMA_VERSION}, - core::trident::mock::{connect_mock_client, MockTridentdConfig, Outcome}, + core::{ + config::CookieLifetime, + trident::mock::{connect_mock_client, MockTridentdConfig, Outcome}, + }, }; async fn cookie_orchestrator( @@ -2379,7 +2322,7 @@ mod tests { let mut server = mockito::Server::new_async().await; for (enabled, lifetime) in [ (true, CookieLifetime::PerUpdate), - (true, CookieLifetime::Forever), + (true, CookieLifetime::PerMessage), (false, CookieLifetime::PerUpdate), ] { let dir = tempfile::tempdir().unwrap(); @@ -2388,7 +2331,7 @@ mod tests { .nebraska_cookie .save("ASLBSA=previous") .unwrap(); - let expected_cookie = if enabled && lifetime == CookieLifetime::Forever { + let expected_cookie = if enabled && lifetime == CookieLifetime::PerMessage { Matcher::Exact("ASLBSA=previous".to_string()) } else { Matcher::Missing @@ -2403,27 +2346,8 @@ mod tests { .expect(1) .create_async() .await; - let initial_cookie = - orchestrator.load_nebraska_cookie(lifetime == CookieLifetime::PerUpdate); let endpoint = Url::parse(&format!("{}/v1/update/", server.url())).unwrap(); - let check_endpoint = endpoint.clone(); - let cookie = task::spawn_blocking(move || { - let mut client = NebraskaClient::new( - check_endpoint, - "app-1", - "stable", - crate::core::nebraska::MachineId::new("mid-1").unwrap(), - ) - .with_cookie_enabled(enabled); - if let Some(cookie) = initial_cookie { - client = client.with_cookie(cookie); - } - client.check_for_update(&Version::new(1, 0, 0)).unwrap(); - client.cookie() - }) - .await - .unwrap(); - orchestrator.save_nebraska_cookie(cookie); + initial_cookie_check(&orchestrator, endpoint.clone()).await; check.assert_async().await; let progress = server @@ -2431,6 +2355,7 @@ mod tests { .match_header("cookie", expected_cookie) .match_body(Matcher::Regex(""#; let dir = tempfile::tempdir().unwrap(); let orchestrator = cookie_orchestrator(dir.path(), true, CookieLifetime::PerUpdate).await; - assert_eq!(orchestrator.load_nebraska_cookie(true), None); - orchestrator.save_nebraska_cookie(Some("ASLBSA=first-attempt".to_string())); - assert_eq!( - orchestrator.load_nebraska_cookie(false).as_deref(), - Some("ASLBSA=first-attempt"), - ); + let first = server + .mock("POST", "/v1/update/") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_header("set-cookie", "ASLBSA=first-attempt; Path=/") + .with_body(response) + .expect(1) + .create_async() + .await; + initial_cookie_check(&orchestrator, endpoint.clone()).await; + first.assert_async().await; drop(orchestrator); let restarted = cookie_orchestrator(dir.path(), true, CookieLifetime::PerUpdate).await; @@ -2477,11 +2432,47 @@ mod tests { restarted.nebraska_cookie.load().unwrap().as_deref(), Some("ASLBSA=first-attempt"), ); - assert_eq!(restarted.load_nebraska_cookie(true), None); - assert_eq!(restarted.nebraska_cookie.load().unwrap(), None); - restarted.save_nebraska_cookie(Some("ASLBSA=second-attempt".to_string())); + let completion = server + .mock("POST", "/v1/update/") + .match_header("cookie", "ASLBSA=first-attempt") + .with_status(200) + .with_header("set-cookie", "ASLBSA=ignored; Path=/") + .with_body(response) + .expect(1) + .create_async() + .await; + let mut request = request(RequestedOperation::Finalize); + request.server = Some(endpoint.clone()); + request.app_id = Some("app-1".to_string()); + request.track = Some("stable".to_string()); + restarted + .report_nebraska_event( + &request, + NebraskaReport::Completed { + previous: Version::new(1, 0, 0), + current: Version::new(2, 0, 0), + }, + ) + .await; + completion.assert_async().await; assert_eq!( - restarted.load_nebraska_cookie(false).as_deref(), + restarted.nebraska_cookie.load().unwrap().as_deref(), + Some("ASLBSA=first-attempt") + ); + + let next = server + .mock("POST", "/v1/update/") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_header("set-cookie", "ASLBSA=second-attempt; Path=/") + .with_body(response) + .expect(1) + .create_async() + .await; + initial_cookie_check(&restarted, endpoint).await; + next.assert_async().await; + assert_eq!( + restarted.nebraska_cookie.load().unwrap().as_deref(), Some("ASLBSA=second-attempt"), ); } diff --git a/crates/trident-acl-agent/src/connection_check.rs b/crates/trident-acl-agent/src/connection_check.rs index 2530f9273e..c53b390af8 100644 --- a/crates/trident-acl-agent/src/connection_check.rs +++ b/crates/trident-acl-agent/src/connection_check.rs @@ -77,16 +77,12 @@ pub async fn validate_connection( // rather than the pure reachability check documented on // ConnectionTarget::Nebraska above. // - // Shares the same persisted CookieStore (cookie_path) the - // orchestrator uses: this probe talks to the same Nebraska - // instance real traffic does, so it must participate in the - // affinity-cookie dance too - otherwise every - // `--validate-connection nebraska` run is its own permanent - // "first contact" that never adopts or persists a cookie. + // Probes share affinity without starting a new update interval. let endpoint_for_task = endpoint.clone(); let track = config.nebraska.track.clone(); let cookie_store = nebraska::CookieStore::new(config.orchestration.cookie_path.clone()); let cookie_enabled = config.orchestration.cookie_enabled; + let cookie_lifetime = config.orchestration.cookie_lifetime; task::spawn_blocking(move || { check_nebraska_reachable( &endpoint_for_task, @@ -95,6 +91,7 @@ pub async fn validate_connection( IdSource::MachineIdHashed, &cookie_store, cookie_enabled, + cookie_lifetime, ) }) .await diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index 2e27ed3158..b890b95433 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -34,6 +34,8 @@ use url::Url; use crate::{cosi_sha384::CosiSha384Source, DEFAULT_NEBRASKA_APP_ID, DEFAULT_NEBRASKA_TRACK}; +pub use super::nebraska::CookieLifetime; + const ENV_PREFIX_NEBRASKA: &str = "TRIDENT_ACL_AGENT_NEBRASKA_"; const ENV_PREFIX_COSI_SHA384: &str = "TRIDENT_ACL_AGENT_COSI_SHA384_"; const ENV_PREFIX_KUBERNETES: &str = "TRIDENT_ACL_AGENT_KUBERNETES_"; @@ -494,49 +496,13 @@ impl FromStr for Mode { } } -/// Controls whether a persisted Nebraska affinity cookie is reused across -/// staging attempts, or reset before each attempt's initial check. See -/// [`OrchestrationConfig::cookie_lifetime`]. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum CookieLifetime { - /// Reuse a persisted cookie across every update until Nebraska itself - /// replaces it. This is the originally-shipped behavior. - #[default] - Forever, - /// Start each staging attempt without a cookie: clear persisted - /// affinity before `handle_stage`'s initial Nebraska check, then persist - /// and reuse any cookie it establishes for progress, finalize, and - /// post-reboot completion. - /// - /// Restarting an unfinished stage starts a new attempt, even when the - /// target version, `nodeUpdateId`, and `operationId` are unchanged. - /// Its initial check also starts cookie-less; affinity is not preserved - /// across staging attempts. Resuming finalize/commit does not reset it. - PerUpdate, -} - -impl FromStr for CookieLifetime { - type Err = Error; - - fn from_str(s: &str) -> Result { - match s { - "forever" => Ok(CookieLifetime::Forever), - "per-update" => Ok(CookieLifetime::PerUpdate), - other => Err(anyhow!( - "unknown cookie lifetime {other:?} (expected \"forever\" or \"per-update\")" - )), - } - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct OrchestrationConfig { pub mode: Mode, pub state_path: PathBuf, /// Path to the persisted Nebraska affinity cookie (see - /// `core::nebraska::CookieStore`). Deliberately a separate file/field - /// from `state_path`: the cookie has no relationship to any particular - /// update operation, unlike everything `state_path` persists. + /// `core::nebraska::CookieStore`). Kept separate from `state_path` so + /// cookie persistence does not change the serialized update-state schema. pub cookie_path: PathBuf, /// Blanket kill switch for the whole cookie mechanism. When `false`, no /// `Cookie:` header is ever sent and no `Set-Cookie` is ever adopted - @@ -545,10 +511,10 @@ pub struct OrchestrationConfig { /// later just picks the stale file back up as-is. Override with /// `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED`. pub cookie_enabled: bool, - /// Whether a persisted cookie is reused forever, or deliberately - /// reset before every staging attempt. See [`CookieLifetime`]. + /// Whether only the initial update response or every message response + /// may replace the cookie. Defaults to per-update. See [`CookieLifetime`]. /// Override with `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME` - /// (`"forever"` or `"per-update"`). + /// (`"per-update"` or `"per-message"`). pub cookie_lifetime: CookieLifetime, /// Placeholder default pending real data from storm aclagent scenario runs. pub stage_timeout: Duration, @@ -579,7 +545,7 @@ impl Default for OrchestrationConfig { state_path: PathBuf::from(DEFAULT_STATE_PATH), cookie_path: PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH), cookie_enabled: DEFAULT_COOKIE_ENABLED, - cookie_lifetime: CookieLifetime::Forever, + cookie_lifetime: CookieLifetime::PerUpdate, stage_timeout: DEFAULT_STAGE_TIMEOUT, finalize_timeout: DEFAULT_FINALIZE_TIMEOUT, heartbeat_interval: DEFAULT_HEARTBEAT_INTERVAL, @@ -621,6 +587,8 @@ mod tests { fn startup_summary_includes_effective_defaults() { let config = AgentConfig::from_vars(vec![]).unwrap(); let summary = format!("{:?}", config.redacted_for_logging()); + assert!(summary.contains("cookie_enabled: false"), "{summary}"); + assert!(summary.contains("cookie_lifetime: PerUpdate"), "{summary}"); for expected in [ "AgentConfig", "nebraska.endpoint", @@ -743,7 +711,7 @@ mod tests { ); assert_eq!( config.orchestration.cookie_lifetime, - CookieLifetime::Forever + CookieLifetime::PerUpdate ); assert_eq!(config.orchestration.stage_timeout, DEFAULT_STAGE_TIMEOUT); assert_eq!( @@ -799,7 +767,7 @@ mod tests { ("TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", "true"), ( "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", - "per-update", + "per-message", ), ("TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT", "21m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT", "11m"), @@ -841,7 +809,7 @@ mod tests { assert!(config.orchestration.cookie_enabled); assert_eq!( config.orchestration.cookie_lifetime, - CookieLifetime::PerUpdate + CookieLifetime::PerMessage ); assert_eq!( config.orchestration.stage_timeout, @@ -933,6 +901,38 @@ mod tests { assert!(format!("{err:#}").contains("yes please"), "{err:#}"); } + #[test] + fn cookie_lifetime_modes_and_empty_default() { + assert_eq!(CookieLifetime::default(), CookieLifetime::PerUpdate); + for (value, expected) in [ + ("", CookieLifetime::PerUpdate), + ("per-update", CookieLifetime::PerUpdate), + ("per-message", CookieLifetime::PerMessage), + ] { + let config = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + value, + )])) + .unwrap(); + assert_eq!(config.orchestration.cookie_lifetime, expected); + let summary = format!("{:?}", config.redacted_for_logging()); + assert!(summary.contains(&format!("cookie_lifetime: {expected:?}"))); + } + } + + #[test] + fn removed_forever_cookie_lifetime_is_rejected() { + let err = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + "forever", + )])) + .unwrap_err(); + let message = format!("{err:#}"); + assert!(message.contains("forever"), "{message}"); + assert!(message.contains("per-update"), "{message}"); + assert!(message.contains("per-message"), "{message}"); + } + #[test] fn malformed_cookie_lifetime_is_a_parse_error() { let err = AgentConfig::from_vars(vars(&[( diff --git a/crates/trident-acl-agent/src/core/nebraska/README.md b/crates/trident-acl-agent/src/core/nebraska/README.md index 702a9540b0..d8e92f7109 100644 --- a/crates/trident-acl-agent/src/core/nebraska/README.md +++ b/crates/trident-acl-agent/src/core/nebraska/README.md @@ -48,6 +48,39 @@ new ones does not break parsing. ## Usage +### Affinity cookies + +Cookies are opt-in. The agent uses these environment variables: + +| Variable | Default | +| --- | --- | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED` | `false` | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH` | `/var/lib/trident-acl-agent/nebraska-cookie` | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME` | `per-update` | + +`per-update` sends each initial `check_for_update` without a cookie, clears the +previous stored value, and persists only that response's cookie. Progress, +failure, completion, and diagnostic requests reuse it without adopting +replacements. An initial response without a cookie leaves the update cookie-less. +A restarted staging attempt starts a new interval, even with unchanged request +IDs; resuming finalize/completion does not. + +`per-message` sends the stored cookie on every request and persists each +response's replacement for the next request, including completion retries. +Responses without `Set-Cookie` leave the file unchanged. `forever` is not supported. + +Standalone clients opt in with `Client::with_cookie_store(store, lifetime)`, using +`CookieStore` and `CookieLifetime` from this module. A client without a store +neither sends nor adopts cookies. Both agent orchestration modes and connectivity +probes honor the enabled switch and policy. + +Only HTTP-success responses with readable bodies can supply a cookie; XML +validation happens afterward. Cookies are stored with owner-only permissions and +atomic replacement, never logged. I/O failures warn without failing the update; +the affinity guarantees are therefore best-effort when storage fails. Concurrent +cookie-bearing responses use last-successful-write-wins ordering. Configure +independent processes consistently when sharing a cookie path. + ### Poll for an update ```rust,no_run diff --git a/crates/trident-acl-agent/src/core/nebraska/client.rs b/crates/trident-acl-agent/src/core/nebraska/client.rs index 09c5a081aa..7a6cf59e2d 100644 --- a/crates/trident-acl-agent/src/core/nebraska/client.rs +++ b/crates/trident-acl-agent/src/core/nebraska/client.rs @@ -1,12 +1,13 @@ //! The high-level [`Client`] for talking to a Nebraska server. -use std::{cell::RefCell, collections::BTreeMap, thread, time::Duration}; +use std::{collections::BTreeMap, thread, time::Duration}; use log::{debug, trace, warn}; use semver::Version; use url::Url; use super::{ + cookie_store::{CookieLifetime, CookieStore}, error::NebraskaError, event::{ProgressEvent, TerminalEvent}, id::MachineId, @@ -331,17 +332,8 @@ pub struct Client { track: String, machine_id: MachineId, client_version: Option, - /// Nebraska's load-balancer affinity cookie (e.g. an App-Service-style - /// `ASLBSA` cookie), carried on every request once known and updated - /// from every response that sets one. `RefCell` because every method - /// here takes `&self` (see the type docs), but a single `Client` still - /// needs to pick up a cookie mid-sequence (e.g. during - /// [`check_for_update`](Client::check_for_update)) and carry it into its - /// own later calls. Seeding it across *separate* `Client`s - including - /// across a reboot - is the caller's job: see [`with_cookie`](Client::with_cookie) - /// and [`cookie`](Client::cookie). - cookie: RefCell>, - cookie_enabled: bool, + cookie_store: Option, + cookie_lifetime: CookieLifetime, transport: T, } @@ -375,43 +367,20 @@ impl Client { track: track.into(), machine_id, client_version: None, - cookie: RefCell::new(None), - cookie_enabled: true, + cookie_store: None, + cookie_lifetime: CookieLifetime::default(), transport, } } - /// Enables or disables affinity cookies, including response adoption - /// and retries. Disabling also discards any previously seeded cookie. - pub fn with_cookie_enabled(mut self, enabled: bool) -> Self { - self.cookie_enabled = enabled; - if !enabled { - *self.cookie.borrow_mut() = None; - } + /// Enables best-effort file-backed cookies. Without a store, cookies + /// are neither sent nor adopted. Each request reads the store anew. + pub fn with_cookie_store(mut self, store: CookieStore, lifetime: CookieLifetime) -> Self { + self.cookie_store = Some(store); + self.cookie_lifetime = lifetime; self } - /// Seeds this client with a previously-persisted Nebraska affinity - /// cookie (see the [`cookie`](Client::cookie) getter), so it is sent on - /// this client's very first request rather than only picked up from a - /// response. Ignored when cookies are disabled. - pub fn with_cookie(self, cookie: impl Into) -> Self { - if self.cookie_enabled { - *self.cookie.borrow_mut() = Some(cookie.into()); - } - self - } - - /// The most recent Nebraska affinity cookie this client has seen, if - /// any - either seeded via [`with_cookie`](Client::with_cookie) or - /// picked up from a response. Callers that need the cookie to outlive - /// this `Client` (across a reboot, or just a later process invocation) - /// must read this after use and persist it themselves; nothing here - /// does that automatically. - pub fn cookie(&self) -> Option { - self.cookie.borrow().clone() - } - /// Names the **updater** — the program driving the update — in /// ``. /// @@ -438,12 +407,15 @@ impl Client { /// `current_version` **must be the real version** and valid semver: a client /// reporting `0.0.0` is offered an update on every poll forever, and a /// non-semver version fails instance registration server-side. + /// + /// This is the initial update communication: per-update cookies reset + /// here, not on completion's batched update check or connectivity probes. pub fn check_for_update( &self, current_version: &Version, ) -> Result { let app = self.app(current_version).with_update_check(); - let response = self.send(app)?; + let response = self.send(app, true)?; self.interpret_check(response) } @@ -475,7 +447,7 @@ impl Client { /// the server is reachable and resolved the app/track. pub fn probe(&self, current_version: &Version) -> Result<(), NebraskaError> { let app = self.app(current_version); - let response = self.send(app)?; + let response = self.send(app, false)?; let app_response = self.app_response(&response)?; if app_response.status.is_ok() || app_response.status.is_update_in_progress() { Ok(()) @@ -496,7 +468,7 @@ impl Client { event: ProgressEvent, ) -> Result<(), NebraskaError> { let app = self.app(current_version).with_event(event.wire()); - let response = self.send(app)?; + let response = self.send(app, false)?; self.require_events_accepted(&response, 1)?; Ok(()) } @@ -602,7 +574,7 @@ impl Client { ) .with_ping() .with_update_check(); - let response = self.send(app)?; + let response = self.send(app, false)?; // Landing the terminal event is what this request exists for, and a // dropped event is the only failure here that can wedge the instance — @@ -639,7 +611,7 @@ impl Client { let app = self .app(current_version) .with_event_from_version(TerminalEvent::Failed.wire(), previous_version.to_string()); - let response = self.send(app)?; + let response = self.send(app, false)?; self.require_events_accepted(&response, 1)?; Ok(()) } @@ -655,7 +627,7 @@ impl Client { } /// Serializes, sends, and parses a request/response round-trip. - fn send(&self, app: App) -> Result { + fn send(&self, app: App, initial_update: bool) -> Result { let request = wire::request_for(app, self.client_version.clone()); let body = request .to_xml() @@ -667,17 +639,29 @@ impl Client { ); // Deliberately never logged: this is a session-affinity credential, // not protocol content. - let cookie = if self.cookie_enabled { - self.cookie.borrow().clone() - } else { - None - }; + let cookie = self.cookie_store.as_ref().and_then(|store| { + if initial_update && self.cookie_lifetime == CookieLifetime::PerUpdate { + if let Err(err) = store.save("") { + warn!("failed to clear Nebraska cookie for staging attempt: {err}"); + } + return None; + } + match store.load() { + Ok(cookie) => cookie, + Err(err) => { + warn!("failed to load Nebraska cookie, proceeding without one: {err}"); + None + } + } + }); let (text, set_cookie) = self.transport .post_xml(&self.endpoint, &body, cookie.as_deref())?; - if self.cookie_enabled { - if let Some(set_cookie) = set_cookie { - *self.cookie.borrow_mut() = Some(set_cookie); + if initial_update || self.cookie_lifetime == CookieLifetime::PerMessage { + if let (Some(store), Some(set_cookie)) = (&self.cookie_store, set_cookie) { + if let Err(err) = store.save(&set_cookie) { + warn!("failed to persist Nebraska cookie: {err}"); + } } } trace!("Nebraska response: {} bytes", text.len()); @@ -1438,46 +1422,58 @@ mod tests { #[test] fn completion_retries_honor_cookie_enabled() { - for enabled in [true, false] { - let mut client = client_with_responses(vec![IN_PROGRESS, NO_UPDATE]) - .with_cookie("ASLBSA=seeded") - .with_cookie_enabled(enabled); - client.transport.set_cookie = Some("ASLBSA=fresh".to_string()); - - let outcome = client - .complete_after_reboot_with_policy( - fast_policy(), - &Version::new(1, 0, 0), - &Version::new(2, 0, 0), - ) - .unwrap(); - - assert_eq!(outcome, CheckOutcome::UpToDate); - assert_eq!(client.transport.calls.get(), 2); - let expected = if enabled { - vec![ - Some("ASLBSA=seeded".to_string()), - Some("ASLBSA=fresh".to_string()), - ] - } else { - vec![None, None] - }; - assert_eq!(*client.transport.cookies_sent.borrow(), expected); - assert_eq!( - client.cookie().as_deref(), - enabled.then_some("ASLBSA=fresh") - ); + for lifetime in [CookieLifetime::PerUpdate, CookieLifetime::PerMessage] { + for enabled in [true, false] { + let (_dir, store) = cookie_store(); + store.save("ASLBSA=seeded").unwrap(); + let mut client = client_with_responses(vec![IN_PROGRESS, NO_UPDATE]); + if enabled { + client = client.with_cookie_store(store.clone(), lifetime); + } + client.transport.set_cookie = Some("ASLBSA=fresh".to_string()); + + let outcome = client + .complete_after_reboot_with_policy( + fast_policy(), + &Version::new(1, 0, 0), + &Version::new(2, 0, 0), + ) + .unwrap(); + + assert_eq!(outcome, CheckOutcome::UpToDate); + assert_eq!(client.transport.calls.get(), 2); + let expected = if enabled { + vec![ + Some("ASLBSA=seeded".to_string()), + Some(if lifetime == CookieLifetime::PerMessage { + "ASLBSA=fresh".to_string() + } else { + "ASLBSA=seeded".to_string() + }), + ] + } else { + vec![None, None] + }; + assert_eq!(*client.transport.cookies_sent.borrow(), expected); + assert_eq!( + store.load().unwrap().as_deref(), + Some(if enabled && lifetime == CookieLifetime::PerMessage { + "ASLBSA=fresh" + } else { + "ASLBSA=seeded" + }) + ); + } } } #[test] - fn disabled_cookies_cannot_be_seeded() { - let client = client_with(NO_UPDATE) - .with_cookie_enabled(false) - .with_cookie("ASLBSA=seeded"); + fn cookies_without_a_store_are_not_adopted() { + let client = client_with(NO_UPDATE); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=fresh".to_string()); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); client.check_for_update(&Version::new(1, 0, 0)).unwrap(); assert_eq!(*client.transport.last_cookie_sent.borrow(), None); - assert_eq!(client.cookie(), None); } #[test] @@ -1661,36 +1657,52 @@ mod tests { ); } + fn cookie_store() -> (tempfile::TempDir, CookieStore) { + let dir = tempfile::tempdir().unwrap(); + let store = CookieStore::new(dir.path().join("cookie")); + (dir, store) + } + #[test] fn no_cookie_is_sent_until_one_is_known() { - let client = client_with(NO_UPDATE); + let (_dir, store) = cookie_store(); + let client = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerMessage); client.check_for_update(&Version::new(1, 0, 0)).unwrap(); assert_eq!(*client.transport.last_cookie_sent.borrow(), None); - assert_eq!(client.cookie(), None); + assert_eq!(store.load().unwrap(), None); + assert!(!store.path().exists()); } #[test] - fn with_cookie_seeds_the_first_request() { - let client = client_with(NO_UPDATE).with_cookie("ASLBSA=seeded"); + fn stored_cookie_seeds_the_first_per_message_request() { + let (_dir, store) = cookie_store(); + store.save("ASLBSA=seeded").unwrap(); + let client = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerMessage); + let modified = std::fs::metadata(store.path()).unwrap().modified().unwrap(); client.check_for_update(&Version::new(1, 0, 0)).unwrap(); assert_eq!( client.transport.last_cookie_sent.borrow().as_deref(), Some("ASLBSA=seeded") ); + assert_eq!( + std::fs::metadata(store.path()).unwrap().modified().unwrap(), + modified + ); } #[test] fn cookie_from_a_response_is_echoed_on_the_next_request() { - // Two sequential calls on the *same* Client, mirroring what happens - // mid-operation (e.g. check_for_update picking up a cookie that a - // later report_progress on the same Client must then carry). + let (_dir, store) = cookie_store(); let client = Client::with_transport( Url::parse("https://nebraska.example/v1/update/").unwrap(), "app-1", "stable", MachineId::new("mid-1").unwrap(), MockTransport::new(NO_UPDATE).with_set_cookie("ASLBSA=fresh"), - ); + ) + .with_cookie_store(store.clone(), CookieLifetime::PerMessage); client.check_for_update(&Version::new(1, 0, 0)).unwrap(); assert_eq!( @@ -1698,14 +1710,162 @@ mod tests { None, "first request has no cookie to send yet" ); - assert_eq!(client.cookie().as_deref(), Some("ASLBSA=fresh")); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=fresh")); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=replacement".to_string()); client.check_for_update(&Version::new(1, 0, 0)).unwrap(); assert_eq!( client.transport.last_cookie_sent.borrow().as_deref(), Some("ASLBSA=fresh"), "second request must echo back the cookie the first response set" ); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=replacement")); + *client.transport.set_cookie.borrow_mut() = None; + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=replacement") + ); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=replacement")); + } + + #[test] + fn per_update_only_initial_response_sets_cookie() { + let (_dir, store) = cookie_store(); + store.save("ASLBSA=previous").unwrap(); + let client = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerUpdate); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=initial".to_string()); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=initial")); + + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=ignored".to_string()); + client + .report_progress(&Version::new(1, 0, 0), ProgressEvent::DownloadStarted) + .unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=initial") + ); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=initial")); + + let resumed = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerUpdate); + *resumed.transport.set_cookie.borrow_mut() = Some("ASLBSA=also-ignored".to_string()); + resumed + .try_complete_after_reboot(&Version::new(1, 0, 0), &Version::new(2, 0, 0)) + .unwrap(); + assert_eq!( + resumed.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=initial") + ); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=initial")); + + *resumed.transport.set_cookie.borrow_mut() = Some("ASLBSA=next-update".to_string()); + resumed.check_for_update(&Version::new(2, 0, 0)).unwrap(); + assert_eq!(*resumed.transport.last_cookie_sent.borrow(), None); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=next-update")); + } + + #[test] + fn per_update_missing_initial_cookie_cannot_be_established_later() { + let (_dir, store) = cookie_store(); + store.save("ASLBSA=previous").unwrap(); + let client = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerUpdate); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert_eq!(store.load().unwrap(), None); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=ignored".to_string()); + client.probe(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert_eq!(store.load().unwrap(), None); + } + + #[test] + fn cookie_store_edits_and_deletion_are_visible_to_existing_client() { + let (_dir, store) = cookie_store(); + let client = + client_with(NO_UPDATE).with_cookie_store(store.clone(), CookieLifetime::PerMessage); + store.save("ASLBSA=external").unwrap(); + client.probe(&Version::new(1, 0, 0)).unwrap(); + assert_eq!( + client.transport.last_cookie_sent.borrow().as_deref(), + Some("ASLBSA=external") + ); + std::fs::remove_file(store.path()).unwrap(); + client.probe(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + assert!(!store.path().exists()); + } + + #[test] + fn response_without_cookie_does_not_overwrite_another_requests_cookie() { + struct UpdatingTransport(CookieStore); + impl Transport for UpdatingTransport { + fn post_xml( + &self, + _endpoint: &Url, + _body: &[u8], + cookie: Option<&str>, + ) -> Result<(String, Option), NebraskaError> { + assert_eq!(cookie, Some("ASLBSA=old")); + self.0.save("ASLBSA=concurrent").unwrap(); + Ok((NO_UPDATE.to_string(), None)) + } + } + let (_dir, store) = cookie_store(); + store.save("ASLBSA=old").unwrap(); + let client = Client::with_transport( + Url::parse("https://nebraska.example/v1/update/").unwrap(), + "app-1", + "stable", + MachineId::new("mid-1").unwrap(), + UpdatingTransport(store.clone()), + ) + .with_cookie_store(store.clone(), CookieLifetime::PerMessage); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=concurrent")); + } + + #[test] + fn invalid_stored_cookie_does_not_prevent_the_request() { + let (_dir, store) = cookie_store(); + std::fs::write(store.path(), "ASLBSA=bad\r\nInjected: value").unwrap(); + let client = client_with(NO_UPDATE).with_cookie_store(store, CookieLifetime::PerMessage); + client.check_for_update(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + } + + #[test] + fn cookie_io_failures_do_not_fail_requests() { + let dir = tempfile::tempdir().unwrap(); + let parent = dir.path().join("not-a-directory"); + std::fs::write(&parent, "block parent creation").unwrap(); + for lifetime in [CookieLifetime::PerUpdate, CookieLifetime::PerMessage] { + let client = client_with(NO_UPDATE) + .with_cookie_store(CookieStore::new(parent.join("cookie")), lifetime); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=new".to_string()); + assert_eq!( + client.check_for_update(&Version::new(1, 0, 0)).unwrap(), + CheckOutcome::UpToDate + ); + client.probe(&Version::new(1, 0, 0)).unwrap(); + assert_eq!(*client.transport.last_cookie_sent.borrow(), None); + } + } + + #[test] + fn eligible_cookie_is_saved_before_xml_parse_error() { + for lifetime in [CookieLifetime::PerUpdate, CookieLifetime::PerMessage] { + let (_dir, store) = cookie_store(); + let client = client_with("not XML").with_cookie_store(store.clone(), lifetime); + *client.transport.set_cookie.borrow_mut() = Some("ASLBSA=new".to_string()); + let err = client.check_for_update(&Version::new(1, 0, 0)).unwrap_err(); + assert!(matches!(err, NebraskaError::Parse(_)), "{err:?}"); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=new")); + } } #[test] diff --git a/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs index ed66963ec3..acbb080749 100644 --- a/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs +++ b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs @@ -6,26 +6,56 @@ //! `state.json`: that file exists to bridge one update *operation's* //! pre-/post-reboot halves (every field there is keyed by `operation_id` or //! is the single in-flight `pendingCommit`), and is read with -//! `deny_unknown_fields`. Nebraska's affinity cookie is a host/connection -//! -level attribute with no relationship to any particular operation - it -//! must outlive operations completing, failing, or being reset, and -//! vice versa - so it gets its own small file instead. +//! `deny_unknown_fields`. A separate cookie file preserves affinity across +//! process restarts and reboots without coupling its replacement policy to +//! the serialized update-state schema. use std::{ fs, io::{ErrorKind, Write}, os::unix::fs::OpenOptionsExt, path::{Path, PathBuf}, - process, + str::FromStr, }; -use anyhow::{Context, Error}; +use anyhow::{anyhow, Context, Error}; +use log::warn; +use reqwest::header::HeaderValue; +use uuid::Uuid; /// The cookie is a bearer-style session-affinity credential, so store it /// with owner-only permissions - the same posture `annotations::state` /// already applies to `state.json`. const COOKIE_FILE_MODE: u32 = 0o600; +/// Controls which responses may replace the persisted affinity cookie. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum CookieLifetime { + /// Start each initial update check without a cookie and retain only + /// that response's cookie until the next initial check. A restarted + /// initial check starts fresh; resuming completion does not reset it. + /// If the initial response sets none, later responses cannot establish one. + #[default] + PerUpdate, + /// Each response that supplies a cookie replaces the stored value for + /// the next request. Responses without a cookie leave it unchanged. + PerMessage, +} + +impl FromStr for CookieLifetime { + type Err = Error; + + fn from_str(s: &str) -> Result { + match s { + "per-update" => Ok(Self::PerUpdate), + "per-message" => Ok(Self::PerMessage), + other => Err(anyhow!( + "unknown cookie lifetime {other:?} (expected \"per-update\" or \"per-message\")" + )), + } + } +} + /// Durable store for the single Nebraska affinity cookie this host has most /// recently seen (e.g. an App-Service-style `ASLBSA`/`ASLBSACORS` pair, /// already reduced to a `Cookie:`-header-ready string by @@ -53,6 +83,7 @@ impl CookieStore { Ok(if trimmed.is_empty() { None } else { + HeaderValue::from_str(trimmed).context("invalid Nebraska cookie header")?; Some(trimmed.to_string()) }) } @@ -67,6 +98,7 @@ impl CookieStore { /// to a sibling temp file at owner-only (0600) permissions, fsync it, /// then atomically rename over the real path. pub fn save(&self, cookie: &str) -> Result<(), Error> { + HeaderValue::from_str(cookie).context("invalid Nebraska cookie header")?; let parent = match self.path.parent() { Some(parent) => { fs::create_dir_all(parent) @@ -82,58 +114,47 @@ impl CookieStore { .file_name() .and_then(|name| name.to_str()) .unwrap_or("nebraska-cookie"), - process::id() + Uuid::new_v4() )); - { - // create_new (not create+truncate) so the file is always newly - // created with COOKIE_FILE_MODE - see StateStore::save for why - // truncating a stale temp file left behind by a prior crash - // would silently keep that file's existing permissions. A - // leftover temp file (same process::id() reused across - // reboots) can only be this process's own abandoned write, - // never another process's live file, so it is removed and - // retried once. - let mut open_opts = fs::OpenOptions::new(); - open_opts - .write(true) - .create_new(true) - .mode(COOKIE_FILE_MODE); - let mut file = match open_opts.open(&temp_path) { - Ok(file) => file, - Err(err) if err.kind() == ErrorKind::AlreadyExists => { - fs::remove_file(&temp_path).with_context(|| { - format!("failed to remove stale {}", temp_path.display()) - })?; - open_opts - .open(&temp_path) - .with_context(|| format!("failed to create {}", temp_path.display()))? - } - Err(err) => { - return Err(err) - .with_context(|| format!("failed to create {}", temp_path.display())) - } - }; + // A per-write name prevents concurrent clients in this process + // from removing or renaming each other's live temporary file. + let mut file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .mode(COOKIE_FILE_MODE) + .open(&temp_path) + .with_context(|| format!("failed to create {}", temp_path.display()))?; + let result = (|| { file.write_all(cookie.as_bytes()) .with_context(|| format!("failed to write {}", temp_path.display()))?; file.sync_all() .with_context(|| format!("failed to fsync {}", temp_path.display()))?; - } + drop(file); - fs::rename(&temp_path, &self.path).with_context(|| { - format!( - "failed to atomically replace {} with {}", - self.path.display(), - temp_path.display() - ) - })?; + fs::rename(&temp_path, &self.path).with_context(|| { + format!( + "failed to atomically replace {} with {}", + self.path.display(), + temp_path.display() + ) + }) + })(); + if result.is_err() { + if let Err(err) = fs::remove_file(&temp_path) { + if err.kind() != ErrorKind::NotFound { + warn!("failed to remove Nebraska cookie temporary file: {err}"); + } + } + } + result?; // Best-effort: POSIX doesn't guarantee a rename is durable until the // containing directory's metadata is also synced, so fsync it too. // Not fatal if this fails (e.g. unsupported on some filesystems) - // the rename itself has already succeeded. - if let Ok(dir) = fs::File::open(parent) { - let _ = dir.sync_all(); + if let Err(err) = fs::File::open(parent).and_then(|dir| dir.sync_all()) { + warn!("Nebraska cookie was saved but directory sync failed: {err}"); } Ok(()) @@ -144,6 +165,11 @@ impl CookieStore { mod tests { use super::*; + use std::{ + sync::{Arc, Barrier}, + thread, + }; + fn store() -> (tempfile::TempDir, CookieStore) { let dir = tempfile::tempdir().expect("failed to create temp dir"); let path = dir.path().join("nebraska-cookie"); @@ -157,6 +183,68 @@ mod tests { assert_eq!(store.load().unwrap(), None); } + #[test] + fn empty_cookie_loads_as_absent() { + let (_dir, store) = store(); + store.save("").unwrap(); + assert_eq!(store.load().unwrap(), None); + } + + #[test] + fn invalid_cookie_headers_are_rejected_without_replacing_store() { + let (_dir, store) = store(); + store.save("ASLBSA=valid").unwrap(); + store.save("ASLBSA=bad\r\nInjected: value").unwrap_err(); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=valid")); + fs::write(store.path(), "ASLBSA=bad\r\nInjected: value").unwrap(); + store.load().unwrap_err(); + } + + #[test] + fn concurrent_saves_publish_whole_values_without_temporary_file_collisions() { + let (dir, store) = store(); + let values = Arc::new( + (0..8) + .map(|i| format!("ASLBSA={}", i.to_string().repeat(8192))) + .collect::>(), + ); + let barrier = Arc::new(Barrier::new(values.len())); + let threads = (0..values.len()) + .map(|index| { + let store = store.clone(); + let barrier = barrier.clone(); + let values = values.clone(); + thread::spawn(move || { + barrier.wait(); + for _ in 0..8 { + store.save(&values[index]).unwrap(); + assert!(values.contains(&store.load().unwrap().unwrap())); + } + }) + }) + .collect::>(); + for thread in threads { + thread.join().unwrap(); + } + assert!(values.contains(&store.load().unwrap().unwrap())); + assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1); + } + + #[test] + fn failed_rename_cleans_up_only_its_temporary_file() { + let dir = tempfile::tempdir().unwrap(); + let destination = dir.path().join("cookie"); + fs::create_dir(&destination).unwrap(); + let stale = dir.path().join("cookie.tmp-unrelated"); + fs::write(&stale, "keep this file").unwrap(); + CookieStore::new(&destination) + .save("ASLBSA=new") + .unwrap_err(); + assert!(destination.is_dir()); + assert_eq!(fs::read_to_string(stale).unwrap(), "keep this file"); + assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 2); + } + #[test] fn save_then_load_round_trips() { let (_dir, store) = store(); diff --git a/crates/trident-acl-agent/src/core/nebraska/mod.rs b/crates/trident-acl-agent/src/core/nebraska/mod.rs index 68b43a6efe..83ab1e36b4 100644 --- a/crates/trident-acl-agent/src/core/nebraska/mod.rs +++ b/crates/trident-acl-agent/src/core/nebraska/mod.rs @@ -85,7 +85,7 @@ mod transport; mod wire; pub use client::{redacted, CheckOutcome, Client, PackageFile, PackageHash, UpdateOffer}; -pub use cookie_store::CookieStore; +pub use cookie_store::{CookieLifetime, CookieStore}; pub use error::NebraskaError; pub use event::ProgressEvent; pub use id::MachineId; diff --git a/crates/trident-acl-agent/src/lib.rs b/crates/trident-acl-agent/src/lib.rs index f08cd78960..d5f7be079e 100644 --- a/crates/trident-acl-agent/src/lib.rs +++ b/crates/trident-acl-agent/src/lib.rs @@ -14,11 +14,11 @@ //! both by this crate's agent binary and by a future Trident ACL Agent that //! orchestrates updates differently. -use log::warn; use semver::Version; use url::Url; use crate::core::{ + config::CookieLifetime, error::AgentError, nebraska::{Client, CookieStore, MachineId, NebraskaError}, version::FALLBACK_ALWAYS_VERSION, @@ -78,22 +78,10 @@ fn build_machine_id(source: IdSource) -> Result { /// (hence this being named `probe`, not `ping`). See [`Client::probe`]'s doc /// comment for the full explanation. /// -/// Also loads/saves `cookie_store`: this probe talks to the very same -/// Nebraska instance (and, in production, the same load-balancer-fronted -/// deployment) that `check_for_update`/`report_progress`/etc. do, so it -/// must participate in the shared affinity-cookie dance too -- otherwise -/// every invocation of `--validate-connection nebraska` is itself a fresh -/// "first contact" that never adopts or persists a cookie, silently -/// defeating the whole mechanism's point for one more request each time -/// it's run. Best-effort, like every other cookie_store use: a load/save -/// failure only logs a warning and the probe proceeds/returns normally. -/// -/// `cookie_enabled` is the same blanket kill switch -/// (`OrchestrationConfig::cookie_enabled`) the orchestrator honors: when -/// `false`, this probe never touches `cookie_store` at all, in either -/// direction. There's no `CookieLifetime` concept here -- a diagnostic -/// probe isn't part of any update, so "first call of a new update" never -/// applies to it. +/// When enabled, probes reuse the shared cookie store. Per-update probes +/// cannot initialize or replace the update's cookie; per-message probes +/// persist response cookies. Storage failures only warn, and disabled +/// probes never access the store. pub fn check_nebraska_reachable( url: &Url, app_id: &str, @@ -101,37 +89,18 @@ pub fn check_nebraska_reachable( machine_id_source: IdSource, cookie_store: &CookieStore, cookie_enabled: bool, + cookie_lifetime: CookieLifetime, ) -> Result<(), AgentError> { let machine_id = build_machine_id(machine_id_source)?; - let mut client = - Client::new(url.clone(), app_id, track, machine_id).with_cookie_enabled(cookie_enabled); - let initial_cookie = cookie_enabled - .then(|| match cookie_store.load() { - Ok(cookie) => cookie, - Err(err) => { - warn!("failed to load Nebraska cookie, proceeding without one: {err}"); - None - } - }) - .flatten(); - if let Some(cookie) = initial_cookie { - client = client.with_cookie(cookie); + let mut client = Client::new(url.clone(), app_id, track, machine_id); + if cookie_enabled { + client = client.with_cookie_store(cookie_store.clone(), cookie_lifetime); } - let result = client.probe( + match client.probe( &Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver"), - ); - - if cookie_enabled { - if let Some(cookie) = client.cookie() { - if let Err(err) = cookie_store.save(&cookie) { - warn!("failed to persist Nebraska cookie: {err}"); - } - } - } - - match result { + ) { Ok(()) => Ok(()), // A well-formed response reporting a non-OK app status still proves // the server is reachable and speaking Omaha; only a @@ -199,6 +168,7 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, true, + CookieLifetime::PerMessage, ) .unwrap(); @@ -232,6 +202,7 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, true, + CookieLifetime::PerUpdate, ) .unwrap(); @@ -249,6 +220,7 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, true, + CookieLifetime::PerUpdate, ) .unwrap_err(); assert!(matches!(err, AgentError::Nebraska(_))); @@ -284,6 +256,7 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, true, + CookieLifetime::PerMessage, ) .unwrap(); first_mock.assert(); @@ -313,11 +286,69 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, true, + CookieLifetime::PerMessage, ) .unwrap(); second_mock.assert(); } + #[test] + fn per_update_probe_reuses_but_never_establishes_or_replaces_cookie() { + let mut server = Server::new(); + let (_dir, store) = temp_cookie_store(); + for cookie in [None, Some("ASLBSA=initial")] { + if let Some(cookie) = cookie { + store.save(cookie).unwrap(); + } + let mock = server.mock("POST", "/") + .match_header("cookie", cookie.map(|v| Matcher::Exact(v.to_string())).unwrap_or(Matcher::Missing)) + .with_status(200) + .with_header("set-cookie", "ASLBSA=ignored; Path=/") + .with_body(r#""#) + .expect(1) + .create(); + check_nebraska_reachable( + &Url::parse(&server.url()).unwrap(), + "test", + "track", + IdSource::MachineIdHashed, + &store, + true, + CookieLifetime::PerUpdate, + ) + .unwrap(); + mock.assert(); + assert_eq!(store.load().unwrap().as_deref(), cookie); + } + } + + #[test] + fn http_error_cookie_is_not_adopted() { + let mut server = Server::new(); + let (_dir, store) = temp_cookie_store(); + store.save("ASLBSA=previous").unwrap(); + let mock = server + .mock("POST", "/") + .match_header("cookie", "ASLBSA=previous") + .with_status(503) + .with_header("set-cookie", "ASLBSA=error-cookie; Path=/") + .expect(1) + .create(); + let err = check_nebraska_reachable( + &Url::parse(&server.url()).unwrap(), + "test", + "track", + IdSource::MachineIdHashed, + &store, + true, + CookieLifetime::PerMessage, + ) + .unwrap_err(); + assert!(matches!(err, AgentError::Nebraska(_)), "{err:?}"); + mock.assert(); + assert_eq!(store.load().unwrap().as_deref(), Some("ASLBSA=previous")); + } + #[test] fn test_check_nebraska_reachable_cookie_enabled_false_never_touches_store() { // The blanket kill switch: with cookie_enabled=false, no Cookie @@ -349,6 +380,7 @@ mod tests { IdSource::MachineIdHashed, &cookie_store, false, + CookieLifetime::PerUpdate, ) .unwrap(); diff --git a/crates/trident-acl-agent/src/omahaonly/mod.rs b/crates/trident-acl-agent/src/omahaonly/mod.rs index 04d8c0e937..07a03ae336 100644 --- a/crates/trident-acl-agent/src/omahaonly/mod.rs +++ b/crates/trident-acl-agent/src/omahaonly/mod.rs @@ -11,7 +11,7 @@ use crate::{ core::{ config::AgentConfig, cosi_sha384, - nebraska::{CheckOutcome, Client}, + nebraska::{CheckOutcome, Client, CookieStore}, trident::TridentClient, version::{self, FALLBACK_ALWAYS_VERSION}, }, @@ -43,8 +43,14 @@ pub async fn run_omaha_only(config: &AgentConfig) -> Result<(), Error> { ); Version::parse(FALLBACK_ALWAYS_VERSION).expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); + let cookie_enabled = config.orchestration.cookie_enabled; + let cookie_store = CookieStore::new(config.orchestration.cookie_path.clone()); + let cookie_lifetime = config.orchestration.cookie_lifetime; let outcome = task::spawn_blocking(move || { - let client = Client::new(endpoint, app_id, track, machine_id); + let mut client = Client::new(endpoint, app_id, track, machine_id); + if cookie_enabled { + client = client.with_cookie_store(cookie_store, cookie_lifetime); + } client.check_for_update(¤t_version) }) .await diff --git a/tools/storm/aclagent/proxies/nebraska.go b/tools/storm/aclagent/proxies/nebraska.go index d5029221d2..a476b8859f 100644 --- a/tools/storm/aclagent/proxies/nebraska.go +++ b/tools/storm/aclagent/proxies/nebraska.go @@ -215,10 +215,8 @@ func (p *NebraskaProxy) AffinityCookie() string { } // RequestsMissingCookie returns how many requests this proxy has served -// with no Cookie header at all - i.e. before trident-acl-agent has ever -// picked up the simulated affinity cookie. A healthy end-to-end run expects -// exactly 1: the agent's very first-ever contact with this Nebraska -// instance. +// with no Cookie header at all, including diagnostic probes before affinity +// is established and the initial check of each per-update staging attempt. func (p *NebraskaProxy) RequestsMissingCookie() int { p.cookieMu.Lock() defer p.cookieMu.Unlock() diff --git a/tools/storm/aclagent/tests/update.go b/tools/storm/aclagent/tests/update.go index 5e756aa9dd..e0c05e352d 100644 --- a/tools/storm/aclagent/tests/update.go +++ b/tools/storm/aclagent/tests/update.go @@ -469,9 +469,11 @@ func RunABUpdate(testConfig stormaclconfig.TestConfig, vmConfig stormvmconfig.Al // best-effort hint, not something Nebraska requires to drive the update // correctly, so these checks run after - not instead of - the real // pass/fail assertions above. - if got := nebraska.RequestsMissingCookie(); got != 1 { + // The successful diagnostic probe cannot establish per-update affinity; + // the initial staging check must independently start without a cookie. + if got := nebraska.RequestsMissingCookie(); got != 2 { collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) - return fmt.Errorf("expected exactly 1 Nebraska request with no affinity cookie (the agent's first-ever contact), got %d", got) + return fmt.Errorf("expected exactly 2 Nebraska requests with no affinity cookie (diagnostic probe and initial staging check), got %d", got) } if got := nebraska.RequestsWithWrongCookie(); got != 0 { collectAclArtifactsBestEffort(vmConfig.VMConfig, vmIP, testConfig.OutputPath) From eb873915c59489f855a281bb0d37be7cda8f6361 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Thu, 22 Oct 2026 20:27:38 +0000 Subject: [PATCH 09/13] acl: document only supported cookie policies Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5638630c-63c7-440d-b89f-d5add6423071 --- crates/trident-acl-agent/src/core/config.rs | 16 +++------------- .../src/core/nebraska/README.md | 2 +- 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index b890b95433..50ddafa229 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -921,28 +921,18 @@ mod tests { } #[test] - fn removed_forever_cookie_lifetime_is_rejected() { + fn malformed_cookie_lifetime_is_a_parse_error() { let err = AgentConfig::from_vars(vars(&[( "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", - "forever", + "bogus", )])) .unwrap_err(); let message = format!("{err:#}"); - assert!(message.contains("forever"), "{message}"); + assert!(message.contains("bogus"), "{message}"); assert!(message.contains("per-update"), "{message}"); assert!(message.contains("per-message"), "{message}"); } - #[test] - fn malformed_cookie_lifetime_is_a_parse_error() { - let err = AgentConfig::from_vars(vars(&[( - "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", - "bogus", - )])) - .unwrap_err(); - assert!(format!("{err:#}").contains("bogus"), "{err:#}"); - } - #[test] fn malformed_duration_is_a_parse_error() { let err = AgentConfig::from_vars(vars(&[( diff --git a/crates/trident-acl-agent/src/core/nebraska/README.md b/crates/trident-acl-agent/src/core/nebraska/README.md index d8e92f7109..ff474da128 100644 --- a/crates/trident-acl-agent/src/core/nebraska/README.md +++ b/crates/trident-acl-agent/src/core/nebraska/README.md @@ -67,7 +67,7 @@ IDs; resuming finalize/completion does not. `per-message` sends the stored cookie on every request and persists each response's replacement for the next request, including completion retries. -Responses without `Set-Cookie` leave the file unchanged. `forever` is not supported. +Responses without `Set-Cookie` leave the file unchanged. Standalone clients opt in with `Client::with_cookie_store(store, lifetime)`, using `CookieStore` and `CookieLifetime` from this module. A client without a store From b064a2ed08a646c1a4f0778b19ef96f3c103597f Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 23 Oct 2026 17:17:38 +0000 Subject: [PATCH 10/13] acl: import cookie header constants directly Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0eeaf50a-88e0-4d3b-af9d-9a3bb53fd7c6 --- crates/trident-acl-agent/src/core/nebraska/transport.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/trident-acl-agent/src/core/nebraska/transport.rs b/crates/trident-acl-agent/src/core/nebraska/transport.rs index 241a3ad683..5efd33dd3a 100644 --- a/crates/trident-acl-agent/src/core/nebraska/transport.rs +++ b/crates/trident-acl-agent/src/core/nebraska/transport.rs @@ -24,6 +24,7 @@ use std::time::{Duration, Instant}; use log::trace; +use reqwest::header::{COOKIE, SET_COOKIE}; use url::Url; use super::{client::redacted, error::NebraskaError}; @@ -114,7 +115,7 @@ impl Transport for ReqwestTransport { .header("Content-Type", "application/xml") .timeout(self.timeout); if let Some(cookie) = cookie { - request = request.header(reqwest::header::COOKIE, cookie); + request = request.header(COOKIE, cookie); } // `reqwest::Error`'s `Display` embeds the request URL, which may carry an @@ -178,7 +179,7 @@ impl Transport for ReqwestTransport { fn extract_cookie(response: &reqwest::blocking::Response) -> Option { let pairs: Vec = response .headers() - .get_all(reqwest::header::SET_COOKIE) + .get_all(SET_COOKIE) .iter() .filter_map(|value| value.to_str().ok()) .filter_map(|raw| raw.split(';').next()) From 43192d65621d49ecc249d842dd9647ef85649b5d Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 23 Oct 2026 17:20:54 +0000 Subject: [PATCH 11/13] acl: derive cookie lifetime parsing with strum Keep configuration errors descriptive and reject noncanonical policies. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0eeaf50a-88e0-4d3b-af9d-9a3bb53fd7c6 --- Cargo.lock | 2 ++ crates/trident-acl-agent/Cargo.toml | 2 ++ crates/trident-acl-agent/src/core/config.rs | 28 ++++++++++++------- .../src/core/nebraska/cookie_store.rs | 21 +++----------- 4 files changed, 26 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 59f83d90a9..eb8f095ef0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4047,6 +4047,8 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_yaml", + "strum 0.26.3", + "strum_macros 0.26.4", "sysdefs", "systemd-journal-logger", "tempfile", diff --git a/crates/trident-acl-agent/Cargo.toml b/crates/trident-acl-agent/Cargo.toml index 31340a2574..e5cc64e98c 100644 --- a/crates/trident-acl-agent/Cargo.toml +++ b/crates/trident-acl-agent/Cargo.toml @@ -28,6 +28,8 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } serde_yaml = { workspace = true } serde_path_to_error = { workspace = true } +strum = { workspace = true } +strum_macros = { workspace = true } systemd-journal-logger = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } diff --git a/crates/trident-acl-agent/src/core/config.rs b/crates/trident-acl-agent/src/core/config.rs index d396b7c9e8..5f88d25f2b 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -429,7 +429,13 @@ where D: Deserializer<'de>, { empty_as_none(String::deserialize(deserializer)?) - .map(|value| value.parse::().map_err(D::Error::custom)) + .map(|value| { + value.parse::().map_err(|_| { + D::Error::custom(format!( + "unknown cookie lifetime {value:?} (expected \"per-update\" or \"per-message\")" + )) + }) + }) .transpose() } @@ -1161,15 +1167,17 @@ mod tests { #[test] fn malformed_cookie_lifetime_is_a_parse_error() { - let err = AgentConfig::from_vars(vars(&[( - "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", - "bogus", - )])) - .unwrap_err(); - let message = format!("{err:#}"); - assert!(message.contains("bogus"), "{message}"); - assert!(message.contains("per-update"), "{message}"); - assert!(message.contains("per-message"), "{message}"); + for value in ["bogus", "forever", "PerUpdate", "per_update", "PER-MESSAGE"] { + let err = AgentConfig::from_vars(vars(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + value, + )])) + .unwrap_err(); + let message = format!("{err:#}"); + assert!(message.contains(value), "{message}"); + assert!(message.contains("per-update"), "{message}"); + assert!(message.contains("per-message"), "{message}"); + } } #[test] diff --git a/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs index acbb080749..b8318ae082 100644 --- a/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs +++ b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs @@ -15,12 +15,12 @@ use std::{ io::{ErrorKind, Write}, os::unix::fs::OpenOptionsExt, path::{Path, PathBuf}, - str::FromStr, }; -use anyhow::{anyhow, Context, Error}; +use anyhow::{Context, Error}; use log::warn; use reqwest::header::HeaderValue; +use strum_macros::EnumString; use uuid::Uuid; /// The cookie is a bearer-style session-affinity credential, so store it @@ -29,7 +29,8 @@ use uuid::Uuid; const COOKIE_FILE_MODE: u32 = 0o600; /// Controls which responses may replace the persisted affinity cookie. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, EnumString)] +#[strum(serialize_all = "kebab-case")] pub enum CookieLifetime { /// Start each initial update check without a cookie and retain only /// that response's cookie until the next initial check. A restarted @@ -42,20 +43,6 @@ pub enum CookieLifetime { PerMessage, } -impl FromStr for CookieLifetime { - type Err = Error; - - fn from_str(s: &str) -> Result { - match s { - "per-update" => Ok(Self::PerUpdate), - "per-message" => Ok(Self::PerMessage), - other => Err(anyhow!( - "unknown cookie lifetime {other:?} (expected \"per-update\" or \"per-message\")" - )), - } - } -} - /// Durable store for the single Nebraska affinity cookie this host has most /// recently seen (e.g. an App-Service-style `ASLBSA`/`ASLBSACORS` pair, /// already reduced to a `Cookie:`-header-ready string by From c5966f1e8ceeac42abadba5820544db3f02830dd Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 23 Oct 2026 17:49:22 +0000 Subject: [PATCH 12/13] docs: keep cookie lifetime configuration internal Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0eeaf50a-88e0-4d3b-af9d-9a3bb53fd7c6 --- .../src/core/nebraska/README.md | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/crates/trident-acl-agent/src/core/nebraska/README.md b/crates/trident-acl-agent/src/core/nebraska/README.md index ff474da128..264c3d90d7 100644 --- a/crates/trident-acl-agent/src/core/nebraska/README.md +++ b/crates/trident-acl-agent/src/core/nebraska/README.md @@ -56,30 +56,22 @@ Cookies are opt-in. The agent uses these environment variables: | --- | --- | | `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED` | `false` | | `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH` | `/var/lib/trident-acl-agent/nebraska-cookie` | -| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME` | `per-update` | -`per-update` sends each initial `check_for_update` without a cookie, clears the -previous stored value, and persists only that response's cookie. Progress, +Each initial `check_for_update` starts without a cookie, clears the previous +stored value, and persists only that response's cookie. Progress, failure, completion, and diagnostic requests reuse it without adopting replacements. An initial response without a cookie leaves the update cookie-less. A restarted staging attempt starts a new interval, even with unchanged request IDs; resuming finalize/completion does not. -`per-message` sends the stored cookie on every request and persists each -response's replacement for the next request, including completion retries. -Responses without `Set-Cookie` leave the file unchanged. - -Standalone clients opt in with `Client::with_cookie_store(store, lifetime)`, using -`CookieStore` and `CookieLifetime` from this module. A client without a store -neither sends nor adopts cookies. Both agent orchestration modes and connectivity -probes honor the enabled switch and policy. +When cookies are disabled, clients neither send nor adopt them. Both agent +orchestration modes and connectivity probes honor the enabled switch. Only HTTP-success responses with readable bodies can supply a cookie; XML validation happens afterward. Cookies are stored with owner-only permissions and atomic replacement, never logged. I/O failures warn without failing the update; the affinity guarantees are therefore best-effort when storage fails. Concurrent -cookie-bearing responses use last-successful-write-wins ordering. Configure -independent processes consistently when sharing a cookie path. +cookie-bearing responses use last-successful-write-wins ordering. ### Poll for an update From d8b6ca41c07aeaa5814b24fc98664ca273babcc3 Mon Sep 17 00:00:00 2001 From: Brian Fjeldstad Date: Fri, 23 Oct 2026 17:58:31 +0000 Subject: [PATCH 13/13] docs: align ACL agent guide with cookie support and main Document cookie enable/path settings and default affinity behavior without exposing lifetime configuration. Correct staging, diagnostics, timeout, and deployment details against the implementation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0eeaf50a-88e0-4d3b-af9d-9a3bb53fd7c6 --- docs/Explanation/Trident-ACL-Agent.md | 80 ++++++++++++++++++--------- 1 file changed, 55 insertions(+), 25 deletions(-) diff --git a/docs/Explanation/Trident-ACL-Agent.md b/docs/Explanation/Trident-ACL-Agent.md index a3f95c25c4..958e29db49 100644 --- a/docs/Explanation/Trident-ACL-Agent.md +++ b/docs/Explanation/Trident-ACL-Agent.md @@ -7,9 +7,9 @@ half of Azure Container Linux (ACL)'s update mechanism. Any Kubernetes control-plane component (a custom controller, an operator, or a script driven by `kubectl patch`) can orchestrate updates across a fleet of nodes by writing to the annotation -contract described below, provided it is willing to speak the -[Omaha](https://github.com/omaha-consortium/omaha) protocol for image -distribution and honors the agent's per-node protocol. +contract described below. It supplies the image source on a +[Nebraska](https://github.com/flatcar/nebraska) server and honors the agent's +per-node protocol; the agent speaks Omaha to Nebraska for image distribution. ## Deployment @@ -29,11 +29,12 @@ files. Installing the package does not by itself enable or start the service — a deployment decides when that happens, e.g. by running `systemctl enable --now trident-acl-agent.service` on the node, or by baking that enablement into the image build (as this repo's own -`updateimg-acl-agent.yaml` test image does via Image Customizer's -`services: enable` list). +`tests/images/trident-vm-testimage/base/updateimg-acl-agent.yaml` does via +Image Customizer's `services: enable` list). -The shipped unit carries no `Environment=` lines beyond `ExecStart`, so -every deployment-specific choice — which annotation prefix to watch, where +The shipped unit invokes the agent via `ExecStart` and has no +`Environment=` lines, so every deployment-specific choice — which annotation +prefix to watch, where to read the current version from, which Kubernetes API server to talk to, and so on — is supplied the same way any other systemd service is configured: standard `Environment=`/`EnvironmentFile=` constructs, most @@ -102,7 +103,7 @@ A request annotation looks like: | `operation` | Trident invocation | Effect | |---|---|---| -| `stage` | `trident update --allowed-operations=stage` | Queries the `server`/`appId`/`track` Omaha endpoint for `targetVersion`, then streams the resulting image to the inactive partition. No reboot. | +| `stage` | `trident update --allowed-operations=stage` | Queries the `server`/`appId`/`track` Omaha endpoint with the node's current version, requires the offered version to match `targetVersion`, then streams the resulting image to the inactive partition. No reboot. | | `finalize` | `trident update --allowed-operations=finalize` (gRPC `UpdateFinalize`, caller-handled reboot) | Arms boot for the staged target, writes a terminal `finalize` status, then triggers the reboot. | | `rollback` | `trident rollback --ab` (gRPC `RollbackStage`/`RollbackFinalize`, caller-handled reboot) | Swaps back to the previous partition, mirroring `finalize` on the return path. Only the last update can be undone this way. | @@ -123,9 +124,9 @@ sequenceDiagram participant Trident Orchestrator->>API: 1. PATCH request: stage (opId A) - Note over Agent: agent picks up the request
on its next poll + Note over Agent: agent picks up the request
through its Node watch Agent->>API: 2. read request, PATCH status: stage InProgress - Agent->>Nebraska: 3. query targetVersion (server/appId/track) + Agent->>Nebraska: 3. check for update using current version (server/appId/track) Nebraska-->>Agent: image location Agent->>Trident: 4. Stage (image to inactive partition) Trident-->>Agent: staged | error @@ -223,8 +224,8 @@ A status annotation (`/update-status` or | `InvalidRequest` | Yes | The request annotation itself was rejected before any action was taken — malformed JSON, a schema/version mismatch, a missing required field (`server`/`appId`/`track`/`targetVersion`), a `finalize` whose `targetVersion` doesn't match what was staged, or a second `finalize`/`rollback` submitted while one is already pending its post-reboot `commit`. No Trident operation runs. | See the request/status schema types in -`crates/trident-acl-agent/src/annotations/protocol.rs` for the full contract, -including the formal JSON Schema both sides validate against. +`crates/trident-acl-agent/src/annotations/protocol.rs` for the validation +rules and the formal JSON Schemas pinned by the protocol conformance tests. ## Pre/post-reboot state and the watchdog @@ -234,18 +235,40 @@ commit is pending and a marker for "a boot happened after this point". On restart, the agent checks this state to resume the post-reboot `commit` step rather than re-running `finalize` from scratch. -While an operation is in flight, the agent refreshes the `InProgress` -status's `lastUpdatedUtc` on a heartbeat cadence +While a stage, finalize, or rollback Trident call is in flight, the agent +refreshes the `InProgress` status's `lastUpdatedUtc` on a heartbeat cadence (`TRIDENT_ACL_AGENT_ORCHESTRATION_HEARTBEAT_INTERVAL`), so an external watchdog can distinguish a working agent from a stuck one and reprovision a node that never reports a terminal `commit` status within its SLA. +## Nebraska affinity cookies + +Cookie handling is opt-in via +`TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED=true`. When disabled (the +default), the agent neither sends nor adopts cookies and leaves any existing +cookie file untouched. + +When enabled, each initial staging update check starts without a cookie and +clears the previous stored value. The cookie returned by that check is +persisted separately from `state.json` and reused for progress, failure, +finalize, and post-reboot completion requests. Later responses do not replace +it. If the initial response sets no cookie, subsequent requests remain +cookie-less. A restarted staging attempt starts fresh, even with unchanged +request IDs; resuming finalize or completion retains the stored cookie. + +The agent reads the file for each request. Cookies are saved only from +successful HTTP responses with readable bodies, before Omaha XML validation. +The file uses owner-only permissions and atomic replacement; cookie values +are never logged. Cookie I/O failures warn without failing the update, so +affinity is best-effort if storage fails. + ## Configuration At startup, the agent logs its effective `AgentConfig` at INFO after applying environment overrides and defaults, including for `--validate-connection`. Network URL credentials, paths, query strings, and fragments are redacted. -Kubeconfig and state paths are logged, but their contents are not read for logging. +Kubeconfig, state, and cookie file paths are logged, but their contents are +not read for logging. There is no config file. Every setting is an environment variable prefixed `TRIDENT_ACL_AGENT_`, systemd-style: set it in the unit's own @@ -254,7 +277,7 @@ sets the process's environment before it starts. A variable that is unset, or set to the empty string, falls back to its default. A variable set to a malformed value (a bad URL, a bad duration, an -unrecognized `mode`) causes the agent to fail to start with an error +invalid boolean) causes the agent to fail to start with an error naming the offending variable. | Variable | Default | Description | @@ -263,7 +286,7 @@ naming the offending variable. | `TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX` | `acl.microsoft.com` | The annotation-key prefix for the request/status/commit-status annotations (e.g. the `acl.microsoft.com` in `acl.microsoft.com/update-request`). Any orchestrator can pick its own namespace here so its annotations don't collide with another controller's. | | `TRIDENT_ACL_AGENT_CURRENT_VERSION_PATH` | `/etc/os-release` | The file the agent reads to determine the node's currently running version. Any file works, as long as it follows the os-release format (`KEY=VALUE` lines, optionally single- or double-quoted, blank lines and `#` comments ignored) — see [below](#configuring-the-on-disk-version). | | `TRIDENT_ACL_AGENT_CURRENT_VERSION_KEY` | `VERSION_ID` | The key the agent looks up in `TRIDENT_ACL_AGENT_CURRENT_VERSION_PATH` (`/etc/os-release` by default) to determine the node's currently running version, used to compare against a request's `targetVersion` (e.g. to short-circuit to `AlreadyAtTarget`). `VERSION_ID` is the standard `os-release` field most images already stamp; a deployment that instead carries an ACL-specific `IMAGE_VERSION` field can point this variable at that key instead — see [below](#configuring-the-on-disk-version). | -| `TRIDENT_ACL_AGENT_CURRENT_VERSION_FALLBACK` | `always` | Controls what happens when `TRIDENT_ACL_AGENT_CURRENT_VERSION_KEY` isn't present at `TRIDENT_ACL_AGENT_CURRENT_VERSION_PATH` (e.g. a dev/test host, or an image that hasn't started stamping that key yet). `always` reports `0.0.0` as the node's current version — a sentinel that can never collide with a real release version and cause a false `AlreadyAtTarget`. `error` fails the operation instead of using a placeholder version. Any other value is used verbatim as the current version, with no format validation. | +| `TRIDENT_ACL_AGENT_CURRENT_VERSION_FALLBACK` | `always` | Controls what happens when the configured version file is unreadable, the configured key is missing, or its value is empty. `always` reports `0.0.0` as the node's current version. `error` fails the operation instead of using a placeholder version. Any other value is used verbatim as the current version, with no format validation. | | `TRIDENT_ACL_AGENT_KUBERNETES_API_SERVER` | unset | Explicit override for the Kubernetes API server URL. When unset, the server embedded in the kubeconfig is used as-is. | | `TRIDENT_ACL_AGENT_KUBERNETES_KUBECONFIG` | `/var/lib/kubelet/kubeconfig` | Path to the kubeconfig used to reach the Kubernetes API server and authenticate as this node. | | `TRIDENT_ACL_AGENT_KUBERNETES_NODE_NAME` | derived from `TRIDENT_ACL_AGENT_KUBERNETES_NODE_NAME_SOURCE` | The Node object this agent watches/patches. | @@ -271,8 +294,10 @@ naming the offending variable. | `TRIDENT_ACL_AGENT_KUBERNETES_KUBELET_CERT_PATH` | `/var/lib/kubelet/pki/kubelet-client-current.pem` | Path to kubelet's own client certificate, consulted when `TRIDENT_ACL_AGENT_KUBERNETES_NODE_NAME_SOURCE=kubelet-cert`. Override when kubelet's cert lives somewhere other than the standard path (e.g. a non-default kubelet `--cert-dir`). Ignored when `node_name_source` is `hostname` or `TRIDENT_ACL_AGENT_KUBERNETES_NODE_NAME` is set. | | `TRIDENT_ACL_AGENT_TRIDENT_SOCKET` | `unix:///run/trident/trident.sock` | The gRPC Unix socket URI used to reach `tridentd`. | | `TRIDENT_ACL_AGENT_ORCHESTRATION_STATE_PATH` | `/var/lib/trident-acl-agent/state.json` | Path to the agent's persistent state file bridging the pre-reboot and post-reboot halves of `finalize`/`rollback` across the reboot. | -| `TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT` | `20m` | How long a `stage` is allowed to run (a [`humantime`](https://docs.rs/humantime) duration, e.g. `20m`, `1h`) before it's considered failed. | -| `TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT` | `10m` | How long a `finalize` is allowed to run before it's considered failed. | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED` | `false` | Enable Nebraska affinity cookie transmission and persistence. When disabled, existing cookie files are left untouched. See [Nebraska affinity cookies](#nebraska-affinity-cookies). | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_PATH` | `/var/lib/trident-acl-agent/nebraska-cookie` | Path to the persisted Nebraska affinity cookie, separate from the agent's state file. | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT` | `20m` | Timeout for each Trident update-stage or rollback-stage call (a [`humantime`](https://docs.rs/humantime) duration, e.g. `20m`, `1h`), not the entire staging workflow. | +| `TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT` | `10m` | Timeout for each Trident update-finalize, rollback-finalize, or post-reboot commit call. | | `TRIDENT_ACL_AGENT_ORCHESTRATION_HEARTBEAT_INTERVAL` | `60s` | Refresh cadence for the `InProgress` status heartbeat. | | `TRIDENT_ACL_AGENT_ORCHESTRATION_NODE_GONE_MAX_WAIT` | unset (wait forever) | How long the agent will keep waiting for its own Node object to reappear after a 404 before giving up and exiting (a [`humantime`](https://docs.rs/humantime) duration, e.g. `30m`, `1h`). Unset means the agent never gives up on its own — see below. | | `TRIDENT_ACL_AGENT_KUBERNETES_VALIDATE_NODE_UUID` | `false` | When `true`, every Node read verifies `status.nodeInfo.systemUUID` against this machine's own `/sys/class/dmi/id/product_uuid`, treating a mismatch the same as the Node not existing (`NodeGone`) — see below. | @@ -329,8 +354,8 @@ deployment. ### Setting env vars via a systemd drop-in -The agent ships as `trident-acl-agent.service`, with no `Environment=` -lines of its own beyond `ExecStart`. Any setting is overridden with a +The agent ships as `trident-acl-agent.service`, with an `ExecStart` command +but no `Environment=` lines of its own. Any setting is overridden with a drop-in file, without editing the packaged unit: ```console @@ -383,11 +408,10 @@ compares it against a request's `targetVersion` the same way it would for `VERSION_ID`/`/etc/os-release` — including short-circuiting to `AlreadyAtTarget` when they already match. -If the configured key is absent from the configured file (for example, on -a dev/test host with a minimal `os-release`), the agent consults +If the configured file is unreadable, the key is missing, or its value is +empty, the agent consults `TRIDENT_ACL_AGENT_CURRENT_VERSION_FALLBACK` (`always` by default): `always` -reports `0.0.0` as the current version — a sentinel that can never -accidentally match a real requested version; `error` fails the operation +reports `0.0.0` as the current version; `error` fails the operation instead of guessing; any other value is used verbatim as the current version, unvalidated. @@ -398,6 +422,12 @@ checks connectivity to a single dependency using the current environment and exits immediately — useful for a systemd `ExecStartPre` check or manual on-node troubleshooting without running the full orchestrator loop. +The Nebraska probe honors the cookie enable/path settings. When enabled, +it reuses an existing persisted cookie without clearing or replacing it; +a probe cannot establish affinity. A standalone invocation does not inherit +a systemd service's environment, so supply those settings explicitly when +checking cookie reuse. + `--validate-connection nebraska` is the one place `TRIDENT_ACL_AGENT_NEBRASKA_ENDPOINT`, `TRIDENT_ACL_AGENT_NEBRASKA_APP_ID`, and `TRIDENT_ACL_AGENT_NEBRASKA_TRACK` are used: it sends a non-mutating