diff --git a/src/agent/onefuzz-supervisor/src/config.rs b/src/agent/onefuzz-supervisor/src/config.rs index 6237e251f..1c0a62de1 100644 --- a/src/agent/onefuzz-supervisor/src/config.rs +++ b/src/agent/onefuzz-supervisor/src/config.rs @@ -1,13 +1,13 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use anyhow::Result; use reqwest::StatusCode; use std::{ - path::Path, + path::{Path, PathBuf}, time::{Duration, Instant}, }; - -use anyhow::Result; +use tokio::fs; use url::Url; use uuid::Uuid; @@ -85,7 +85,7 @@ impl StaticConfig { } } -#[derive(Clone, Debug, Deserialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct DynamicConfig { /// Queried to get pending commands for the machine. pub commands_url: Url, @@ -97,6 +97,30 @@ pub struct DynamicConfig { pub work_queue: Url, } +impl DynamicConfig { + pub async fn save(&self) -> Result<()> { + let path = Self::save_path()?; + let data = serde_json::to_vec(&self)?; + fs::write(&path, &data).await?; + info!("saved dynamic-config: {}", path.display()); + Ok(()) + } + + pub async fn load() -> Result { + let path = Self::save_path()?; + let data = fs::read(&path).await?; + let ctx: Self = serde_json::from_slice(&data)?; + info!("loaded dynamic-config: {}", path.display()); + Ok(ctx) + } + + fn save_path() -> Result { + Ok(onefuzz::fs::onefuzz_root()? + .join("etc") + .join("dynamic-config.json")) + } +} + #[derive(Clone, Debug)] pub struct Registration { pub config: StaticConfig, @@ -140,7 +164,8 @@ impl Registration { match response { Ok(response) => { - let dynamic_config = response.json().await?; + let dynamic_config: DynamicConfig = response.json().await?; + dynamic_config.save().await?; return Ok(Self { config, dynamic_config, @@ -162,6 +187,18 @@ impl Registration { anyhow::bail!("Unable to register agent") } + pub async fn load_existing(config: StaticConfig) -> Result { + let dynamic_config = DynamicConfig::load().await?; + let machine_id = onefuzz::machine_id::get_machine_id().await?; + let mut registration = Self { + config, + dynamic_config, + machine_id, + }; + registration.renew().await?; + Ok(registration) + } + pub async fn create_managed(config: StaticConfig) -> Result { Self::create(config, true, DEFAULT_REGISTRATION_CREATE_TIMEOUT).await } @@ -186,6 +223,7 @@ impl Registration { .error_for_status()?; self.dynamic_config = response.json().await?; + self.dynamic_config.save().await?; Ok(()) } diff --git a/src/agent/onefuzz-supervisor/src/main.rs b/src/agent/onefuzz-supervisor/src/main.rs index 6f63d8df3..91949c91b 100644 --- a/src/agent/onefuzz-supervisor/src/main.rs +++ b/src/agent/onefuzz-supervisor/src/main.rs @@ -126,8 +126,11 @@ async fn run_agent(config: StaticConfig) -> Result<()> { telemetry::set_property(EventData::ScalesetId(scaleset)); } - let registration = config::Registration::create_managed(config.clone()).await?; - verbose!("created managed registration: {:?}", registration); + let registration = match config::Registration::load_existing(config.clone()).await { + Ok(registration) => registration, + Err(_) => config::Registration::create_managed(config.clone()).await?, + }; + verbose!("current registration: {:?}", registration); let coordinator = coordinator::Coordinator::new(registration.clone()).await?; verbose!("initialized coordinator");