diff --git a/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md new file mode 100644 index 0000000000..c5f2430f84 --- /dev/null +++ b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md @@ -0,0 +1,212 @@ +# NATS integration tests + +Tests that run real client code against a real NATS server on your machine — +subjects, JSON payloads, JetStream publishes and acks all go over the wire. +No mocks of NATS. + +## How it works + +``` +test + ├─ NatsServer::start() spawns `nats-server -a 127.0.0.1 -p -js -sd ` + │ waits until the port accepts connections, kills it on drop + ├─ builds the agent side real services on a temp DirectoryManager, + │ NatsConnectionManager given a plain nats:// client + ├─ "tester" async-nats client plays the backend: publishes commands, subscribes to results + └─ asserts on what comes back ACKED / DONE / FAILED, files on disk, signals +``` + +- Every test starts its **own** server on a free port with its own temp storage, + so tests run in parallel and leave nothing behind. +- JetStream is on (`-js`), so code that publishes to a stream (`publish_acked`) + works. The test creates the streams it needs. +- The agent connects without TLS and without a token: the test hands + `NatsConnectionManager` a ready client instead of calling `connect()`. + +### Pieces + +| What | Where | Use | +|---|---|---| +| `NatsServer` | `src/test_support.rs` (`cfg(test)`) | `NatsServer::start() -> Option`, `.url()` | +| `use_client_for_tests` | `src/services/nats_connection_manager.rs` (`cfg(test)`) | puts a connected `async_nats::Client` into the manager | +| Example suite | `src/listener/command/integration_tests.rs` | the feature-flags listener end to end | + +## Running + +Install the server once (single ~1 MB binary): + +```bash +brew install nats-server +``` + +On Windows or Linux, download a release from + and put it on `PATH`. + +Then the usual: + +```bash +make test +``` + +runs them together with the unit tests. + +Only the integration suite, failing if the server is missing: + +```bash +REQUIRE_NATS_SERVER=1 OPENFRAME_VERSION=0.0.0-dev cargo test --lib -- integration_tests +``` + +| Variable | Effect | +|---|---| +| *(none)* | no `nats-server` on `PATH` → each test returns early and **passes without running** | +| `REQUIRE_NATS_SERVER=1` | a missing binary is a test failure — use in CI so the suite can't silently skip | +| `NATS_SERVER_BIN=/path/to/nats-server` | use a binary that isn't on `PATH` | + +A skipped test prints ``skipping: `nats-server` is not installed (brew install nats-server)`` to stderr +(visible with `-- --nocapture`). + +### CI + +CI does not install `nats-server` yet, so there the suite skips and nothing +fails. To make it run: add an install step before `make test` in the Rust job +(`brew install nats-server` on macOS, a release download or `choco install nats-server` +on Windows) and set `REQUIRE_NATS_SERVER: "1"` on that step. + +## Writing a test + +### 1. Start a server and skip if there is none + +```rust +use crate::test_support::NatsServer; + +#[tokio::test] +async fn my_command_is_acked_and_done() { + let Some(server) = NatsServer::start() else { + return; + }; + // ... +} +``` + +Always use the `let … else { return }` form: that is what makes the test skip +on machines without the binary. + +### 2. Build the agent side with real services + +Everything is built on a temp directory, the same way `Client::new` does it: + +```rust +let dir = tempfile::tempdir().unwrap(); +let directories = DirectoryManager::with_custom_dirs( + dir.path().join("logs"), + dir.path().join("app"), + dir.path().join("secured"), +); +let config = AgentConfigurationService::new(directories.clone()).unwrap(); +config + .save_registration_data("it-machine".into(), "client".into(), "secret".into()) + .await + .unwrap(); + +let manager = NatsConnectionManager::new(/* config, initial config, auth, tls, ... */); +manager.use_client_for_tests(async_nats::connect(server.url()).await.unwrap()); +``` + +The machine id saved in the config is what listeners put into their subjects +(`machine..…`). See `Agent::start()` in `listener/command/integration_tests.rs` +for the full list of services and copy it. + +### 3. Create the streams your code publishes to + +`nats-server` starts empty. If the code under test publishes with an ack, create +the stream first, or the publish fails with "no responders": + +```rust +let tester = async_nats::connect(server.url()).await.unwrap(); +async_nats::jetstream::new(tester.clone()) + .create_stream(async_nats::jetstream::stream::Config { + name: "DELIVERY_RESULT".into(), + subjects: vec!["machine.*.delivery.result".into()], + ..Default::default() + }) + .await + .unwrap(); +``` + +### 4. Subscribe to what you assert on before sending + +```rust +let mut results = tester + .subscribe("machine.it-machine.delivery.result") + .await + .unwrap(); +tester.flush().await.unwrap(); +``` + +Core NATS has no backlog: anything published before the subscription exists is +lost. The same applies to the agent's listener — see the next step. + +### 5. Wait until the listener is subscribed + +A listener subscribes asynchronously after `start()`. Don't `sleep`; send a probe +command and repeat until it is answered: + +```rust +timeout(BOUND, async { + loop { + send(probe()).await; + if timeout(Duration::from_millis(200), next_result("probe")).await.is_ok() { + return; + } + } +}) +.await +.expect("listener never answered"); +``` + +Pick a probe that has no side effects (the flags suite uses a stale `sequence`). +Later assertions must ignore extra probe answers — filter results by `dispatchId`. + +### 6. Assert with bounded waits + +- **Something must happen:** wait for it with `timeout(BOUND, …)` (10 s in the + example suite) and fail with a message that says what was expected. +- **Something must not happen:** there is no event to wait for, so wait a short + fixed window (`QUIET`, 500 ms) and assert nothing arrived. Keep these few. +- **Results arrive out of order or interleaved:** don't drop messages you don't + want yet — buffer them (`Agent::next_result` keeps an `unclaimed` list). + Dropping them is the classic flaky-test bug: a later `expect` waits for a + message the helper already threw away. + +### 7. Assert on real side effects + +Read the files the code wrote (`feature_flags.json`), wait on the signals it +raises (`CommandContext::restart_requested()`), check what reached the "backend" +subscription. Prefer these over inspecting internal state. + +## Reusing the example harness + +`listener/command/integration_tests.rs` wraps all of the above in a small +`Agent` struct: + +| Helper | Does | +|---|---| +| `Agent::start()` | server + agent + streams + result subscription + listener up, or `None` to skip | +| `agent.launch()` | starts another listener on the same agent state — "the next process" after a restart | +| `agent.send(payload)` | publishes a command to the listener's subject | +| `agent.expect(dispatch_id, result)` | waits for `ACKED` / `DONE` / `FAILED` for that dispatch | +| `agent.expect_silence(dispatch_id)` | asserts nothing arrives for that dispatch within `QUIET` | +| `agent.persisted()` | reads the state file the listener wrote | +| `command(dispatch_id, sequence, body)` | builds a payload with a `delivery` block | + +For a new delivery command, copy the struct, swap the handler in `launch()` +and the subject in `send()`, and keep the helpers. + +## Limits + +- What the server enforces in production but the local server doesn't: auth, + TLS, and the per-user subject permissions. A `Permissions Violation` on a real + environment is not caught here — check the NATS config when you add a subject. +- Process-level effects (the service manager restarting the client, `exit`) are + out of reach; tests stop at the signal the code raises. +- Each test starts its own server; the example suite (11 tests) takes about a second in total. Fine for dozens of tests, not thousands. diff --git a/clients/openframe-client/src/cli.rs b/clients/openframe-client/src/cli.rs index 3666a4d836..aa0e92a6ba 100644 --- a/clients/openframe-client/src/cli.rs +++ b/clients/openframe-client/src/cli.rs @@ -3,6 +3,7 @@ mod auth; mod healthcheck; mod install; mod permissions; +mod restart; mod uninstall; use crate::doctor::healing::{HealOutcome, HealResult}; @@ -14,6 +15,8 @@ use std::process; use tokio::runtime::Runtime; +pub const RESTART_SERVICE_SUBCOMMAND: &str = "restart-service"; + #[derive(Parser)] #[command(name = "openframe-client", author, version = env!("OPENFRAME_VERSION"), about, long_about = None)] struct Cli { @@ -74,6 +77,9 @@ enum Commands { CheckPermissions, /// Run environment health check (reads config from installed agent) Doctor, + /// Stop and start the client service (run detached by the client to apply feature flags) + #[command(hide = true, name = RESTART_SERVICE_SUBCOMMAND)] + RestartService, } /// Parse CLI arguments and run the requested OpenFrame client command. @@ -91,6 +97,7 @@ pub fn run() -> Result<()> { Some(Commands::Run) => agent::run_direct(&rt), Some(Commands::RunAsService) => agent::run_as_service(), Some(Commands::CheckPermissions) => permissions::run(), + Some(Commands::RestartService) => restart::run(&rt), None => agent::run_legacy(&rt), } diff --git a/clients/openframe-client/src/cli/restart.rs b/clients/openframe-client/src/cli/restart.rs new file mode 100644 index 0000000000..f6e7f93843 --- /dev/null +++ b/clients/openframe-client/src/cli/restart.rs @@ -0,0 +1,22 @@ +use std::process; + +use tokio::runtime::Runtime; +use tracing::{error, info}; + +use crate::service::Service; + +pub(super) fn run(rt: &Runtime) -> ! { + super::init_file_logging(); + info!("Restarting the OpenFrame client service"); + + match rt.block_on(Service::nudge_restart()) { + Ok(()) => { + info!("OpenFrame client service restarted"); + process::exit(0); + } + Err(e) => { + error!("Failed to restart the OpenFrame client service: {:#}", e); + process::exit(1); + } + } +} diff --git a/clients/openframe-client/src/config/update_config.rs b/clients/openframe-client/src/config/update_config.rs index 1932368a0e..78ff5a91dd 100644 --- a/clients/openframe-client/src/config/update_config.rs +++ b/clients/openframe-client/src/config/update_config.rs @@ -17,6 +17,12 @@ pub const CONSUMER_CYCLE_PAUSE_MS: u64 = 30000; // 30 seconds pause between retr // Reconnection pub const RECONNECTION_DELAY_MS: u64 = 5000; // 5 seconds +pub const SELF_RESTART_EXIT_CODE: i32 = 75; +pub const SELF_RESTART_IDLE_WAIT_SECS: u64 = 600; +pub const SELF_RESTART_IDLE_POLL_SECS: u64 = 5; +pub const SELF_RESTART_HANDOFF_SECS: u64 = 900; +pub const DEV_CLIENT_VERSION: &str = "0.0.0-dev"; + // Durable result outbox pub const OUTBOX_MAX_ENTRIES: usize = 500; pub const OUTBOX_FLUSH_INTERVAL_SECS: u64 = 30; diff --git a/clients/openframe-client/src/lib.rs b/clients/openframe-client/src/lib.rs index 152c07126e..ee0458d34f 100644 --- a/clients/openframe-client/src/lib.rs +++ b/clients/openframe-client/src/lib.rs @@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize}; use std::sync::Arc; use tokio::sync::RwLock; use tokio::time::Duration; -use tracing::{error, info}; +use tracing::{error, info, warn}; pub mod clients; mod config; @@ -29,6 +29,9 @@ pub mod service_adapter; pub mod system; pub mod utils; +#[cfg(test)] +mod test_support; + pub mod cli; pub use cli::run; @@ -38,8 +41,12 @@ pub mod executor; use crate::clients::build_agent_http_client; use crate::clients::tool_agent_file_client::ToolAgentFileClient; use crate::clients::{AuthClient, RegistrationClient, ToolApiClient}; -use crate::config::update_config::{DOWNLOAD_CLIENT_TIMEOUT_SECS, HTTP_CLIENT_TIMEOUT_SECS}; +use crate::config::update_config::{ + DOWNLOAD_CLIENT_TIMEOUT_SECS, HTTP_CLIENT_TIMEOUT_SECS, SELF_RESTART_EXIT_CODE, + SELF_RESTART_HANDOFF_SECS, +}; use crate::listener::client_uninstall_message_listener::ClientUninstallMessageListener; +use crate::listener::command::{CommandContext, CommandListener, FeatureFlagsHandler, RestartGate}; use crate::listener::execution_listener::ExecutionListener; use crate::listener::openframe_client_update_listener::OpenFrameClientUpdateListener; use crate::listener::remote_access_message_listener::RemoteAccessMessageListener; @@ -59,6 +66,7 @@ use crate::services::deactivation_service::DeactivationService; use crate::services::device_data_fetcher::DeviceDataFetcher; use crate::services::encryption_service::EncryptionService; use crate::services::execution_service::ExecutionService; +use crate::services::feature_flags_service::FeatureFlagsService; use crate::services::github_download_service::GithubDownloadService; use crate::services::hostname_report_publisher::HostnameReportPublisher; use crate::services::initial_authentication_processor::InitialAuthenticationProcessor; @@ -172,6 +180,8 @@ pub struct Client { tool_agent_update_listener: ToolAgentUpdateListener, client_uninstall_message_listener: ClientUninstallMessageListener, remote_access_message_listener: RemoteAccessMessageListener, + feature_flags_listener: CommandListener, + restart_gate: RestartGate, command_execution_listener: ExecutionListener, script_execution_listener: ExecutionListener, script_bootstrap_execution_listener: ExecutionListener, @@ -366,6 +376,11 @@ impl Client { let tool_ops = ToolOps::default(); + let feature_flags_service = Arc::new(FeatureFlagsService::load( + directory_manager.secured_dir().join("feature_flags.json"), + env!("OPENFRAME_VERSION").to_string(), + )); + // Initialize tool run manager let tool_run_manager = ToolRunManager::new( installed_tools_service.clone(), @@ -553,6 +568,24 @@ impl Client { ); let result_store_for_recovery = result_store.clone(); + let restart_gate = RestartGate::new( + tool_ops.busy_tools(), + execution_service.running_scripts(), + tool_run_manager.client_update_pending(), + ); + let command_context = CommandContext::new( + nats_connection_manager.clone(), + config_service.clone(), + restart_gate.clone(), + nats_message_publisher.clone(), + result_store.clone(), + flush_notify.clone(), + ); + let feature_flags_listener = CommandListener::new( + FeatureFlagsHandler::new(feature_flags_service), + command_context, + ); + let command_execution_listener = ExecutionListener::::new( nats_connection_manager.clone(), nats_message_publisher.clone(), @@ -635,6 +668,8 @@ impl Client { tool_agent_update_listener, client_uninstall_message_listener, remote_access_message_listener, + feature_flags_listener, + restart_gate, command_execution_listener, script_execution_listener, script_bootstrap_execution_listener, @@ -719,6 +754,59 @@ impl Client { // Continue startup even if update handling fails - don't block the client } + self.feature_flags_listener.start().await?; + let restart_gate = self.restart_gate.clone(); + tokio::spawn(async move { + restart_gate.requested().await; + match crate::platform::detached_process::spawn_detached( + &crate::service::Service::get_install_location(), + cli::RESTART_SERVICE_SUBCOMMAND, + ) { + Ok(mut child) => { + info!( + pid = child.id(), + "Service restart handed to a detached process" + ); + let finished = tokio::time::timeout( + Duration::from_secs(SELF_RESTART_HANDOFF_SECS), + tokio::task::spawn_blocking(move || child.wait()), + ) + .await; + match finished { + Ok(Ok(Ok(status))) => warn!( + %status, + "Detached restart finished but did not stop this process, exiting instead" + ), + Ok(Ok(Err(e))) => { + warn!( + "Cannot wait for the detached restart, exiting instead: {:#}", + e + ) + } + Ok(Err(e)) => { + warn!( + "Cannot wait for the detached restart, exiting instead: {:#}", + e + ) + } + Err(_) => warn!( + "Detached restart still running after {}s, exiting instead", + SELF_RESTART_HANDOFF_SECS + ), + } + } + Err(e) => warn!( + "Cannot hand the service restart off, exiting instead: {:#}", + e + ), + } + info!( + code = SELF_RESTART_EXIT_CODE, + "Exiting so the service manager restarts the client" + ); + std::process::exit(SELF_RESTART_EXIT_CODE); + }); + // Start machine heartbeat run manager self.machine_heartbeat_run_manager.start(); self.machine_timezone_run_manager.start(); diff --git a/clients/openframe-client/src/listener/client_update_gate.rs b/clients/openframe-client/src/listener/client_update_gate.rs index d5943ec41f..d0cc76860c 100644 --- a/clients/openframe-client/src/listener/client_update_gate.rs +++ b/clients/openframe-client/src/listener/client_update_gate.rs @@ -14,7 +14,7 @@ pub async fn park_or_dispatch( F: FnOnce(Message) -> Fut + Send + 'static, Fut: Future + Send + 'static, { - if !manager.is_client_update_pending().await { + if !manager.is_client_update_pending() { dispatch(message).await; return; } @@ -25,7 +25,7 @@ pub async fn park_or_dispatch( ); tokio::spawn(async move { - while manager.is_client_update_pending().await { + while manager.is_client_update_pending() { if let Err(e) = message.ack_with(AckKind::Progress).await { warn!( "Failed to send Progress ack for parked {}: {} — abandoning park, redelivery takes over", diff --git a/clients/openframe-client/src/listener/command/core_nats.rs b/clients/openframe-client/src/listener/command/core_nats.rs new file mode 100644 index 0000000000..93639cf605 --- /dev/null +++ b/clients/openframe-client/src/listener/command/core_nats.rs @@ -0,0 +1,130 @@ +use anyhow::{anyhow, bail, Result}; +use async_nats::Message; +use futures::StreamExt; +use serde::Deserialize; +use tracing::{error, info, warn}; + +use super::{CommandHandler, CommandListener, Outcome}; +use crate::models::DeliveryRef; + +#[derive(Deserialize)] +struct Envelope { + delivery: DeliveryRef, +} + +pub(super) async fn report_restart_completion(listener: &CommandListener) { + let Some(delivery) = listener.handler.awaiting_restart() else { + return; + }; + let machine_id = match listener.context.config_service.get_machine_id() { + Ok(machine_id) => machine_id, + Err(e) => { + error!(kind = H::KIND, dispatch_id = %delivery.dispatch_id, "Cannot resolve machine id to report the restart: {:#}", e); + return; + } + }; + info!(kind = H::KIND, dispatch_id = %delivery.dispatch_id, "Applied after restart, reporting done"); + if !listener + .context + .delivery + .settle(&machine_id, delivery, Outcome::Done) + .await + { + return; + } + if let Err(e) = listener.handler.restart_reported() { + warn!( + kind = H::KIND, + "Failed to clear the restart marker, done is re-sent on next start: {:#}", e + ); + } +} + +pub(super) async fn listen(listener: &CommandListener) -> Result<()> { + let nats = &listener.context.nats_connection_manager; + let mut client_rx = nats.on_client_replaced(); + let client = nats.get_client().await?; + let machine_id = listener.context.config_service.get_machine_id()?; + + let subject = H::subject(&machine_id); + let mut subscriber = client + .subscribe(subject.clone()) + .await + .map_err(|e| anyhow!("failed to subscribe to {}: {}", subject, e))?; + info!(kind = H::KIND, subject = %subject, "Command listener active"); + + loop { + tokio::select! { + message = subscriber.next() => match message { + Some(message) => handle(listener, message, &machine_id).await, + None => bail!("subscription to {} ended", subject), + }, + _ = client_rx.changed() => { + info!(kind = H::KIND, "NATS client replaced, resubscribing"); + return Ok(()); + } + } + } +} + +async fn handle( + listener: &CommandListener, + message: Message, + machine_id: &str, +) { + let payload = String::from_utf8_lossy(&message.payload); + info!(kind = H::KIND, "Received command message: {:?}", payload); + + let delivery = match serde_json::from_slice::(&message.payload) { + Ok(envelope) => envelope.delivery, + Err(e) => { + error!( + kind = H::KIND, + "Command message without a valid delivery block, dropping: {:#}", e + ); + return; + } + }; + let context = &listener.context; + let dispatch_id = delivery.dispatch_id.clone(); + + let Some(admitted) = context.restart.admit() else { + info!(kind = H::KIND, dispatch_id = %dispatch_id, "Restart pending, leaving the command to the backend retry"); + return; + }; + + let command = match serde_json::from_slice::(&message.payload) { + Ok(command) => command, + Err(e) => { + error!(kind = H::KIND, dispatch_id = %dispatch_id, "Failed to parse command message: {:#}", e); + context + .delivery + .settle( + machine_id, + delivery, + Outcome::Failed(format!("malformed payload: {e}")), + ) + .await; + return; + } + }; + + context.delivery.acked(machine_id, &delivery).await; + let outcome = listener.handler.run(machine_id, &delivery, command).await; + match &outcome { + Outcome::Done => info!(kind = H::KIND, dispatch_id = %dispatch_id, "Command done"), + Outcome::RestartRequired => { + info!(kind = H::KIND, dispatch_id = %dispatch_id, "Command applied, restart required") + } + Outcome::Failed(reason) => { + error!(kind = H::KIND, dispatch_id = %dispatch_id, "Command failed: {}", reason) + } + } + + if outcome != Outcome::RestartRequired { + context.delivery.settle(machine_id, delivery, outcome).await; + return; + } + drop(admitted); + context.restart.drain_and_request().await; +} diff --git a/clients/openframe-client/src/listener/command/delivery.rs b/clients/openframe-client/src/listener/command/delivery.rs new file mode 100644 index 0000000000..2a7b0b7e7f --- /dev/null +++ b/clients/openframe-client/src/listener/command/delivery.rs @@ -0,0 +1,126 @@ +use std::sync::Arc; + +use tokio::sync::Notify; +use tokio::time::{timeout, Duration}; +use tracing::{error, info, warn}; + +use super::Outcome; +use crate::config::update_config::FLUSH_PUBLISH_TIMEOUT_SECS; +use crate::models::{delivery_result_subject, DeliveryRef, DeliveryResult, DeliveryResultMessage}; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_store::ResultStore; + +#[derive(Clone)] +pub(super) struct Delivery { + publisher: NatsMessagePublisher, + store: Arc, + flush_notify: Arc, +} + +impl Delivery { + pub(super) fn new( + publisher: NatsMessagePublisher, + store: Arc, + flush_notify: Arc, + ) -> Self { + Self { + publisher, + store, + flush_notify, + } + } + + pub(super) async fn acked(&self, machine_id: &str, delivery: &DeliveryRef) { + let message = DeliveryResultMessage { + delivery: delivery.clone(), + result: DeliveryResult::Acked, + error: None, + }; + let subject = delivery_result_subject(machine_id); + let published = match serde_json::to_vec(&message) { + Ok(bytes) => timeout( + Duration::from_secs(FLUSH_PUBLISH_TIMEOUT_SECS), + self.publisher.publish_raw(&subject, &bytes), + ) + .await + .unwrap_or_else(|_| Err(anyhow::anyhow!("publish timed out"))), + Err(e) => Err(e.into()), + }; + match published { + Ok(()) => info!(dispatch_id = %delivery.dispatch_id, "Delivery ACKED"), + Err(e) => warn!( + dispatch_id = %delivery.dispatch_id, + "Failed to report ACKED, the backend retry re-sends the command: {:#}", e + ), + } + } + + pub(super) async fn settle( + &self, + machine_id: &str, + delivery: DeliveryRef, + outcome: Outcome, + ) -> bool { + let dispatch_id = delivery.dispatch_id.clone(); + let (result, error) = match outcome { + Outcome::Done | Outcome::RestartRequired => (DeliveryResult::Done, None), + Outcome::Failed(reason) => (DeliveryResult::Failed, Some(reason)), + }; + let message = DeliveryResultMessage { + delivery, + result, + error, + }; + let subject = delivery_result_subject(machine_id); + + if self.store.enabled() { + let enqueued = match serde_json::to_vec(&message) { + Ok(bytes) => { + self.store + .enqueue(format!("delivery:{}", dispatch_id), subject, bytes) + .await + } + Err(e) => Err(e.into()), + }; + return match enqueued { + Ok(()) => { + self.flush_notify.notify_one(); + info!(dispatch_id = %dispatch_id, ?result, "Delivery result queued"); + true + } + Err(e) => { + error!( + dispatch_id = %dispatch_id, + ?result, + "Failed to persist delivery result, the backend closes it on timeout: {:#}", e + ); + false + } + }; + } + + let published = match serde_json::to_vec(&message) { + Ok(bytes) => timeout( + Duration::from_secs(FLUSH_PUBLISH_TIMEOUT_SECS), + self.publisher.publish_acked(&subject, &bytes), + ) + .await + .unwrap_or_else(|_| Err(anyhow::anyhow!("publish timed out"))), + Err(e) => Err(e.into()), + }; + match published { + Ok(()) => { + info!(dispatch_id = %dispatch_id, ?result, "Delivery result published"); + true + } + Err(e) => { + error!( + dispatch_id = %dispatch_id, + ?result, + "Failed to publish delivery result without a durable store, the backend closes it on timeout: {:#}", e + ); + false + } + } + } +} diff --git a/clients/openframe-client/src/listener/command/handlers.rs b/clients/openframe-client/src/listener/command/handlers.rs new file mode 100644 index 0000000000..3c730de1f7 --- /dev/null +++ b/clients/openframe-client/src/listener/command/handlers.rs @@ -0,0 +1,65 @@ +use std::sync::Arc; + +use anyhow::Result; +use tracing::info; + +use super::{CommandHandler, Outcome}; +use crate::models::{DeliveryRef, FeatureFlagsMessage}; +use crate::services::feature_flags_service::{FeatureFlagsService, FlagsApplied}; + +#[derive(Clone)] +pub struct FeatureFlagsHandler { + service: Arc, +} + +impl FeatureFlagsHandler { + pub fn new(service: Arc) -> Self { + Self { service } + } +} + +impl CommandHandler for FeatureFlagsHandler { + type Command = FeatureFlagsMessage; + + const KIND: &'static str = "feature-flags"; + + fn subject(machine_id: &str) -> String { + format!("machine.{}.feature-flags", machine_id) + } + + async fn run( + &self, + machine_id: &str, + delivery: &DeliveryRef, + command: Self::Command, + ) -> Outcome { + let Some(sequence) = delivery.sequence else { + return Outcome::Failed("feature flags without a sequence".to_string()); + }; + let applied_sequence = self.service.applied_sequence(machine_id); + if sequence <= applied_sequence { + info!( + sequence, + applied_sequence, "Stale feature flags, already applied a newer set" + ); + return Outcome::Done; + } + match self.service.apply(machine_id, delivery, sequence, command) { + Ok(FlagsApplied::Unchanged) => Outcome::Done, + Ok(FlagsApplied::RestartRequired) => Outcome::RestartRequired, + Err(e) => Outcome::Failed(format!("{e:#}")), + } + } + + fn awaiting_restart(&self) -> Option { + self.service.awaiting_restart() + } + + fn restart_reported(&self) -> Result<()> { + self.service.restart_reported() + } +} + +#[cfg(test)] +#[path = "handlers_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/listener/command/handlers_tests.rs b/clients/openframe-client/src/listener/command/handlers_tests.rs new file mode 100644 index 0000000000..50a5fb9bb3 --- /dev/null +++ b/clients/openframe-client/src/listener/command/handlers_tests.rs @@ -0,0 +1,125 @@ +use super::*; +use serde_json::{json, Value}; +use tempfile::TempDir; + +const MACHINE: &str = "machine-1"; +const VERSION: &str = "1.5.1"; + +fn handler() -> (FeatureFlagsHandler, TempDir) { + let dir = tempfile::tempdir().unwrap(); + let service = + FeatureFlagsService::load(dir.path().join("feature_flags.json"), VERSION.to_string()); + (FeatureFlagsHandler::new(Arc::new(service)), dir) +} + +fn delivery(sequence: Option) -> DeliveryRef { + DeliveryRef { + delivery_type: "CLIENT_FEATURE_FLAGS".to_string(), + target_id: "tenant".to_string(), + dispatch_id: "dispatch".to_string(), + sequence, + } +} + +fn tool_installation(value: bool, min_version: &str) -> FeatureFlagsMessage { + message( + json!({ "delivery-tool-installation": { "enabled": value, "minVersion": min_version } }), + ) +} + +fn message(flags: Value) -> FeatureFlagsMessage { + serde_json::from_value(json!({ "flags": flags })).unwrap() +} + +#[test] +fn the_subject_is_machine_scoped() { + assert_eq!( + FeatureFlagsHandler::subject("abc123"), + "machine.abc123.feature-flags" + ); +} + +#[tokio::test] +async fn flags_without_a_sequence_fail() { + let (handler, _dir) = handler(); + + let outcome = handler + .run(MACHINE, &delivery(None), tool_installation(true, "1.0.0")) + .await; + + assert!(matches!(outcome, Outcome::Failed(_))); + assert_eq!(handler.awaiting_restart(), None); +} + +#[tokio::test] +async fn an_applicable_change_restarts_and_anything_else_is_done() { + let cases = [ + ( + tool_installation(false, "1.0.0"), + Outcome::Done, + "switched off, as running", + ), + ( + tool_installation(true, "2.0.0"), + Outcome::Done, + "for a newer client", + ), + ( + tool_installation(true, "1.0.0"), + Outcome::RestartRequired, + "switched on for this client", + ), + ]; + for (message, expected, why) in cases { + let (handler, _dir) = handler(); + + let outcome = handler.run(MACHINE, &delivery(Some(1)), message).await; + + assert_eq!(outcome, expected, "{why}"); + } +} + +#[tokio::test] +async fn a_stale_or_repeated_sequence_is_done_without_applying() { + let (handler, _dir) = handler(); + handler + .run( + MACHINE, + &delivery(Some(5)), + tool_installation(true, "1.0.0"), + ) + .await; + + for stale in [5, 4] { + let outcome = handler + .run(MACHINE, &delivery(Some(stale)), message(json!({}))) + .await; + + assert_eq!(outcome, Outcome::Done, "sequence {stale}"); + assert_eq!(handler.service.applied_sequence(MACHINE), 5); + assert!(handler.awaiting_restart().is_some()); + } +} + +#[tokio::test] +async fn a_reregistered_machine_applies_a_lower_sequence() { + let (handler, _dir) = handler(); + handler + .run( + MACHINE, + &delivery(Some(50)), + tool_installation(true, "1.0.0"), + ) + .await; + + let outcome = handler + .run( + "machine-2", + &delivery(Some(3)), + tool_installation(true, "1.0.0"), + ) + .await; + + assert_eq!(outcome, Outcome::RestartRequired); + assert_eq!(handler.service.applied_sequence("machine-2"), 3); +} diff --git a/clients/openframe-client/src/listener/command/integration_tests.rs b/clients/openframe-client/src/listener/command/integration_tests.rs new file mode 100644 index 0000000000..b7e8036728 --- /dev/null +++ b/clients/openframe-client/src/listener/command/integration_tests.rs @@ -0,0 +1,552 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use futures::StreamExt; +use serde_json::json; +use tempfile::TempDir; +use tokio::sync::Notify; +use tokio::time::{timeout, Duration}; + +use super::{CommandContext, CommandListener, FeatureFlagsHandler, RestartGate}; +use crate::clients::AuthClient; +use crate::config::update_config::DEV_CLIENT_VERSION; +use crate::models::{DeliveryResult, DeliveryResultMessage, Feature, FeatureFlagsState}; +use crate::platform::DirectoryManager; +use crate::services::encryption_service::EncryptionService; +use crate::services::execution_service::ExecutionService; +use crate::services::feature_flags_service::FeatureFlagsService; +use crate::services::local_tls_config_provider::LocalTlsConfigProvider; +use crate::services::nats_connection_manager::NatsConnectionManager; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_outbox_run_manager::ResultOutboxRunManager; +use crate::services::result_store::ResultStore; +use crate::services::shared_token_service::SharedTokenService; +use crate::services::tool_run_manager::ToolRunManager; +use crate::services::{ + AgentAuthService, AgentConfigurationService, DeactivationService, InitialConfigurationService, + InstalledToolsService, MachineIdService, ToolCommandParamsResolver, ToolKillService, ToolOps, +}; +use crate::test_support::NatsServer; + +const MACHINE: &str = "it-machine"; +const BOUND: Duration = Duration::from_secs(10); +const QUIET: Duration = Duration::from_millis(500); +const PROBE: &str = "probe"; +const VERSION: &str = "1.5.1"; + +struct Agent { + _server: NatsServer, + _dir: TempDir, + tester: async_nats::Client, + results: async_nats::Subscriber, + unclaimed: Vec, + manager: NatsConnectionManager, + config: AgentConfigurationService, + tool_run_manager: ToolRunManager, + publisher: NatsMessagePublisher, + store: Arc, + flush: Arc, + flags_path: PathBuf, + version: String, +} + +impl Agent { + async fn start() -> Option<(Self, RestartGate)> { + Self::start_as(VERSION).await + } + + async fn start_as(version: &str) -> Option<(Self, RestartGate)> { + let server = NatsServer::start()?; + let dir = tempfile::tempdir().unwrap(); + let directories = DirectoryManager::with_custom_dirs( + dir.path().join("logs"), + dir.path().join("app"), + dir.path().join("secured"), + ); + let config = AgentConfigurationService::new(directories.clone()).unwrap(); + config + .save_registration_data(MACHINE.into(), "client".into(), "secret".into()) + .await + .unwrap(); + let initial = InitialConfigurationService::new(directories.clone()).unwrap(); + let tool_run_manager = ToolRunManager::new( + InstalledToolsService::new(directories.clone()).unwrap(), + ToolCommandParamsResolver::new(directories.clone(), initial.clone(), config.clone()), + ToolKillService::new(), + ToolOps::default(), + ); + let deactivation = DeactivationService::new(&directories); + let auth = AgentAuthService::new( + AuthClient::new( + "http://127.0.0.1:9".into(), + reqwest::Client::new(), + deactivation.clone(), + ), + config.clone(), + SharedTokenService::new(directories.clone(), EncryptionService::new()), + ); + let manager = NatsConnectionManager::new( + server.url(), + config.clone(), + initial.clone(), + auth, + LocalTlsConfigProvider::new(initial), + deactivation, + MachineIdService::new(&directories), + ); + manager.use_client_for_tests(async_nats::connect(server.url()).await.unwrap()); + + let tester = async_nats::connect(server.url()).await.unwrap(); + async_nats::jetstream::new(tester.clone()) + .create_stream(async_nats::jetstream::stream::Config { + name: "DELIVERY_RESULT".into(), + subjects: vec!["machine.*.delivery.result".into()], + ..Default::default() + }) + .await + .unwrap(); + let results = tester + .subscribe(format!("machine.{MACHINE}.delivery.result")) + .await + .unwrap(); + tester.flush().await.unwrap(); + + let publisher = NatsMessagePublisher::new(manager.clone()); + let store = Arc::new(ResultStore::open_or_degrade( + directories.secured_dir().join("scheduled_outbox.redb"), + )); + let flush = Arc::new(Notify::new()); + ResultOutboxRunManager::new(store.clone(), Arc::new(publisher.clone()), flush.clone()) + .start(); + + let mut agent = Self { + _server: server, + flags_path: directories.secured_dir().join("feature_flags.json"), + version: version.to_string(), + _dir: dir, + tester, + results, + unclaimed: Vec::new(), + manager, + config, + tool_run_manager, + publisher, + store, + flush, + }; + let (gate, _service) = agent.launch().await; + Some((agent, gate)) + } + + async fn launch(&mut self) -> (RestartGate, Arc) { + let gate = RestartGate::new( + ToolOps::default().busy_tools(), + ExecutionService::new().running_scripts(), + self.tool_run_manager.client_update_pending(), + ); + let context = CommandContext::new( + self.manager.clone(), + self.config.clone(), + gate.clone(), + self.publisher.clone(), + self.store.clone(), + self.flush.clone(), + ); + let service = Arc::new(FeatureFlagsService::load( + self.flags_path.clone(), + self.version.clone(), + )); + CommandListener::new(FeatureFlagsHandler::new(service.clone()), context) + .start() + .await + .unwrap(); + self.until_listening().await; + (gate, service) + } + + async fn launch_as(&mut self, version: &str) -> (RestartGate, Arc) { + self.version = version.to_string(); + self.launch().await + } + + async fn until_listening(&mut self) { + timeout(BOUND, async { + loop { + self.send(command(PROBE, Some(0), json!({}))).await; + if timeout(Duration::from_millis(200), self.next_result(PROBE)) + .await + .is_ok() + { + return; + } + } + }) + .await + .expect("the flags listener never answered"); + } + + async fn send(&self, payload: serde_json::Value) { + self.tester + .publish( + format!("machine.{MACHINE}.feature-flags"), + serde_json::to_vec(&payload).unwrap().into(), + ) + .await + .unwrap(); + self.tester.flush().await.unwrap(); + } + + async fn next_result(&mut self, dispatch_id: &str) -> DeliveryResultMessage { + if let Some(index) = self + .unclaimed + .iter() + .position(|result| result.delivery.dispatch_id == dispatch_id) + { + return self.unclaimed.remove(index); + } + loop { + let message = self + .results + .next() + .await + .expect("result subscription ended"); + let result: DeliveryResultMessage = serde_json::from_slice(&message.payload).unwrap(); + if result.delivery.dispatch_id == dispatch_id { + return result; + } + self.unclaimed.push(result); + } + } + + async fn expect( + &mut self, + dispatch_id: &str, + expected: DeliveryResult, + ) -> DeliveryResultMessage { + let result = timeout(BOUND, self.next_result(dispatch_id)) + .await + .unwrap_or_else(|_| panic!("no {expected:?} for {dispatch_id}")); + assert_eq!(result.result, expected, "{dispatch_id}: {result:?}"); + result + } + + async fn expect_silence(&mut self, dispatch_id: &str) { + let heard = timeout(QUIET, self.next_result(dispatch_id)).await; + assert!( + heard.is_err(), + "unexpected result for {dispatch_id}: {heard:?}" + ); + } + + fn persisted(&self) -> FeatureFlagsState { + serde_json::from_slice(&std::fs::read(&self.flags_path).unwrap()).unwrap() + } +} + +fn entry(value: bool, min_version: &str) -> serde_json::Value { + json!({ "enabled": value, "minVersion": min_version }) +} + +async fn expect_no_restart(gate: &RestartGate) { + let requested = timeout(QUIET, gate.requested()).await; + assert!(requested.is_err(), "a restart was requested"); +} + +fn command( + dispatch_id: &str, + sequence: Option, + flags: serde_json::Value, +) -> serde_json::Value { + let mut delivery = json!({ + "type": "CLIENT_FEATURE_FLAGS", + "targetId": "tenant", + "dispatchId": dispatch_id, + }); + if let Some(sequence) = sequence { + delivery["sequence"] = json!(sequence); + } + json!({ "delivery": delivery, "flags": flags }) +} + +#[tokio::test] +async fn unchanged_flags_are_acked_done_and_persisted() { + let Some((mut agent, _gate)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d1", + Some(1), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), + )) + .await; + + agent.expect("d1", DeliveryResult::Acked).await; + agent.expect("d1", DeliveryResult::Done).await; + let persisted = agent.persisted(); + assert_eq!(persisted.sequence, 1); + assert_eq!(persisted.awaiting_restart, None); +} + +#[tokio::test] +async fn changed_flags_restart_and_the_next_process_reports_done() { + let Some((mut agent, gate)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d2", + Some(2), + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), + )) + .await; + + agent.expect("d2", DeliveryResult::Acked).await; + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); + agent.expect_silence("d2").await; + assert_eq!( + agent + .persisted() + .awaiting_restart + .map(|delivery| delivery.dispatch_id), + Some("d2".to_string()) + ); + + agent.launch().await; + + agent.expect("d2", DeliveryResult::Done).await; + timeout(BOUND, async { + while agent.persisted().awaiting_restart.is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the restart marker was never cleared"); +} + +#[tokio::test] +async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { + let Some((mut agent, gate)) = Agent::start().await else { + return; + }; + agent + .send(command( + "d3", + Some(3), + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), + )) + .await; + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); + + agent + .send(command( + "d4", + Some(4), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), + )) + .await; + + agent.expect_silence("d4").await; + assert_eq!(agent.persisted().sequence, 3); +} + +#[tokio::test] +async fn stale_sequence_is_done_without_applying() { + let Some((mut agent, _gate)) = Agent::start().await else { + return; + }; + agent.send(command("d5", Some(5), json!({}))).await; + agent.expect("d5", DeliveryResult::Acked).await; + agent.expect("d5", DeliveryResult::Done).await; + + agent + .send(command( + "d6", + Some(4), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), + )) + .await; + + agent.expect("d6", DeliveryResult::Acked).await; + agent.expect("d6", DeliveryResult::Done).await; + assert_eq!(agent.persisted().sequence, 5); +} + +#[tokio::test] +async fn malformed_flags_fail_with_a_reason() { + let Some((mut agent, _gate)) = Agent::start().await else { + return; + }; + + for (dispatch_id, payload) in [ + ("d7", command("d7", Some(7), json!("not-a-map"))), + ("d7-null", command("d7-null", Some(7), json!(null))), + ( + "d7-missing", + json!({ "delivery": { "type": "CLIENT_FEATURE_FLAGS", "targetId": "tenant", "dispatchId": "d7-missing", "sequence": 7 } }), + ), + ] { + agent.send(payload).await; + + let failed = agent.expect(dispatch_id, DeliveryResult::Failed).await; + assert!( + failed.error.unwrap().contains("malformed payload"), + "{dispatch_id}" + ); + } + assert!(!agent.flags_path.exists()); +} + +#[tokio::test] +async fn a_message_without_delivery_gets_no_answer() { + let Some((mut agent, _gate)) = Agent::start().await else { + return; + }; + + agent + .send(json!({ "flags": { "delivery-tool-installation": entry(true, "1.0.0") } })) + .await; + agent + .send(command( + "d8", + Some(8), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), + )) + .await; + + let first = timeout(BOUND, async { + loop { + let message = agent.results.next().await.unwrap(); + let result: DeliveryResultMessage = serde_json::from_slice(&message.payload).unwrap(); + if result.delivery.dispatch_id != PROBE { + return result; + } + } + }) + .await + .unwrap(); + assert_eq!(first.delivery.dispatch_id, "d8"); + assert_eq!(agent.persisted().sequence, 8); +} + +#[tokio::test] +async fn flags_without_a_sequence_fail() { + let Some((mut agent, _gate)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d9", + None, + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), + )) + .await; + + agent.expect("d9", DeliveryResult::Acked).await; + let failed = agent.expect("d9", DeliveryResult::Failed).await; + assert!(failed.error.unwrap().contains("sequence")); +} + +#[tokio::test] +async fn a_flag_for_a_newer_client_is_stored_without_a_restart_and_turns_on_after_the_update() { + let Some((mut agent, gate)) = Agent::start_as("1.5.9").await else { + return; + }; + + agent + .send(command( + "d20", + Some(20), + json!({ "delivery-tool-installation": entry(true, "1.6.0") }), + )) + .await; + + agent.expect("d20", DeliveryResult::Acked).await; + agent.expect("d20", DeliveryResult::Done).await; + expect_no_restart(&gate).await; + assert!(agent + .persisted() + .flags + .contains_key("delivery-tool-installation")); + + let (_updated, service) = agent.launch_as("1.6.0").await; + + assert!(service + .flags() + .is_enabled(Feature::DeliveryToolInstallation)); + agent.expect_silence("d20").await; +} + +#[tokio::test] +async fn a_rollback_below_the_min_version_runs_with_the_flag_off() { + let Some((mut agent, gate)) = Agent::start().await else { + return; + }; + agent + .send(command( + "d21", + Some(21), + json!({ "delivery-tool-installation": entry(true, "1.5.0") }), + )) + .await; + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); + + let (_rolled_back, service) = agent.launch_as("1.4.0").await; + + assert!(!service + .flags() + .is_enabled(Feature::DeliveryToolInstallation)); +} + +#[tokio::test] +async fn a_dev_build_applies_a_true_flag_whatever_its_min_version() { + let Some((mut agent, gate)) = Agent::start_as(DEV_CLIENT_VERSION).await else { + return; + }; + + agent + .send(command( + "d22", + Some(22), + json!({ "delivery-tool-installation": entry(true, "99.0.0") }), + )) + .await; + + agent.expect("d22", DeliveryResult::Acked).await; + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); +} + +#[tokio::test] +async fn invalid_flag_entries_are_done_without_a_restart() { + let Some((mut agent, gate)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d23", + Some(23), + json!({ + "delivery-tool-installation": true, + "delivery-client-uninstall": { "enabled": true }, + "delivery-client-update": { "enabled": true, "minVersion": "latest" } + }), + )) + .await; + + agent.expect("d23", DeliveryResult::Acked).await; + agent.expect("d23", DeliveryResult::Done).await; + expect_no_restart(&gate).await; + assert!(agent.persisted().flags.is_empty()); +} diff --git a/clients/openframe-client/src/listener/command/mod.rs b/clients/openframe-client/src/listener/command/mod.rs new file mode 100644 index 0000000000..6b8c13b773 --- /dev/null +++ b/clients/openframe-client/src/listener/command/mod.rs @@ -0,0 +1,118 @@ +mod core_nats; +mod delivery; +mod handlers; +mod restart; + +use std::future::Future; +use std::sync::Arc; + +use anyhow::Result; +use serde::de::DeserializeOwned; +use tokio::sync::Notify; +use tokio::time::Duration; +use tracing::{error, info}; + +use crate::config::update_config::RECONNECTION_DELAY_MS; +use crate::models::DeliveryRef; +use crate::services::nats_connection_manager::NatsConnectionManager; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_store::ResultStore; +use crate::services::AgentConfigurationService; + +use delivery::Delivery; + +pub use handlers::FeatureFlagsHandler; +pub use restart::RestartGate; + +pub trait CommandHandler: Clone + Send + Sync + 'static { + type Command: DeserializeOwned + Send + 'static; + + const KIND: &'static str; + + fn subject(machine_id: &str) -> String; + + fn run( + &self, + machine_id: &str, + delivery: &DeliveryRef, + command: Self::Command, + ) -> impl Future + Send; + + fn awaiting_restart(&self) -> Option { + None + } + + fn restart_reported(&self) -> Result<()> { + Ok(()) + } +} + +#[derive(Debug, PartialEq)] +pub enum Outcome { + Done, + RestartRequired, + Failed(String), +} + +#[derive(Clone)] +pub struct CommandContext { + nats_connection_manager: NatsConnectionManager, + config_service: AgentConfigurationService, + restart: RestartGate, + delivery: Delivery, +} + +impl CommandContext { + pub fn new( + nats_connection_manager: NatsConnectionManager, + config_service: AgentConfigurationService, + restart: RestartGate, + publisher: NatsMessagePublisher, + result_store: Arc, + flush_notify: Arc, + ) -> Self { + Self { + nats_connection_manager, + config_service, + restart, + delivery: Delivery::new(publisher, result_store, flush_notify), + } + } +} + +#[derive(Clone)] +pub struct CommandListener { + handler: H, + context: CommandContext, +} + +impl CommandListener { + pub fn new(handler: H, context: CommandContext) -> Self { + Self { handler, context } + } + + pub async fn start(&self) -> Result> { + let listener = self.clone(); + let handle = tokio::spawn(async move { + core_nats::report_restart_completion(&listener).await; + loop { + info!(kind = H::KIND, "Starting command listener..."); + match core_nats::listen(&listener).await { + Ok(()) => info!(kind = H::KIND, "Command listener rebinding"), + Err(e) => error!(kind = H::KIND, "Command listener error: {:#}", e), + } + info!( + kind = H::KIND, + "Reconnecting command listener in {} seconds...", + RECONNECTION_DELAY_MS / 1000 + ); + tokio::time::sleep(Duration::from_millis(RECONNECTION_DELAY_MS)).await; + } + }); + Ok(handle) + } +} + +#[cfg(test)] +#[path = "integration_tests.rs"] +mod integration_tests; diff --git a/clients/openframe-client/src/listener/command/restart.rs b/clients/openframe-client/src/listener/command/restart.rs new file mode 100644 index 0000000000..d16a655078 --- /dev/null +++ b/clients/openframe-client/src/listener/command/restart.rs @@ -0,0 +1,87 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +use tokio::sync::{Notify, OwnedRwLockReadGuard, RwLock}; +use tokio::time::{sleep, timeout_at, Duration, Instant}; +use tracing::{info, warn}; + +use crate::config::update_config::{SELF_RESTART_IDLE_POLL_SECS, SELF_RESTART_IDLE_WAIT_SECS}; +use crate::services::execution_service::RunningScripts; +use crate::services::tool_ops::BusyTools; +use crate::services::tool_run_manager::ClientUpdatePending; + +#[derive(Clone)] +pub struct RestartGate { + closing: Arc, + work: Arc>, + requested: Arc, + busy_tools: BusyTools, + running_scripts: RunningScripts, + client_update: ClientUpdatePending, +} + +impl RestartGate { + pub fn new( + busy_tools: BusyTools, + running_scripts: RunningScripts, + client_update: ClientUpdatePending, + ) -> Self { + Self { + closing: Arc::new(AtomicBool::new(false)), + work: Arc::new(RwLock::new(())), + requested: Arc::new(Notify::new()), + busy_tools, + running_scripts, + client_update, + } + } + + pub(super) fn admit(&self) -> Option> { + let admitted = self.work.clone().try_read_owned().ok()?; + if self.closing.load(Ordering::Acquire) { + return None; + } + Some(admitted) + } + + pub(super) async fn drain_and_request(&self) { + self.closing.store(true, Ordering::Release); + let deadline = Instant::now() + Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS); + info!("Restart scheduled: no new commands are taken, waiting for running ones to finish"); + + if timeout_at(deadline, self.work.write()).await.is_err() { + warn!( + "Commands still running after {}s, restarting anyway", + SELF_RESTART_IDLE_WAIT_SECS + ); + } + while self.busy() && Instant::now() < deadline { + sleep(Duration::from_secs(SELF_RESTART_IDLE_POLL_SECS)).await; + } + if self.busy() { + warn!( + "Tool operation, script or client update still running after {}s, restarting anyway", + SELF_RESTART_IDLE_WAIT_SECS + ); + } + // A running client update restarts the client itself; stopping the service under it could + // interrupt the binary swap. With a tool still busy the update is only deferred and keeps + // re-marking the flag on redelivery, so it is not worth waiting for. + while self.client_update.is_pending() && !self.busy_tools.any() { + sleep(Duration::from_secs(SELF_RESTART_IDLE_POLL_SECS)).await; + } + self.requested.notify_one(); + } + + fn busy(&self) -> bool { + self.busy_tools.any() || self.running_scripts.any() || self.client_update.is_pending() + } + + pub async fn requested(&self) { + self.requested.notified().await; + } +} + +#[cfg(test)] +#[path = "restart_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/listener/command/restart_tests.rs b/clients/openframe-client/src/listener/command/restart_tests.rs new file mode 100644 index 0000000000..0f09de6a02 --- /dev/null +++ b/clients/openframe-client/src/listener/command/restart_tests.rs @@ -0,0 +1,154 @@ +use super::*; +use crate::config::update_config::SELF_RESTART_IDLE_WAIT_SECS; +use crate::services::execution_service::ExecutionService; +use crate::services::tool_ops::ToolOps; +use crate::services::tool_run_manager::ClientUpdatePendingFlag; +use tokio::time::timeout; + +const BOUND: Duration = Duration::from_secs(10); + +#[derive(Default)] +struct Work { + tool_ops: ToolOps, + scripts: ExecutionService, + client_update: ClientUpdatePendingFlag, +} + +impl Work { + fn gate(&self) -> RestartGate { + RestartGate::new( + self.tool_ops.busy_tools(), + self.scripts.running_scripts(), + self.client_update.view(), + ) + } +} + +async fn until_closed(gate: &RestartGate) { + timeout(BOUND, async { + while gate.admit().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the gate never closed"); +} + +fn start_drain(gate: &RestartGate) -> tokio::task::JoinHandle<()> { + let gate = gate.clone(); + tokio::spawn(async move { gate.drain_and_request().await }) +} + +async fn until_requested(gate: &RestartGate) { + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); +} + +async fn drain_waits_until_released(gate: &RestartGate, release: impl FnOnce()) { + let draining = start_drain(gate); + until_closed(gate).await; + assert!(!draining.is_finished()); + + release(); + until_requested(gate).await; +} + +#[tokio::test] +async fn admits_concurrent_commands_while_no_restart_is_pending() { + let gate = Work::default().gate(); + + let first = gate.admit(); + let second = gate.admit(); + + assert!(first.is_some()); + assert!(second.is_some()); +} + +#[tokio::test] +async fn drain_refuses_new_commands_and_waits_for_the_admitted_one() { + let gate = Work::default().gate(); + let running = gate.admit().expect("admitted before the restart"); + + drain_waits_until_released(&gate, || drop(running)).await; + + assert!(gate.admit().is_none()); +} + +#[tokio::test] +async fn drain_waits_for_a_busy_tool() { + let work = Work::default(); + let gate = work.gate(); + let op = work + .tool_ops + .try_lock("fleetmdm-agent") + .unwrap() + .mark_busy(); + + drain_waits_until_released(&gate, || drop(op)).await; +} + +#[tokio::test] +async fn drain_waits_for_a_running_script() { + let work = Work::default(); + let gate = work.gate(); + let script = work.scripts.running_for_tests(); + + drain_waits_until_released(&gate, || drop(script)).await; +} + +#[tokio::test] +async fn drain_waits_for_a_pending_client_update() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); + + drain_waits_until_released(&gate, || work.client_update.clear()).await; +} + +#[tokio::test] +async fn the_gate_stays_closed_after_the_restart_is_requested() { + let gate = Work::default().gate(); + + start_drain(&gate); + until_requested(&gate).await; + + assert!(gate.admit().is_none()); +} + +#[tokio::test(start_paused = true)] +async fn a_running_client_update_holds_the_restart_past_the_deadline() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); + + let draining = start_drain(&gate); + sleep(Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS * 2)).await; + assert!(!draining.is_finished()); + + work.client_update.clear(); + until_requested(&gate).await; +} + +#[tokio::test(start_paused = true)] +async fn a_client_update_deferred_by_a_busy_tool_does_not_hold_the_restart() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); + let _op = work + .tool_ops + .try_lock("fleetmdm-agent") + .unwrap() + .mark_busy(); + + let draining = start_drain(&gate); + timeout( + Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS * 2), + gate.requested(), + ) + .await + .expect("restart was never requested"); + + assert!(work.client_update.view().is_pending()); + draining.await.unwrap(); +} diff --git a/clients/openframe-client/src/listener/execution_listener.rs b/clients/openframe-client/src/listener/execution_listener.rs index 6e4e5a9af6..958d562527 100644 --- a/clients/openframe-client/src/listener/execution_listener.rs +++ b/clients/openframe-client/src/listener/execution_listener.rs @@ -13,7 +13,8 @@ use crate::config::update_config::{ FALLBACK_PUBLISH_MAX_RETRY_DELAY_MS, RECONNECTION_DELAY_MS, }; use crate::models::{ - is_ack_subject, ExecutionAck, ExecutionMessage, ExecutionRequest, RmmResult, EXECUTION_ACK_KIND, + requires_stream_ack, ExecutionAck, ExecutionMessage, ExecutionRequest, RmmResult, + EXECUTION_ACK_KIND, }; use crate::services::execution_service::ExecutionService; use crate::services::nats_connection_manager::NatsConnectionManager; @@ -254,7 +255,7 @@ impl ExecutionListener { fn spawn_fallback_publish(&self, subject: String, bytes: Vec) { let publisher = self.nats_message_publisher.clone(); - let acked = is_ack_subject(&subject); + let acked = requires_stream_ack(&subject); tokio::spawn(async move { let max_backoff = Duration::from_millis(FALLBACK_PUBLISH_MAX_RETRY_DELAY_MS); let mut backoff = Duration::from_millis(FALLBACK_PUBLISH_INITIAL_RETRY_DELAY_MS); diff --git a/clients/openframe-client/src/listener/mod.rs b/clients/openframe-client/src/listener/mod.rs index bbf2880fe4..b8cac5820c 100644 --- a/clients/openframe-client/src/listener/mod.rs +++ b/clients/openframe-client/src/listener/mod.rs @@ -1,5 +1,6 @@ pub mod client_uninstall_message_listener; pub mod client_update_gate; +pub mod command; pub mod execution_listener; pub mod openframe_client_update_listener; pub mod remote_access_message_listener; @@ -9,6 +10,7 @@ pub mod tool_restart_message_listener; pub mod tool_uninstall_message_listener; pub use client_uninstall_message_listener::ClientUninstallMessageListener; +pub use command::CommandListener; pub use execution_listener::ExecutionListener; pub use openframe_client_update_listener::OpenFrameClientUpdateListener; pub use remote_access_message_listener::RemoteAccessMessageListener; diff --git a/clients/openframe-client/src/models/delivery.rs b/clients/openframe-client/src/models/delivery.rs new file mode 100644 index 0000000000..a4cc814dd7 --- /dev/null +++ b/clients/openframe-client/src/models/delivery.rs @@ -0,0 +1,34 @@ +use serde::{Deserialize, Serialize}; + +pub const DELIVERY_RESULT_KIND: &str = "delivery.result"; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DeliveryRef { + #[serde(rename = "type")] + pub delivery_type: String, + pub target_id: String, + pub dispatch_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sequence: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum DeliveryResult { + Acked, + Done, + Failed, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct DeliveryResultMessage { + pub delivery: DeliveryRef, + pub result: DeliveryResult, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +pub fn delivery_result_subject(machine_id: &str) -> String { + format!("machine.{}.{}", machine_id, DELIVERY_RESULT_KIND) +} diff --git a/clients/openframe-client/src/models/execution.rs b/clients/openframe-client/src/models/execution.rs index 8214e1118c..43505f6fab 100644 --- a/clients/openframe-client/src/models/execution.rs +++ b/clients/openframe-client/src/models/execution.rs @@ -172,8 +172,8 @@ pub struct ExecutionAck { pub const EXECUTION_ACK_KIND: &str = "execution.acknowledge"; -pub fn is_ack_subject(subject: &str) -> bool { - subject.ends_with(EXECUTION_ACK_KIND) +pub fn requires_stream_ack(subject: &str) -> bool { + subject.ends_with(EXECUTION_ACK_KIND) || subject.ends_with(super::DELIVERY_RESULT_KIND) } pub struct ExecutionRequest<'a> { diff --git a/clients/openframe-client/src/models/feature_flags.rs b/clients/openframe-client/src/models/feature_flags.rs new file mode 100644 index 0000000000..1d230351fa --- /dev/null +++ b/clients/openframe-client/src/models/feature_flags.rs @@ -0,0 +1,106 @@ +use std::cmp::Ordering; +use std::collections::{BTreeMap, HashSet}; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::DeliveryRef; +use crate::config::update_config::DEV_CLIENT_VERSION; +use crate::utils::version_comparator::{parse_version, precedence}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Feature { + DeliveryToolInstallation, + DeliveryClientUninstall, + DeliveryClientUpdate, +} + +impl Feature { + pub const ALL: [Feature; 3] = [ + Feature::DeliveryToolInstallation, + Feature::DeliveryClientUninstall, + Feature::DeliveryClientUpdate, + ]; + + pub fn name(self) -> &'static str { + match self { + Feature::DeliveryToolInstallation => "delivery-tool-installation", + Feature::DeliveryClientUninstall => "delivery-client-uninstall", + Feature::DeliveryClientUpdate => "delivery-client-update", + } + } + + pub fn default_enabled(self) -> bool { + match self { + Feature::DeliveryToolInstallation + | Feature::DeliveryClientUninstall + | Feature::DeliveryClientUpdate => false, + } + } +} + +#[derive(Debug, Deserialize)] +pub struct FeatureFlagsMessage { + pub flags: BTreeMap, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct FeatureFlagsState { + #[serde(default)] + pub machine_id: String, + pub sequence: i64, + pub flags: BTreeMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub awaiting_restart: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FeatureFlag { + pub enabled: bool, + pub min_version: String, +} + +impl FeatureFlag { + fn effective(&self, running_version: &str, default: bool) -> bool { + if running_version == DEV_CLIENT_VERSION { + return self.enabled; + } + match ( + parse_version(running_version), + parse_version(&self.min_version), + ) { + (Some(running), Some(min)) if precedence(&running, &min) != Ordering::Less => { + self.enabled + } + _ => default, + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct FeatureFlags { + enabled: HashSet, +} + +impl FeatureFlags { + pub fn resolve(flags: &BTreeMap, running_version: &str) -> Self { + Self { + enabled: Feature::ALL + .into_iter() + .filter(|feature| match flags.get(feature.name()) { + Some(flag) => flag.effective(running_version, feature.default_enabled()), + None => feature.default_enabled(), + }) + .collect(), + } + } + + pub fn is_enabled(&self, feature: Feature) -> bool { + self.enabled.contains(&feature) + } +} + +#[cfg(test)] +#[path = "feature_flags_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/models/feature_flags_tests.rs b/clients/openframe-client/src/models/feature_flags_tests.rs new file mode 100644 index 0000000000..36a6820548 --- /dev/null +++ b/clients/openframe-client/src/models/feature_flags_tests.rs @@ -0,0 +1,209 @@ +use super::*; + +const TOOL_INSTALLATION: &str = "delivery-tool-installation"; + +fn flag(value: bool, min_version: &str) -> FeatureFlag { + FeatureFlag { + enabled: value, + min_version: min_version.to_string(), + } +} + +fn resolve(entries: &[(&str, FeatureFlag)], running_version: &str) -> FeatureFlags { + let flags = entries + .iter() + .map(|(name, flag)| (name.to_string(), flag.clone())) + .collect(); + FeatureFlags::resolve(&flags, running_version) +} + +#[test] +fn a_flag_applies_only_from_its_min_version() { + let cases = [ + (true, "1.5.1", "1.5.1", true, "same version"), + (true, "1.5.1", "1.5.2", true, "newer patch"), + (true, "1.5.1", "1.6.0", true, "newer minor"), + (true, "1.5.1", "2.0.0", true, "newer major"), + (true, "1.5.1", "1.5.0", false, "older patch"), + (true, "1.5.1", "1.4.9", false, "older minor"), + (true, "1.5.1", "0.9.9", false, "older major"), + ( + false, + "1.5.1", + "2.0.0", + false, + "switched off on a supporting version", + ), + ( + false, + "1.5.1", + "1.0.0", + false, + "switched off below the min version", + ), + ( + true, + "1.5.1", + "1.5.1-rc.1", + false, + "a pre-release is below its release", + ), + ( + true, + "1.5.1-rc.1", + "1.5.1-rc.2", + true, + "a later pre-release", + ), + ( + true, + "1.5.1-rc.2", + "1.5.1-rc.1", + false, + "an earlier pre-release", + ), + ( + true, + "1.5.1-rc.1", + "1.5.1", + true, + "a release is above its pre-release", + ), + ( + true, + "1.5.1", + "1.5.1+build.7", + true, + "build metadata carries no order", + ), + ( + true, + "v1.5.1", + "1.5.1", + true, + "a v prefix on the min version", + ), + ( + true, + " 1.5.1 ", + "1.5.1", + true, + "whitespace around the min version", + ), + ( + true, + "latest", + "2.0.0", + false, + "a min version that is not semver", + ), + (true, "", "2.0.0", false, "an empty min version"), + ( + true, + "1.5.1", + "not-a-version", + false, + "a running version that is not semver", + ), + ]; + for (value, min_version, running, expected, why) in cases { + let flags = resolve(&[(TOOL_INSTALLATION, flag(value, min_version))], running); + + assert_eq!( + flags.is_enabled(Feature::DeliveryToolInstallation), + expected, + "{why}: value={value} minVersion={min_version:?} running={running:?}" + ); + } +} + +#[test] +fn the_dev_build_takes_the_value_whatever_the_min_version() { + let cases = [ + (true, "99.0.0", true), + (true, "latest", true), + (true, "0.0.1", true), + (false, "0.0.1", false), + (false, "99.0.0", false), + ]; + for (value, min_version, expected) in cases { + let flags = resolve( + &[(TOOL_INSTALLATION, flag(value, min_version))], + DEV_CLIENT_VERSION, + ); + + assert_eq!( + flags.is_enabled(Feature::DeliveryToolInstallation), + expected, + "value={value} minVersion={min_version:?}" + ); + } +} + +#[test] +fn a_flag_that_does_not_apply_falls_back_to_the_feature_default() { + let below = flag(false, "2.0.0"); + let above = flag(false, "1.0.0"); + + assert!(below.effective("1.5.1", true)); + assert!(!below.effective("1.5.1", false)); + assert!(!above.effective("1.5.1", true)); +} + +#[test] +fn a_missing_flag_resolves_to_the_default() { + let flags = resolve(&[], "1.5.1"); + + for feature in Feature::ALL { + assert_eq!( + flags.is_enabled(feature), + feature.default_enabled(), + "{}", + feature.name() + ); + } +} + +#[test] +fn each_backend_name_switches_only_its_own_feature() { + for feature in Feature::ALL { + let flags = resolve(&[(feature.name(), flag(true, "1.0.0"))], "1.5.1"); + + for other in Feature::ALL { + assert_eq!( + flags.is_enabled(other), + other == feature, + "{} set, {} checked", + feature.name(), + other.name() + ); + } + } +} + +#[test] +fn unknown_flags_do_not_touch_known_features() { + let with_unknown = resolve( + &[ + ("future-feature", flag(true, "1.0.0")), + (TOOL_INSTALLATION, flag(true, "1.0.0")), + ], + "1.5.1", + ); + let without = resolve(&[(TOOL_INSTALLATION, flag(true, "1.0.0"))], "1.5.1"); + + assert_eq!(with_unknown, without); +} + +#[test] +fn flag_entries_use_the_backend_field_names() { + let parsed: FeatureFlag = + serde_json::from_value(serde_json::json!({ "enabled": true, "minVersion": "1.5.1" })) + .unwrap(); + + assert_eq!(parsed, flag(true, "1.5.1")); + assert_eq!( + serde_json::to_value(&parsed).unwrap(), + serde_json::json!({ "enabled": true, "minVersion": "1.5.1" }) + ); +} diff --git a/clients/openframe-client/src/models/mod.rs b/clients/openframe-client/src/models/mod.rs index 28ce7f07a9..6e22ac7240 100644 --- a/clients/openframe-client/src/models/mod.rs +++ b/clients/openframe-client/src/models/mod.rs @@ -2,9 +2,11 @@ pub mod agent_configuration; pub mod agent_registration_request; pub mod agent_registration_response; pub mod agent_token_response; +pub mod delivery; pub mod device_tag; pub mod download_configuration; pub mod execution; +pub mod feature_flags; pub mod hostname_report_message; pub mod initial_configuration; pub mod installed_agent_message; @@ -24,7 +26,7 @@ pub mod tool_version_overrides; pub mod update_state; pub use execution::{ - is_ack_subject, BootstrapRoute, BootstrapScriptMessage, CommandMessage, ExecutionAck, + requires_stream_ack, BootstrapRoute, BootstrapScriptMessage, CommandMessage, ExecutionAck, ExecutionMessage, ExecutionRequest, PrivilegeLevel, PrivilegePolicy, RmmResult, RoutedScriptMessage, ScriptEnvVar, ScriptMessage, ScriptRoute, ScriptScheduleExecutionItem, ScriptScheduleExecutionMessage, ScriptShell, ScriptSpec, SoftwareRoute, SoftwareScriptMessage, @@ -35,8 +37,15 @@ pub use agent_configuration::AgentConfiguration; pub use agent_registration_request::AgentRegistrationRequest; pub use agent_registration_response::AgentRegistrationResponse; pub use agent_token_response::AgentTokenResponse; +pub use delivery::{ + delivery_result_subject, DeliveryRef, DeliveryResult, DeliveryResultMessage, + DELIVERY_RESULT_KIND, +}; pub use device_tag::DeviceTag; pub use download_configuration::{DownloadConfiguration, InstallationType}; +pub use feature_flags::{ + Feature, FeatureFlag, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState, +}; pub use hostname_report_message::HostnameReportMessage; pub use initial_configuration::InitialConfiguration; pub use installed_agent_message::InstalledAgentMessage; diff --git a/clients/openframe-client/src/platform/detached_process.rs b/clients/openframe-client/src/platform/detached_process.rs new file mode 100644 index 0000000000..e08593a188 --- /dev/null +++ b/clients/openframe-client/src/platform/detached_process.rs @@ -0,0 +1,69 @@ +use anyhow::Result; +use std::path::Path; +use std::process::Child; + +/// Spawn ` ` detached, so it survives the service manager stopping +/// our own `com.openframe.client` service. macOS + Windows only. +pub fn spawn_detached(install_path: &Path, subcommand: &str) -> Result { + #[cfg(target_os = "macos")] + { + spawn_detached_macos(install_path, subcommand) + } + #[cfg(target_os = "windows")] + { + spawn_detached_windows(install_path, subcommand) + } + #[cfg(not(any(target_os = "macos", target_os = "windows")))] + { + let _ = (install_path, subcommand); + anyhow::bail!("detached client subcommands are not supported on this platform") + } +} + +/// Detached child in its own session (setsid) so it outlives launchd stopping our own +/// `com.openframe.client` service — no launchd job/plist needed. +#[cfg(target_os = "macos")] +fn spawn_detached_macos(install_path: &Path, subcommand: &str) -> Result { + use anyhow::Context; + use std::os::unix::process::CommandExt; + use std::process::{Command, Stdio}; + + let child = unsafe { + Command::new(install_path) + .arg(subcommand) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + // SAFETY: setsid() is async-signal-safe; it moves the child into a new session so + // launchd stopping our service can't signal it away mid-run. + .pre_exec(|| { + if libc::setsid() == -1 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) + }) + .spawn() + } + .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; + + Ok(child) +} + +/// Detached child (inherits LocalSystem) so it survives the SCM stopping our service. +#[cfg(target_os = "windows")] +fn spawn_detached_windows(install_path: &Path, subcommand: &str) -> Result { + use anyhow::Context; + use std::os::windows::process::CommandExt; + use std::process::Command; + + const CREATE_NO_WINDOW: u32 = 0x0800_0000; + const DETACHED_PROCESS: u32 = 0x0000_0008; + + let child = Command::new(install_path) + .arg(subcommand) + .creation_flags(CREATE_NO_WINDOW | DETACHED_PROCESS) + .spawn() + .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; + + Ok(child) +} diff --git a/clients/openframe-client/src/platform/mod.rs b/clients/openframe-client/src/platform/mod.rs index 8d50da66d8..038b8a207a 100644 --- a/clients/openframe-client/src/platform/mod.rs +++ b/clients/openframe-client/src/platform/mod.rs @@ -1,5 +1,6 @@ pub mod binary_writer; pub mod console; +pub mod detached_process; pub mod directories; pub mod dmg_extractor; pub mod file_acl; diff --git a/clients/openframe-client/src/platform/uninstall.rs b/clients/openframe-client/src/platform/uninstall.rs index 545ab2e963..990cc16c03 100644 --- a/clients/openframe-client/src/platform/uninstall.rs +++ b/clients/openframe-client/src/platform/uninstall.rs @@ -93,65 +93,8 @@ fn remove_update_temp_files(install_path: &Path) { /// Self-uninstall stops the `com.openframe.client` service (our own process), so it must run /// out-of-process. macOS + Windows only — Linux client self-uninstall is unsupported. pub fn spawn_detached_uninstall(install_path: &Path) -> Result<()> { - #[cfg(target_os = "macos")] - { - spawn_detached_uninstall_macos(install_path) - } - #[cfg(target_os = "windows")] - { - spawn_detached_uninstall_windows(install_path) - } - #[cfg(not(any(target_os = "macos", target_os = "windows")))] - { - let _ = install_path; - anyhow::bail!("client self-uninstall is not supported on this platform") - } -} - -/// Detached child in its own session (setsid) so it outlives launchd stopping our own -/// `com.openframe.client` service — no launchd job/plist needed. -#[cfg(target_os = "macos")] -fn spawn_detached_uninstall_macos(install_path: &Path) -> Result<()> { - use std::os::unix::process::CommandExt; - use std::process::{Command, Stdio}; - - let child = unsafe { - Command::new(install_path) - .arg(UNINSTALL_SUBCOMMAND) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - // SAFETY: setsid() is async-signal-safe; it moves the child into a new session so - // launchd stopping our service can't signal it away mid-uninstall. - .pre_exec(|| { - if libc::setsid() == -1 { - return Err(std::io::Error::last_os_error()); - } - Ok(()) - }) - .spawn() - } - .context("Failed to spawn detached self-uninstall process")?; - - info!("Self-uninstall process launched (PID: {})", child.id()); - Ok(()) -} - -/// Detached child (inherits LocalSystem) so it survives the SCM stopping our service. -#[cfg(target_os = "windows")] -fn spawn_detached_uninstall_windows(install_path: &Path) -> Result<()> { - use std::os::windows::process::CommandExt; - use std::process::Command; - - const CREATE_NO_WINDOW: u32 = 0x0800_0000; - const DETACHED_PROCESS: u32 = 0x0000_0008; - - let child = Command::new(install_path) - .arg(UNINSTALL_SUBCOMMAND) - .creation_flags(CREATE_NO_WINDOW | DETACHED_PROCESS) - .spawn() - .context("Failed to spawn detached self-uninstall process")?; - + let child = + crate::platform::detached_process::spawn_detached(install_path, UNINSTALL_SUBCOMMAND)?; info!("Self-uninstall process launched (PID: {})", child.id()); Ok(()) } diff --git a/clients/openframe-client/src/services/execution_service.rs b/clients/openframe-client/src/services/execution_service.rs index 669a950294..6876861aab 100644 --- a/clients/openframe-client/src/services/execution_service.rs +++ b/clients/openframe-client/src/services/execution_service.rs @@ -1,3 +1,5 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; use std::time::Instant; use crate::models::{ExecutionRequest, RmmResult}; @@ -8,14 +10,48 @@ use crate::models::PrivilegePolicy; use crate::{executor::Privilege, models::PrivilegeLevel}; #[derive(Clone, Default)] -pub struct ExecutionService; +pub struct ExecutionService { + running: Arc, +} + +struct Running(Arc); + +#[derive(Clone)] +pub struct RunningScripts(Arc); + +impl RunningScripts { + pub fn any(&self) -> bool { + self.0.load(Ordering::Acquire) > 0 + } +} + +impl Drop for Running { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} impl ExecutionService { pub fn new() -> Self { - Self + Self::default() + } + + pub fn running_scripts(&self) -> RunningScripts { + RunningScripts(self.running.clone()) + } + + fn track(&self) -> Running { + self.running.fetch_add(1, Ordering::AcqRel); + Running(self.running.clone()) + } + + #[cfg(test)] + pub(crate) fn running_for_tests(&self) -> impl Drop { + self.track() } pub async fn execute(&self, req: &ExecutionRequest<'_>, machine_id: &str) -> RmmResult { + let _running = self.track(); let start = Instant::now(); #[cfg(any(unix, windows))] diff --git a/clients/openframe-client/src/services/execution_service_tests.rs b/clients/openframe-client/src/services/execution_service_tests.rs index f2a5ee2de1..9dd4c4a15e 100644 --- a/clients/openframe-client/src/services/execution_service_tests.rs +++ b/clients/openframe-client/src/services/execution_service_tests.rs @@ -47,3 +47,26 @@ async fn maps_spawn_failure_to_error() { assert!(r.error.is_some()); assert!(!r.timed_out); } + +#[tokio::test] +async fn a_script_counts_as_running_until_it_finishes() { + let service = ExecutionService::new(); + let running = service.running_scripts(); + assert!(!running.any()); + + let script = tokio::spawn({ + let service = service.clone(); + async move { service.execute(&req("#!/bin/sh\nsleep 1\n"), "m").await } + }); + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while !running.any() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the script never started"); + + let result = script.await.unwrap(); + assert_eq!(result.exit_code, 0); + assert!(!running.any()); +} diff --git a/clients/openframe-client/src/services/feature_flags_service.rs b/clients/openframe-client/src/services/feature_flags_service.rs new file mode 100644 index 0000000000..3406d85df8 --- /dev/null +++ b/clients/openframe-client/src/services/feature_flags_service.rs @@ -0,0 +1,136 @@ +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::{Mutex, PoisonError}; + +use anyhow::{Context, Result}; +use tracing::{info, warn}; + +use crate::models::{ + DeliveryRef, FeatureFlag, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState, +}; +use crate::utils::fs::atomic_write; +use crate::utils::version_comparator::parse_version; + +pub struct FeatureFlagsService { + path: PathBuf, + version: String, + running: FeatureFlags, + state: Mutex, +} + +#[derive(Debug, PartialEq)] +pub enum FlagsApplied { + Unchanged, + RestartRequired, +} + +impl FeatureFlagsService { + pub fn load(path: PathBuf, version: String) -> Self { + let state = match std::fs::read(&path) { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(state) => state, + Err(e) => { + warn!(path = %path.display(), "Feature flags file unreadable, all features off: {:#}", e); + FeatureFlagsState::default() + } + }, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => FeatureFlagsState::default(), + Err(e) => { + warn!(path = %path.display(), "Failed to read feature flags file, all features off: {:#}", e); + FeatureFlagsState::default() + } + }; + let running = FeatureFlags::resolve(&state.flags, &version); + info!( + version, + sequence = state.sequence, + flags = ?state.flags, + effective = ?running, + "Feature flags loaded" + ); + Self { + path, + version, + running, + state: Mutex::new(state), + } + } + + pub fn flags(&self) -> &FeatureFlags { + &self.running + } + + pub fn applied_sequence(&self, machine_id: &str) -> i64 { + let state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + if state.machine_id == machine_id { + return state.sequence; + } + info!( + applied_for = %state.machine_id, + machine_id, + "Feature flags were applied for another registration, accepting any sequence" + ); + 0 + } + + pub fn apply( + &self, + machine_id: &str, + delivery: &DeliveryRef, + sequence: i64, + message: FeatureFlagsMessage, + ) -> Result { + let mut flags = BTreeMap::new(); + for (name, value) in message.flags { + match serde_json::from_value::(value.clone()) { + Ok(flag) if parse_version(&flag.min_version).is_some() => { + flags.insert(name, flag); + } + _ => { + warn!(flag = %name, %value, "Ignoring feature flag without a boolean enabled and a semver minVersion") + } + } + } + let applied = if FeatureFlags::resolve(&flags, &self.version) == self.running { + FlagsApplied::Unchanged + } else { + FlagsApplied::RestartRequired + }; + let state = FeatureFlagsState { + machine_id: machine_id.to_string(), + sequence, + flags, + awaiting_restart: (applied == FlagsApplied::RestartRequired).then(|| delivery.clone()), + }; + self.persist(&state)?; + info!(sequence, flags = ?state.flags, ?applied, "Feature flags persisted"); + *self.state.lock().unwrap_or_else(PoisonError::into_inner) = state; + Ok(applied) + } + + pub fn awaiting_restart(&self) -> Option { + self.state + .lock() + .unwrap_or_else(PoisonError::into_inner) + .awaiting_restart + .clone() + } + + pub fn restart_reported(&self) -> Result<()> { + let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + let mut reported = state.clone(); + reported.awaiting_restart = None; + self.persist(&reported)?; + *state = reported; + Ok(()) + } + + fn persist(&self, state: &FeatureFlagsState) -> Result<()> { + atomic_write(&self.path, serde_json::to_vec_pretty(state)?) + .with_context(|| format!("failed to persist feature flags to {}", self.path.display())) + } +} + +#[cfg(test)] +#[path = "feature_flags_service_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/services/feature_flags_service_tests.rs b/clients/openframe-client/src/services/feature_flags_service_tests.rs new file mode 100644 index 0000000000..1f09834362 --- /dev/null +++ b/clients/openframe-client/src/services/feature_flags_service_tests.rs @@ -0,0 +1,388 @@ +use super::*; +use crate::config::update_config::DEV_CLIENT_VERSION; +use crate::models::Feature; +use serde_json::{json, Value}; +use tempfile::TempDir; + +const MACHINE: &str = "machine-1"; +const VERSION: &str = "1.5.1"; +const TOOL_INSTALLATION: &str = "delivery-tool-installation"; + +fn flags_path() -> (PathBuf, TempDir) { + let dir = tempfile::tempdir().unwrap(); + (dir.path().join("feature_flags.json"), dir) +} + +fn load(path: &std::path::Path, version: &str) -> FeatureFlagsService { + FeatureFlagsService::load(path.to_path_buf(), version.to_string()) +} + +fn delivery(sequence: i64) -> DeliveryRef { + DeliveryRef { + delivery_type: "CLIENT_FEATURE_FLAGS".to_string(), + target_id: "tenant".to_string(), + dispatch_id: format!("dispatch-{sequence}"), + sequence: Some(sequence), + } +} + +fn entry(value: bool, min_version: &str) -> Value { + json!({ "enabled": value, "minVersion": min_version }) +} + +fn apply(service: &FeatureFlagsService, sequence: i64, flags: Value) -> FlagsApplied { + let message: FeatureFlagsMessage = serde_json::from_value(json!({ "flags": flags })).unwrap(); + service + .apply(MACHINE, &delivery(sequence), sequence, message) + .unwrap() +} + +fn tool_installation(service: &FeatureFlagsService) -> bool { + service + .flags() + .is_enabled(Feature::DeliveryToolInstallation) +} + +fn persisted(path: &std::path::Path) -> FeatureFlagsState { + serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap() +} + +fn running_with_tool_installation_on(path: &std::path::Path) -> FeatureFlagsService { + let first = load(path, VERSION); + apply( + &first, + 10, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ); + let running = load(path, VERSION); + running.restart_reported().unwrap(); + assert!(tool_installation(&running)); + running +} + +#[test] +fn a_missing_file_means_every_feature_off() { + let (path, _dir) = flags_path(); + + let service = load(&path, VERSION); + + assert!(!tool_installation(&service)); + assert_eq!(service.applied_sequence(MACHINE), 0); + assert_eq!(service.awaiting_restart(), None); +} + +#[test] +fn an_unreadable_file_means_every_feature_off() { + for contents in [ + &b"{ not json"[..], + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": true } }"#, + br#"{ "sequence": 9 }"#, + ] { + let (path, _dir) = flags_path(); + std::fs::write(&path, contents).unwrap(); + + let service = load(&path, VERSION); + + assert!( + !tool_installation(&service), + "{}", + String::from_utf8_lossy(contents) + ); + assert_eq!(service.applied_sequence(MACHINE), 0); + } +} + +#[test] +fn turning_on_an_applicable_flag_restarts_and_takes_effect_in_the_next_process() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + let applied = apply( + &service, + 7, + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), + ); + + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(!tool_installation(&service)); + assert_eq!(service.awaiting_restart(), Some(delivery(7))); + + let restarted = load(&path, VERSION); + assert!(tool_installation(&restarted)); + assert_eq!(restarted.applied_sequence(MACHINE), 7); + assert_eq!(restarted.awaiting_restart(), Some(delivery(7))); +} + +#[test] +fn turning_off_a_running_flag_restarts() { + let (path, _dir) = flags_path(); + let service = running_with_tool_installation_on(&path); + + let applied = apply( + &service, + 11, + json!({ TOOL_INSTALLATION: entry(false, "1.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(!tool_installation(&load(&path, VERSION))); +} + +#[test] +fn changes_that_leave_the_effective_set_alone_do_not_restart() { + let cases = [ + ( + "a flag for a newer client", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-uninstall": entry(true, "2.0.0") }), + ), + ( + "the min version moved but still applies", + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), + ), + ( + "an unknown flag added", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "future-feature": entry(true, "1.0.0") }), + ), + ( + "an invalid entry for another feature", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-update": true }), + ), + ( + "a switched-off flag for a newer client", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-update": entry(false, "9.0.0") }), + ), + ( + "the same set with a newer sequence", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ), + ]; + for (why, flags) in cases { + let (path, _dir) = flags_path(); + let service = running_with_tool_installation_on(&path); + + let applied = apply(&service, 11, flags); + + assert_eq!(applied, FlagsApplied::Unchanged, "{why}"); + assert_eq!(service.awaiting_restart(), None, "{why}"); + assert_eq!(service.applied_sequence(MACHINE), 11, "{why}"); + } +} + +#[test] +fn a_flag_for_a_newer_client_is_stored_without_a_restart() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + let applied = apply( + &service, + 3, + json!({ TOOL_INSTALLATION: entry(true, "1.6.0") }), + ); + + assert_eq!(applied, FlagsApplied::Unchanged); + assert_eq!(service.awaiting_restart(), None); + assert_eq!( + persisted(&path).flags.get(TOOL_INSTALLATION), + Some(&FeatureFlag { + enabled: true, + min_version: "1.6.0".to_string() + }) + ); +} + +#[test] +fn the_stored_flag_turns_on_when_the_client_is_updated() { + let (path, _dir) = flags_path(); + let before_update = load(&path, "1.5.9"); + apply( + &before_update, + 3, + json!({ TOOL_INSTALLATION: entry(true, "1.6.0") }), + ); + + let after_update = load(&path, "1.6.0"); + + assert!(tool_installation(&after_update)); + assert_eq!(after_update.awaiting_restart(), None); +} + +#[test] +fn a_rollback_below_the_min_version_turns_the_flag_off() { + let (path, _dir) = flags_path(); + running_with_tool_installation_on(&path); + + let rolled_back = load(&path, "0.9.0"); + + assert!(!tool_installation(&rolled_back)); +} + +#[test] +fn the_dev_build_turns_on_a_flag_whatever_its_min_version() { + let (path, _dir) = flags_path(); + let service = load(&path, DEV_CLIENT_VERSION); + + let applied = apply( + &service, + 1, + json!({ TOOL_INSTALLATION: entry(true, "99.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(tool_installation(&load(&path, DEV_CLIENT_VERSION))); + assert!(!tool_installation(&load(&path, VERSION))); +} + +#[test] +fn invalid_entries_are_dropped_and_resolve_to_the_default() { + let cases = [ + ("a plain boolean", json!(true)), + ("no minVersion", json!({ "enabled": true })), + ("no value", json!({ "minVersion": "1.0.0" })), + ( + "a string value", + json!({ "enabled": "true", "minVersion": "1.0.0" }), + ), + ( + "a numeric minVersion", + json!({ "enabled": true, "minVersion": 1 }), + ), + ( + "a non-semver minVersion", + json!({ "enabled": true, "minVersion": "latest" }), + ), + ("null", json!(null)), + ("a string", json!("on")), + ]; + for (why, invalid) in cases { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + let applied = apply( + &service, + 1, + json!({ TOOL_INSTALLATION: invalid, "delivery-client-uninstall": entry(false, "1.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::Unchanged, "{why}"); + let stored = persisted(&path); + assert_eq!(stored.flags.get(TOOL_INSTALLATION), None, "{why}"); + assert!( + stored.flags.contains_key("delivery-client-uninstall"), + "{why}" + ); + } +} + +#[test] +fn unknown_flags_are_kept_on_disk() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + apply( + &service, + 3, + json!({ "future-feature": entry(true, "1.0.0") }), + ); + + assert!(persisted(&path).flags.contains_key("future-feature")); +} + +#[test] +fn flipping_back_before_the_restart_cancels_it() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + apply( + &service, + 1, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ); + let applied = apply( + &service, + 2, + json!({ TOOL_INSTALLATION: entry(false, "1.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::Unchanged); + assert_eq!(service.awaiting_restart(), None); + assert_eq!(service.applied_sequence(MACHINE), 2); +} + +#[test] +fn a_reported_restart_is_not_reported_again() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + apply( + &service, + 7, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ); + + let restarted = load(&path, VERSION); + restarted.restart_reported().unwrap(); + + assert_eq!(restarted.awaiting_restart(), None); + let next_start = load(&path, VERSION); + assert_eq!(next_start.awaiting_restart(), None); + assert!(tool_installation(&next_start)); +} + +#[test] +fn a_new_registration_accepts_any_sequence_and_keeps_flags_until_then() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + apply( + &service, + 50, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ); + + let restarted = load(&path, VERSION); + + assert_eq!(restarted.applied_sequence(MACHINE), 50); + assert_eq!(restarted.applied_sequence("machine-2"), 0); + assert!(tool_installation(&restarted)); +} + +#[test] +fn a_file_without_machine_id_accepts_the_first_set() { + let (path, _dir) = flags_path(); + std::fs::write( + &path, + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": { "enabled": true, "minVersion": "1.0.0" } } }"#, + ) + .unwrap(); + + let service = load(&path, VERSION); + + assert!(tool_installation(&service)); + assert_eq!(service.applied_sequence(MACHINE), 0); +} + +#[test] +fn the_persisted_state_matches_the_documented_shape() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + apply( + &service, + 42, + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), + ); + + let stored: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + stored, + json!({ + "machine_id": MACHINE, + "sequence": 42, + "flags": { TOOL_INSTALLATION: { "enabled": true, "minVersion": "1.5.1" } }, + "awaiting_restart": { + "type": "CLIENT_FEATURE_FLAGS", + "targetId": "tenant", + "dispatchId": "dispatch-42", + "sequence": 42 + } + }) + ); +} diff --git a/clients/openframe-client/src/services/mod.rs b/clients/openframe-client/src/services/mod.rs index d0a33fa855..277d25ad07 100644 --- a/clients/openframe-client/src/services/mod.rs +++ b/clients/openframe-client/src/services/mod.rs @@ -6,6 +6,7 @@ pub mod deregistration_service; pub mod device_data_fetcher; pub mod encryption_service; pub mod execution_service; +pub mod feature_flags_service; pub mod github_download_service; pub mod hostname_report_publisher; pub mod initial_authentication_processor; diff --git a/clients/openframe-client/src/services/nats_connection_manager.rs b/clients/openframe-client/src/services/nats_connection_manager.rs index 549ceb1dac..53dd87db52 100644 --- a/clients/openframe-client/src/services/nats_connection_manager.rs +++ b/clients/openframe-client/src/services/nats_connection_manager.rs @@ -330,6 +330,11 @@ impl NatsConnectionManager { Ok(format!("{}/ws/nats?authorization={}", host, token)) } + #[cfg(test)] + pub(crate) fn use_client_for_tests(&self, client: Client) { + self.client.send_replace(Some(Arc::new(client))); + } + pub async fn get_client(&self) -> Result> { self.client .borrow() diff --git a/clients/openframe-client/src/services/openframe_client_update_service.rs b/clients/openframe-client/src/services/openframe_client_update_service.rs index 46921d74ec..784f0b6de7 100644 --- a/clients/openframe-client/src/services/openframe_client_update_service.rs +++ b/clients/openframe-client/src/services/openframe_client_update_service.rs @@ -78,7 +78,7 @@ impl OpenFrameClientUpdateService { let requested_version = message.version.trim(); info!("Received update request for version: {}", requested_version); - self.tool_run_manager.mark_client_update_pending().await; + self.tool_run_manager.mark_client_update_pending(); if self.tool_ops.any_busy() { warn!("Tool operation in progress, deferring client update to version {} (will redeliver)", requested_version); @@ -123,7 +123,7 @@ impl OpenFrameClientUpdateService { "Already running version {}, ignoring update request", canonical_version ); - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); return Ok(()); } @@ -148,7 +148,7 @@ impl OpenFrameClientUpdateService { "refusing downgrade to {} — anchored at {}", requested_version, anchor ); - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); return Ok(()); } Ok(_) => {} @@ -262,7 +262,7 @@ impl OpenFrameClientUpdateService { binary_bytes.len() ); - self.tool_run_manager.mark_client_update_pending().await; + self.tool_run_manager.mark_client_update_pending(); // 3. Extract binary update_state.set_phase(UpdatePhase::Extracting); @@ -433,7 +433,7 @@ impl OpenFrameClientUpdateService { ); } - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); if let Err(e) = self.tool_run_manager.resume_after_update_failure().await { warn!( "Failed to resume tool supervision after the aborted update to {}: {:#}", diff --git a/clients/openframe-client/src/services/result_outbox_run_manager.rs b/clients/openframe-client/src/services/result_outbox_run_manager.rs index 3fcde33a62..7155495149 100644 --- a/clients/openframe-client/src/services/result_outbox_run_manager.rs +++ b/clients/openframe-client/src/services/result_outbox_run_manager.rs @@ -5,7 +5,7 @@ use tokio::time::{interval, timeout, Duration}; use tracing::{error, info, warn}; use crate::config::update_config::{FLUSH_PUBLISH_TIMEOUT_SECS, OUTBOX_FLUSH_INTERVAL_SECS}; -use crate::models::is_ack_subject; +use crate::models::requires_stream_ack; use crate::services::result_store::{payload_limit, ResultPublisher, ResultStore}; pub struct ResultOutboxRunManager { @@ -76,7 +76,7 @@ impl ResultOutboxRunManager

