GPA service to use host-date-time for signed http requests

.github/actions/spelling/expect.txt

@@ -1,4 +1,4 @@
 aab
 AAFFBB
 aarch
 abe
@@ -206,6 +206,7 @@
 oneshot
 opencode
 opensource
+parseable
 PERCPU
 pgpkey
 pgrep
@@ -301,6 +302,7 @@
 tgid
 THH
 thiserror
+Thu
 timedout
 timeup
 tlsv
@@ -376,4 +378,4 @@
 xsi
 xxxx
 xxxxxxxx
 xxxxxxxxxxx

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,16 +729,44 @@ 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<KeyStatus> {
     let endpoint = hyper_client::HostEndpoint::new(host, port, STATUS_URL);
     let mut headers = HashMap::new();
     headers.insert(
         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)

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<T>(
     mut response: hyper::Response<hyper::body::Incoming>,
 ) -> Result<T>
@@ -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"
+        );
+    }
 }

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<Vec<time::format_description::FormatI
         .expect("Invalid RFC1123 date format")
     });
 
+struct HostTimeSyncState {
+    synced_host_utc: OffsetDateTime,
+    synced_instant: Instant,
+}
+
+static HOST_TIME_SYNC_STATE: std::sync::LazyLock<RwLock<Option<HostTimeSyncState>>> =
+    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<OffsetDateTime> {
+    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"
+        );
+    }
 }