diff --git a/.github/actions/spelling/expect.txt b/.github/actions/spelling/expect.txt index 8a5c2865..6e010cc7 100644 --- a/.github/actions/spelling/expect.txt +++ b/.github/actions/spelling/expect.txt @@ -206,6 +206,7 @@ openprocess oneshot opencode opensource +parseable PERCPU pgpkey pgrep @@ -301,6 +302,7 @@ testurl tgid THH thiserror +Thu timedout timeup tlsv diff --git a/proxy_agent/src/key_keeper/key.rs b/proxy_agent/src/key_keeper/key.rs index e8c2b61d..cc850e55 100644 --- a/proxy_agent/src/key_keeper/key.rs +++ b/proxy_agent/src/key_keeper/key.rs @@ -36,9 +36,9 @@ use hyper::Uri; use proxy_agent_shared::hyper_client; use proxy_agent_shared::logger::LoggerLevel; use serde_derive::{Deserialize, Serialize}; -use std::ffi::OsString; use std::fmt::{Display, Formatter}; use std::{collections::HashMap, path::PathBuf}; +use std::{ffi::OsString, time::Duration}; const AUDIT_MODE: &str = "audit"; const ENFORCE_MODE: &str = "enforce"; @@ -729,7 +729,11 @@ impl Display for KeyAction { const STATUS_URL: &str = "/secure-channel/status"; const KEY_URL: &str = "/secure-channel/key"; +const HOST_DATE_TIME_DRIFT_MAX_AGE: Duration = Duration::from_secs(60 * 15); +/// Get the current status of the key from the secure channel. +/// This function will perform a single GET request to the secure channel status endpoint. +/// If the host time sync is stale, it will sync the host time from the response headers. pub async fn get_status(host: &str, port: u16) -> Result { let endpoint = hyper_client::HostEndpoint::new(host, port, STATUS_URL); let mut headers = HashMap::new(); @@ -737,8 +741,32 @@ pub async fn get_status(host: &str, port: u16) -> Result { hyper_client::METADATA_HEADER.to_string(), "True ".to_string(), ); - let status: KeyStatus = - hyper_client::get(&endpoint, &headers, None, None, logger::write_warning).await?; + + let request = hyper_client::build_request(Method::GET, &endpoint, &headers, None, None, None)?; + + let response = hyper_client::send_request( + &endpoint.host, + endpoint.port, + request, + logger::write_warning, + ) + .await?; + + if !response.status().is_success() { + return Err(proxy_agent_shared::error::Error::Hyper( + proxy_agent_shared::error::HyperErrorType::ServerError( + endpoint.to_string(), + response.status(), + ), + ))?; + } + + if proxy_agent_shared::misc_helpers::host_time_sync_is_stale(HOST_DATE_TIME_DRIFT_MAX_AGE) { + let host_time_synced = hyper_client::sync_host_time_from_headers(response.headers()); + logger::write(format!("Host time synced: {host_time_synced}")); + } + + let status: KeyStatus = hyper_client::read_response_body(response).await?; status.validate()?; Ok(status) diff --git a/proxy_agent_shared/src/hyper_client.rs b/proxy_agent_shared/src/hyper_client.rs index 8fe933a6..71175448 100644 --- a/proxy_agent_shared/src/hyper_client.rs +++ b/proxy_agent_shared/src/hyper_client.rs @@ -142,6 +142,30 @@ where read_response_body(response).await } +/// Try to sync host time from a response header map. +/// The function first looks for a custom date header (x-ms-azure-host-date) and +/// if not found, falls back to the standard HTTP date header. +/// if custom date header is present but invalid, it will not fall back to standard date header +/// Returns true when sync is updated successfully. +pub fn sync_host_time_from_headers(headers: &hyper::HeaderMap) -> bool { + // first try custom date header + if let Some(host_date) = headers.get(DATE_HEADER) { + if let Ok(host_date_rfc1123) = host_date.to_str() { + return misc_helpers::sync_host_utc_time_from_rfc1123_string(host_date_rfc1123); + } + } + + // fallback to standard HTTP date header + if let Some(host_date) = headers.get(hyper::header::DATE) { + if let Ok(host_date_rfc1123) = host_date.to_str() { + return misc_helpers::sync_host_utc_time_from_rfc1123_string(host_date_rfc1123); + } + } + + // return false if no valid date header found + false +} + pub async fn read_response_body( mut response: hyper::Response, ) -> Result @@ -539,7 +563,7 @@ mod tests { use crate::{ host_clients::{imds_client::ImdsClient, wire_server_client::WireServerClient}, logger::logger_manager, - server_mock, + misc_helpers, server_mock, }; use tokio_util::sync::CancellationToken; @@ -608,4 +632,86 @@ mod tests { cancellation_token.cancel(); } + + #[test] + fn sync_host_time_from_headers_tests() { + // should return false when no date headers are present + let headers = hyper::HeaderMap::new(); + assert!( + !super::sync_host_time_from_headers(&headers), + "should return false when no date headers are present" + ); + + // should return true with valid custom date header + let mut headers = hyper::HeaderMap::new(); + headers.insert( + super::DATE_HEADER, + "Wed, 23 Apr 2025 12:00:00 GMT".parse().unwrap(), + ); + assert!( + super::sync_host_time_from_headers(&headers), + "should return true with valid custom date header" + ); + + // should return true with valid standard Date header + let mut headers = hyper::HeaderMap::new(); + headers.insert( + hyper::header::DATE, + "Wed, 23 Apr 2025 12:00:00 GMT".parse().unwrap(), + ); + assert!( + super::sync_host_time_from_headers(&headers), + "should return true with valid standard Date header" + ); + + // when both headers are present but invalid, should return false without panic (to_str() succeeds but RFC1123 parse fails) + let mut headers = hyper::HeaderMap::new(); + headers.insert(super::DATE_HEADER, "not-a-valid-date".parse().unwrap()); + headers.insert(hyper::header::DATE, "also-not-valid".parse().unwrap()); + assert!( + !super::sync_host_time_from_headers(&headers), + "should return false when both headers have invalid dates" + ); + + // When the custom header is present but has an invalid date string, + // the function returns false immediately (does not fall back to standard Date header) + // because to_str() succeeds but the RFC1123 parse fails. + let mut headers = hyper::HeaderMap::new(); + headers.insert(super::DATE_HEADER, "not-a-valid-date".parse().unwrap()); + headers.insert( + hyper::header::DATE, + "Wed, 23 Apr 2025 12:00:00 GMT".parse().unwrap(), + ); + assert!( + !super::sync_host_time_from_headers(&headers), + "should return false without falling back when custom header has invalid date" + ); + + // when both headers are present and valid, should return true (sync from custom header takes precedence) + // and not fallback to standard Date header + let mut headers = hyper::HeaderMap::new(); + let custom_host_time = "Wed, 23 Apr 2025 12:00:00 GMT"; + headers.insert(super::DATE_HEADER, custom_host_time.parse().unwrap()); + headers.insert( + hyper::header::DATE, + "Thu, 24 Apr 2025 12:00:00 GMT".parse().unwrap(), + ); + // Both are valid; the function should return true + assert!( + super::sync_host_time_from_headers(&headers), + "should return true when both headers are present" + ); + // verify the sync time is from the custom header, not the standard Date header (which is 1 day later) + let sync_time = misc_helpers::parse_rfc1123_to_offset_datetime( + &misc_helpers::get_date_time_rfc1123_string(), + ) + .unwrap(); + let expected_sync_time = + misc_helpers::parse_rfc1123_to_offset_datetime(custom_host_time).unwrap(); + let diff = sync_time - expected_sync_time; + assert!( + diff < time::Duration::seconds(5), + "sync time should be close to the custom header time" + ); + } } diff --git a/proxy_agent_shared/src/misc_helpers.rs b/proxy_agent_shared/src/misc_helpers.rs index 9378c1b0..9d121c0d 100644 --- a/proxy_agent_shared/src/misc_helpers.rs +++ b/proxy_agent_shared/src/misc_helpers.rs @@ -11,6 +11,8 @@ use std::{ fs::{self, File}, path::{Path, PathBuf}, process::Command, + sync::RwLock, + time::Instant, }; use thread_id; use time::{format_description, OffsetDateTime, PrimitiveDateTime}; @@ -48,6 +50,14 @@ static RFC1123_FORMAT: std::sync::LazyLock>> = + std::sync::LazyLock::new(|| RwLock::new(None)); + pub fn get_date_time_string_with_milliseconds() -> String { let time_str = OffsetDateTime::now_utc() .format(&*ISO8601_MILLIS_FORMAT) @@ -63,11 +73,66 @@ pub fn get_date_time_string() -> String { } pub fn get_date_time_rfc1123_string() -> String { - OffsetDateTime::now_utc() + get_current_utc_time_synced() .format(&*RFC1123_FORMAT) .expect("Failed to format RFC1123 date") } +/// Update host-time sync state from a RFC1123 datetime string. +/// Returns true when sync state is updated successfully, false otherwise. +pub fn sync_host_utc_time_from_rfc1123_string(host_utc_rfc1123: &str) -> bool { + let Ok(parsed_host_utc) = parse_rfc1123_to_offset_datetime(host_utc_rfc1123) else { + return false; + }; + + let Ok(mut state) = HOST_TIME_SYNC_STATE.write() else { + return false; + }; + + *state = Some(HostTimeSyncState { + synced_host_utc: parsed_host_utc, + synced_instant: Instant::now(), + }); + true +} + +pub fn parse_rfc1123_to_offset_datetime(rfc1123_str: &str) -> Result { + PrimitiveDateTime::parse(rfc1123_str, &*RFC1123_FORMAT) + .map(|dt| dt.assume_utc()) + .map_err(|e| { + Error::ParseDateTimeStringError(format!( + "Failed to parse RFC1123 datetime string '{rfc1123_str}': {e}" + )) + }) +} + +/// Returns true when current host-time sync state is older than `max_age`. +/// If there is no host-time sync state yet, this returns true. +pub fn host_time_sync_is_stale(max_age: std::time::Duration) -> bool { + let Ok(state) = HOST_TIME_SYNC_STATE.read() else { + return true; + }; + state + .as_ref() + .is_none_or(|synced| synced.synced_instant.elapsed() > max_age) +} + +fn get_current_utc_time_synced() -> OffsetDateTime { + let Ok(state) = HOST_TIME_SYNC_STATE.read() else { + return OffsetDateTime::now_utc(); + }; + + let Some(synced) = state.as_ref() else { + return OffsetDateTime::now_utc(); + }; + + let elapsed = synced.synced_instant.elapsed(); + match time::Duration::try_from(elapsed) { + Ok(elapsed_time) => synced.synced_host_utc + elapsed_time, + Err(_) => OffsetDateTime::now_utc(), + } +} + pub fn get_date_time_unix_nano() -> i128 { OffsetDateTime::now_utc().unix_timestamp_nanos() } @@ -471,6 +536,7 @@ mod tests { use std::env; use std::fs; use std::path::PathBuf; + use std::time::Duration; #[derive(Serialize, Deserialize)] struct TestStruct { @@ -823,4 +889,37 @@ mod tests { "Should fail to parse invalid datetime string" ); } + + #[test] + fn sync_host_utc_time_from_rfc1123_string_test() { + let host_time = "Mon, 01 Jan 2024 00:00:00 GMT"; + assert!( + super::sync_host_utc_time_from_rfc1123_string(host_time), + "Expected valid host RFC1123 time to update sync state" + ); + + assert!( + !super::host_time_sync_is_stale(Duration::from_secs(3600)), + "Sync state should not be stale right after update" + ); + + std::thread::sleep(Duration::from_millis(1)); + assert!( + super::host_time_sync_is_stale(Duration::from_millis(0)), + "Sync state should be stale when max_age is zero" + ); + + // reset sync state to None for other tests + let _ = super::HOST_TIME_SYNC_STATE + .write() + .map(|mut state| *state = None); + } + + #[test] + fn sync_host_utc_time_from_rfc1123_string_invalid_input_test() { + assert!( + !super::sync_host_utc_time_from_rfc1123_string("invalid-rfc1123"), + "Expected invalid host RFC1123 time to fail" + ); + } }