{ continue; } - let outcome = if is_ack_subject(&subject) { + let outcome = if requires_stream_ack(&subject) { timeout(publish_timeout, publisher.publish_acked(&subject, &bytes)).await } else { timeout(publish_timeout, publisher.publish_raw(&subject, &bytes)).await diff --git a/clients/openframe-client/src/services/tool_ops.rs b/clients/openframe-client/src/services/tool_ops.rs index c35d48836a..7b81b6bd2e 100644 --- a/clients/openframe-client/src/services/tool_ops.rs +++ b/clients/openframe-client/src/services/tool_ops.rs @@ -18,6 +18,20 @@ struct OpState { #[derive(Debug)] pub struct Busy; +#[derive(Clone)] +pub struct BusyTools(Arc>); + +impl BusyTools { + pub fn any(&self) -> bool { + !self + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .busy + .is_empty() + } +} + pub struct ToolLock { ops: ToolOps, tool_agent_id: String, @@ -54,12 +68,11 @@ impl ToolOps { } pub fn any_busy(&self) -> bool { - !self - .state - .lock() - .unwrap_or_else(PoisonError::into_inner) - .busy - .is_empty() + self.busy_tools().any() + } + + pub fn busy_tools(&self) -> BusyTools { + BusyTools(self.state.clone()) } pub(crate) fn is_released(&self, tool_agent_id: &str) -> bool { diff --git a/clients/openframe-client/src/services/tool_ops_tests.rs b/clients/openframe-client/src/services/tool_ops_tests.rs index 71b8f92108..85449f25f3 100644 --- a/clients/openframe-client/src/services/tool_ops_tests.rs +++ b/clients/openframe-client/src/services/tool_ops_tests.rs @@ -59,6 +59,18 @@ async fn any_busy_follows_every_tool() { assert!(!ops.any_busy()); } +#[tokio::test] +async fn a_busy_tools_view_follows_ops_started_after_it() { + let ops = ToolOps::default(); + let view = ops.busy_tools(); + + let op = ops.lock(TOOL).await.mark_busy(); + assert!(view.any()); + + drop(op); + assert!(!view.any()); +} + #[tokio::test] async fn waiting_lock_sees_the_flag_cleared() { let ops = ToolOps::default(); diff --git a/clients/openframe-client/src/services/tool_run_manager.rs b/clients/openframe-client/src/services/tool_run_manager.rs index b1a9b6ea51..7b83d1e06f 100644 --- a/clients/openframe-client/src/services/tool_run_manager.rs +++ b/clients/openframe-client/src/services/tool_run_manager.rs @@ -69,20 +69,41 @@ async fn shutdown_break( #[derive(Clone, Default)] pub(crate) struct ClientUpdatePendingFlag { - since: Arc>>, + since: Arc>>, } impl ClientUpdatePendingFlag { - pub(crate) async fn mark(&self) { - *self.since.write().await = Some(std::time::Instant::now()); + pub(crate) fn mark(&self) { + *self.since() = Some(std::time::Instant::now()); } - pub(crate) async fn is_pending(&self, ttl: Duration) -> bool { - matches!(*self.since.read().await, Some(since) if since.elapsed() < ttl) + pub(crate) fn is_pending(&self, ttl: Duration) -> bool { + matches!(*self.since(), Some(since) if since.elapsed() < ttl) } - pub(crate) async fn clear(&self) { - *self.since.write().await = None; + pub(crate) fn clear(&self) { + *self.since() = None; + } + + pub(crate) fn view(&self) -> ClientUpdatePending { + ClientUpdatePending(self.clone()) + } + + fn since(&self) -> std::sync::MutexGuard<'_, Option> { + self.since + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +#[derive(Clone)] +pub struct ClientUpdatePending(ClientUpdatePendingFlag); + +impl ClientUpdatePending { + pub fn is_pending(&self) -> bool { + self.0.is_pending(Duration::from_secs( + crate::config::update_config::CLIENT_UPDATE_PENDING_TTL_SECS, + )) } } @@ -212,21 +233,21 @@ impl ToolRunManager { Ok(()) } - pub async fn mark_client_update_pending(&self) { - self.client_update_pending.mark().await; + pub fn mark_client_update_pending(&self) { + self.client_update_pending.mark(); info!("Client update pending: new tool operations will be parked"); } - pub async fn is_client_update_pending(&self) -> bool { - self.client_update_pending - .is_pending(Duration::from_secs( - crate::config::update_config::CLIENT_UPDATE_PENDING_TTL_SECS, - )) - .await + pub fn is_client_update_pending(&self) -> bool { + self.client_update_pending.view().is_pending() + } + + pub fn client_update_pending(&self) -> ClientUpdatePending { + self.client_update_pending.view() } - pub async fn clear_client_update_pending(&self) { - self.client_update_pending.clear().await; + pub fn clear_client_update_pending(&self) { + self.client_update_pending.clear(); info!("Client update no longer pending: parked tool operations released"); } diff --git a/clients/openframe-client/src/services/tool_run_manager_tests.rs b/clients/openframe-client/src/services/tool_run_manager_tests.rs index 60fd97a0bd..ac088c677e 100644 --- a/clients/openframe-client/src/services/tool_run_manager_tests.rs +++ b/clients/openframe-client/src/services/tool_run_manager_tests.rs @@ -7,42 +7,42 @@ use tokio::sync::RwLock; const LONG_TTL: Duration = Duration::from_secs(3600); -#[tokio::test] -async fn not_pending_before_first_mark() { +#[test] +fn not_pending_before_first_mark() { let flag = ClientUpdatePendingFlag::default(); - assert!(!flag.is_pending(LONG_TTL).await); + assert!(!flag.is_pending(LONG_TTL)); } -#[tokio::test] -async fn pending_after_mark_within_ttl() { +#[test] +fn pending_after_mark_within_ttl() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(flag.is_pending(LONG_TTL).await); + flag.mark(); + assert!(flag.is_pending(LONG_TTL)); } -#[tokio::test] -async fn expired_when_ttl_elapsed() { +#[test] +fn expired_when_ttl_elapsed() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(!flag.is_pending(Duration::ZERO).await); + flag.mark(); + assert!(!flag.is_pending(Duration::ZERO)); } #[tokio::test] async fn remark_refreshes_the_ttl() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; + flag.mark(); tokio::time::sleep(Duration::from_millis(30)).await; - assert!(!flag.is_pending(Duration::from_millis(10)).await); - flag.mark().await; - assert!(flag.is_pending(Duration::from_millis(10)).await); + assert!(!flag.is_pending(Duration::from_millis(10))); + flag.mark(); + assert!(flag.is_pending(Duration::from_millis(10))); } -#[tokio::test] -async fn clones_share_state() { +#[test] +fn clones_share_state() { let flag = ClientUpdatePendingFlag::default(); let clone = flag.clone(); - clone.mark().await; - assert!(flag.is_pending(LONG_TTL).await); + clone.mark(); + assert!(flag.is_pending(LONG_TTL)); } #[tokio::test] @@ -57,13 +57,26 @@ async fn shutdown_break_drops_supervision_only_once_signalled() { assert!(!running_tools.read().await.contains("mesh")); } -#[tokio::test] -async fn clear_releases_the_flag() { +#[test] +fn clear_releases_the_flag() { + let flag = ClientUpdatePendingFlag::default(); + flag.mark(); + assert!(flag.is_pending(LONG_TTL)); + flag.clear(); + assert!(!flag.is_pending(LONG_TTL)); +} + +#[test] +fn the_view_follows_the_flag() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(flag.is_pending(LONG_TTL).await); - flag.clear().await; - assert!(!flag.is_pending(LONG_TTL).await); + let view = flag.view(); + assert!(!view.is_pending()); + + flag.mark(); + assert!(view.is_pending()); + + flag.clear(); + assert!(!view.is_pending()); } fn gui_app_record(state: ToolRecordState) -> InstalledTool { diff --git a/clients/openframe-client/src/test_support.rs b/clients/openframe-client/src/test_support.rs new file mode 100644 index 0000000000..2fb6882f9e --- /dev/null +++ b/clients/openframe-client/src/test_support.rs @@ -0,0 +1,68 @@ +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +use tempfile::TempDir; + +const READY_BOUND: Duration = Duration::from_secs(10); + +pub(crate) struct NatsServer { + child: Child, + port: u16, + _store: TempDir, +} + +impl NatsServer { + pub(crate) fn start() -> Option { + let binary = std::env::var("NATS_SERVER_BIN").unwrap_or_else(|_| "nats-server".to_string()); + let port = TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .expect("no free local port") + .port(); + let store = tempfile::tempdir().unwrap(); + + let spawned = Command::new(&binary) + .args(["-a", "127.0.0.1", "-p", &port.to_string(), "-js", "-sd"]) + .arg(store.path()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn(); + let child = match spawned { + Ok(child) => child, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + if std::env::var("REQUIRE_NATS_SERVER").as_deref() == Ok("1") { + panic!("REQUIRE_NATS_SERVER=1 but `{binary}` is not installed"); + } + eprintln!("skipping: `{binary}` is not installed (brew install nats-server)"); + return None; + } + Err(e) => panic!("failed to start `{binary}`: {e}"), + }; + + let server = Self { + child, + port, + _store: store, + }; + let started = Instant::now(); + while TcpStream::connect(("127.0.0.1", port)).is_err() { + assert!( + started.elapsed() < READY_BOUND, + "nats-server did not start listening on {port}" + ); + std::thread::sleep(Duration::from_millis(20)); + } + Some(server) + } + + pub(crate) fn url(&self) -> String { + format!("nats://127.0.0.1:{}", self.port) + } +} + +impl Drop for NatsServer { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +}