diff --git a/Cargo.lock b/Cargo.lock index 59f83d90a..eb8f095ef 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 31340a257..e5cc64e98 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/annotations/orchestrator.rs b/crates/trident-acl-agent/src/annotations/orchestrator.rs index 65dfafc6b..82e4875aa 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, KubernetesConfig}, 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}, }, @@ -114,6 +114,10 @@ pub struct Orchestrator { node_name: String, 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, } @@ -140,6 +144,7 @@ impl Orchestrator { let annotation_keys = AnnotationKeys::new(&config.kubernetes.annotation_prefix); Ok(Some(Self { state: StateStore::new(config.orchestration.state_path.clone()), + nebraska_cookie: CookieStore::new(config.orchestration.cookie_path.clone()), config, node_name, k8s, @@ -761,10 +766,16 @@ impl Orchestrator { Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver") }); + let cookie_enabled = self.config.orchestration.cookie_enabled; + 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 client = NebraskaClient::new(endpoint, app_id, track, machine_id); + let mut client = NebraskaClient::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) } }) @@ -1684,8 +1695,14 @@ impl Orchestrator { } }; let label = report.label(); + 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 client = NebraskaClient::new(endpoint, app_id, track, machine_id); + let mut client = NebraskaClient::new(endpoint, app_id, track, machine_id); + if cookie_enabled { + client = client.with_cookie_store(cookie_store, cookie_lifetime); + } match report { NebraskaReport::Progress { version, event } => { client.report_progress(&version, event) @@ -2318,11 +2335,254 @@ mod tests { use crate::{ annotations::{RequestedOperation, SCHEMA_VERSION}, core::{ - config::NodeNameSource, + config::{CookieLifetime, NodeNameSource}, 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(); + let 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() + }; + let shutdown = CancellationToken::new(); + Orchestrator::from_config(config, &shutdown) + .await + .unwrap() + .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::PerMessage), + (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::PerMessage { + 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(""#; + let dir = tempfile::tempdir().unwrap(); + let orchestrator = cookie_orchestrator(dir.path(), true, CookieLifetime::PerUpdate).await; + 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; + assert_eq!( + restarted.nebraska_cookie.load().unwrap().as_deref(), + Some("ASLBSA=first-attempt"), + ); + 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.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"), + ); + } + + #[tokio::test] + async fn disabled_cookie_is_not_adopted_by_completion_report_retries() { + let dir = tempfile::tempdir().unwrap(); + let orchestrator = cookie_orchestrator(dir.path(), false, CookieLifetime::PerUpdate).await; + orchestrator + .nebraska_cookie + .save("ASLBSA=previous") + .unwrap(); + let mut server = mockito::Server::new_async().await; + let first = server + .mock("POST", "/v1/update/") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_header("set-cookie", "ASLBSA=fresh; Path=/") + .with_body(r#""#) + .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"), + ); + } + // --- node_name resolution retry (microsoft/trident#839 follow-up) --- /// Proves resolve_node_name_with_retry actually *retries* (keeps diff --git a/crates/trident-acl-agent/src/connection_check.rs b/crates/trident-acl-agent/src/connection_check.rs index dc21e7bd0..65dfaf223 100644 --- a/crates/trident-acl-agent/src/connection_check.rs +++ b/crates/trident-acl-agent/src/connection_check.rs @@ -80,14 +80,22 @@ 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. + // + // 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, &app_id, &track, 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 1808fa362..5f88d25f2 100644 --- a/crates/trident-acl-agent/src/core/config.rs +++ b/crates/trident-acl-agent/src/core/config.rs @@ -40,6 +40,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_"; @@ -67,6 +69,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 path to kubelet's own client certificate, used to derive /// `node_name` when `TRIDENT_ACL_AGENT_KUBERNETES_NODE_NAME_SOURCE` is @@ -86,6 +96,9 @@ const KUBELET_CERT_CN_NODE_PREFIX: &str = "system:node:"; /// Default annotation-key prefix. /// Override with `TRIDENT_ACL_AGENT_KUBERNETES_ANNOTATION_PREFIX`. pub const DEFAULT_ANNOTATION_PREFIX: &str = "acl.microsoft.com"; +/// 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 { @@ -167,6 +180,13 @@ 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)), + 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 @@ -310,6 +330,12 @@ 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_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")] @@ -378,6 +404,9 @@ 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>, @@ -393,6 +422,23 @@ where .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(format!( + "unknown cookie lifetime {value:?} (expected \"per-update\" or \"per-message\")" + )) + }) + }) + .transpose() +} + fn empty_node_name_source_as_none<'de, D>( deserializer: D, ) -> Result, D::Error> @@ -625,6 +671,22 @@ impl FromStr for NodeNameSource { pub struct OrchestrationConfig { pub mode: Mode, pub state_path: PathBuf, + /// Path to the persisted Nebraska affinity cookie (see + /// `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 - + /// 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 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` + /// (`"per-update"` or `"per-message"`). + 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. @@ -652,6 +714,9 @@ impl Default for OrchestrationConfig { Self { 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::PerUpdate, stage_timeout: DEFAULT_STAGE_TIMEOUT, finalize_timeout: DEFAULT_FINALIZE_TIMEOUT, heartbeat_interval: DEFAULT_HEARTBEAT_INTERVAL, @@ -732,6 +797,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", @@ -854,6 +921,18 @@ mod tests { config.orchestration.state_path, PathBuf::from(DEFAULT_STATE_PATH) ); + assert_eq!( + config.orchestration.cookie_path, + PathBuf::from(DEFAULT_NEBRASKA_COOKIE_PATH) + ); + assert!( + !config.orchestration.cookie_enabled, + "cookie_enabled should default to off" + ); + assert_eq!( + config.orchestration.cookie_lifetime, + CookieLifetime::PerUpdate + ); assert_eq!(config.orchestration.stage_timeout, DEFAULT_STAGE_TIMEOUT); assert_eq!( config.orchestration.finalize_timeout, @@ -914,6 +993,15 @@ 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_COOKIE_ENABLED", "true"), + ( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_LIFETIME", + "per-message", + ), ("TRIDENT_ACL_AGENT_ORCHESTRATION_STAGE_TIMEOUT", "21m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_FINALIZE_TIMEOUT", "11m"), ("TRIDENT_ACL_AGENT_ORCHESTRATION_HEARTBEAT_INTERVAL", "45s"), @@ -955,6 +1043,15 @@ 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!(config.orchestration.cookie_enabled); + assert_eq!( + config.orchestration.cookie_lifetime, + CookieLifetime::PerMessage + ); assert_eq!( config.orchestration.stage_timeout, Duration::from_secs(21 * 60) @@ -1023,6 +1120,66 @@ 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(&[( + "TRIDENT_ACL_AGENT_ORCHESTRATION_COOKIE_ENABLED", + "yes please", + )])) + .unwrap_err(); + 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 malformed_cookie_lifetime_is_a_parse_error() { + 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] 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 702a9540b..264c3d90d 100644 --- a/crates/trident-acl-agent/src/core/nebraska/README.md +++ b/crates/trident-acl-agent/src/core/nebraska/README.md @@ -48,6 +48,31 @@ 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` | + +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. + +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. + ### 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 06dcad980..7a6cf59e2 100644 --- a/crates/trident-acl-agent/src/core/nebraska/client.rs +++ b/crates/trident-acl-agent/src/core/nebraska/client.rs @@ -7,6 +7,7 @@ use semver::Version; use url::Url; use super::{ + cookie_store::{CookieLifetime, CookieStore}, error::NebraskaError, event::{ProgressEvent, TerminalEvent}, id::MachineId, @@ -331,6 +332,8 @@ pub struct Client { track: String, machine_id: MachineId, client_version: Option, + cookie_store: Option, + cookie_lifetime: CookieLifetime, transport: T, } @@ -364,10 +367,20 @@ impl Client { track: track.into(), machine_id, client_version: None, + cookie_store: None, + cookie_lifetime: CookieLifetime::default(), transport, } } + /// 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 + } + /// Names the **updater** — the program driving the update — in /// ``. /// @@ -394,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) } @@ -431,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(()) @@ -452,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(()) } @@ -558,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 — @@ -595,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(()) } @@ -611,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() @@ -621,7 +637,33 @@ 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_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 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()); let response = wire::parse_response(&text).map_err(NebraskaError::Parse)?; trace!("Nebraska response: {}", summarize(&response)); @@ -842,11 +884,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 { @@ -854,14 +899,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())) } } @@ -881,13 +941,23 @@ mod tests { struct ScriptedTransport { responses: Vec, calls: Cell, + set_cookie: Option, + cookies_sent: RefCell>>, } 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()) + self.cookies_sent + .borrow_mut() + .push(cookie.map(str::to_string)); + Ok((self.responses[index].clone(), self.set_cookie.clone())) } } @@ -900,6 +970,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()), }, ) } @@ -1348,6 +1420,62 @@ mod tests { assert_eq!(client.transport.calls.get(), 2); } + #[test] + fn completion_retries_honor_cookie_enabled() { + 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 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); + } + #[test] fn complete_after_reboot_fails_when_completion_never_lands() { let client = client_with_responses(vec![IN_PROGRESS]); @@ -1529,6 +1657,217 @@ 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 (_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!(store.load().unwrap(), None); + assert!(!store.path().exists()); + } + + #[test] + 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() { + 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!( + client.transport.last_cookie_sent.borrow().clone(), + None, + "first request has no cookie to send yet" + ); + 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] fn redacted_endpoint_drops_path_and_query() { let url = Url::parse("https://nebraska.example:8443/v1/update/?secret=hunter2").unwrap(); @@ -1625,7 +1964,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 000000000..b8318ae08 --- /dev/null +++ b/crates/trident-acl-agent/src/core/nebraska/cookie_store.rs @@ -0,0 +1,268 @@ +//! 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`. 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}, +}; + +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 +/// 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, 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 + /// 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, +} + +/// 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 { + HeaderValue::from_str(trimmed).context("invalid Nebraska cookie header")?; + 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> { + HeaderValue::from_str(cookie).context("invalid Nebraska cookie header")?; + 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"), + Uuid::new_v4() + )); + + // 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() + ) + }) + })(); + 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 Err(err) = fs::File::open(parent).and_then(|dir| dir.sync_all()) { + warn!("Nebraska cookie was saved but directory sync failed: {err}"); + } + + Ok(()) + } +} + +#[cfg(test)] +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"); + 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 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(); + 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 484efc332..83ab1e36b 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::{CookieLifetime, 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 85d8c053b..5efd33dd3 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}; @@ -41,8 +42,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. @@ -81,7 +96,12 @@ 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 started = Instant::now(); trace!( "Nebraska HTTP POST to '{}': request_bytes={} timeout_ms={}", @@ -89,30 +109,39 @@ impl Transport for ReqwestTransport { body.len(), self.timeout.as_millis() ); - // `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. - let response = self + let mut request = self .client .post(endpoint.as_str()) .header("Content-Type", "application/xml") - .timeout(self.timeout) - .body(body.to_vec()) - .send() - .map_err(|e| { - let message = e.without_url().to_string(); - trace!( - "Nebraska HTTP POST failed: elapsed_ms={} error={message}", - started.elapsed().as_millis() - ); - NebraskaError::Transport(message) - })?; + .timeout(self.timeout); + if let Some(cookie) = cookie { + request = request.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. + let response = request.body(body.to_vec()).send().map_err(|e| { + let message = e.without_url().to_string(); + trace!( + "Nebraska HTTP POST failed: elapsed_ms={} error={message}", + started.elapsed().as_millis() + ); + NebraskaError::Transport(message) + })?; trace!( "Nebraska HTTP response: status={} headers_elapsed_ms={} content_length={:?}", response.status().as_u16(), started.elapsed().as_millis(), response.content_length() ); + + // 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 { @@ -135,7 +164,34 @@ impl Transport for ReqwestTransport { started.elapsed().as_millis() ), } - text + text.map(|body| (body, 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(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("; ")) } } @@ -164,9 +220,9 @@ mod tests { .create(); let endpoint = Url::parse(&format!("{}/v1/update/?secret=hunter2", server.url())).unwrap(); - let result = ReqwestTransport::new().post_xml(&endpoint, b""); + let result = ReqwestTransport::new().post_xml(&endpoint, b"", None); if status == 200 { - assert_eq!(result.unwrap(), ""); + assert_eq!(result.unwrap().0, ""); } else { let error = result.unwrap_err(); assert!(matches!( @@ -197,7 +253,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:?}"); @@ -211,4 +267,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); + } } diff --git a/crates/trident-acl-agent/src/lib.rs b/crates/trident-acl-agent/src/lib.rs index ab48b3156..d5f7be079 100644 --- a/crates/trident-acl-agent/src/lib.rs +++ b/crates/trident-acl-agent/src/lib.rs @@ -18,8 +18,9 @@ use semver::Version; use url::Url; use crate::core::{ + config::CookieLifetime, error::AgentError, - nebraska::{Client, MachineId, NebraskaError}, + nebraska::{Client, CookieStore, MachineId, NebraskaError}, version::FALLBACK_ALWAYS_VERSION, }; @@ -76,14 +77,26 @@ 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. +/// +/// 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, track: &str, 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 client = Client::new(url.clone(), app_id, track, machine_id); + 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); + } + match client.probe( &Version::parse(FALLBACK_ALWAYS_VERSION) .expect("invariant: FALLBACK_ALWAYS_VERSION is valid semver"), @@ -105,6 +118,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 +144,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 +166,9 @@ mod tests { "test", "track", IdSource::MachineIdHashed, + &cookie_store, + true, + CookieLifetime::PerMessage, ) .unwrap(); @@ -155,6 +181,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 +200,9 @@ mod tests { "test", "track", IdSource::MachineIdHashed, + &cookie_store, + true, + CookieLifetime::PerUpdate, ) .unwrap(); @@ -181,14 +211,184 @@ 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, + true, + CookieLifetime::PerUpdate, ) .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, + true, + CookieLifetime::PerMessage, + ) + .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, + 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 + // 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, + CookieLifetime::PerUpdate, + ) + .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" + ); + } } diff --git a/crates/trident-acl-agent/src/omahaonly/mod.rs b/crates/trident-acl-agent/src/omahaonly/mod.rs index 2447d3852..c2c0b5b83 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/docs/Explanation/Trident-ACL-Agent.md b/docs/Explanation/Trident-ACL-Agent.md index a3f95c25c..958e29db4 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 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 2e88631c1..5aff37639 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,6 +8,7 @@ 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 Environment=TRIDENT_ACL_AGENT_KUBERNETES_VALIDATE_NODE_UUID=true diff --git a/tools/storm/aclagent/proxies/nebraska.go b/tools/storm/aclagent/proxies/nebraska.go index 28c8d5c9c..a476b8859 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 @@ -155,6 +188,59 @@ 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, 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() + 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 @@ -219,6 +305,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 745ba8704..b25ab4430 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 @@ -222,4 +257,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 637ea4c9f..8c3d1029f 100644 --- a/tools/storm/aclagent/tests/update.go +++ b/tools/storm/aclagent/tests/update.go @@ -384,17 +384,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) @@ -476,6 +478,42 @@ 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. + // 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 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) + 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) }