Use retry logic for HTTP communication with IMDS

libazureinit/Cargo.toml

@@ -23,6 +23,8 @@ fstab = "0.4.0"
 
 [dev-dependencies]
 tempfile = "3"
+# test-util feature is needed for time::advance in unit tests
+tokio = { version = "1", features = ["full", "test-util"] }
 whoami = "1"
 
 [lib]

libazureinit/src/error.rs

@@ -47,4 +47,6 @@ pub enum Error {
         "Failed to set the user password; none of the provided backends succeeded"
     )]
     NoPasswordProvisioner,
+    #[error("A timeout error occurred")]
+    Timeout(#[from] tokio::time::error::Elapsed),
 }

libazureinit/src/imds.rs

@@ -1,15 +1,18 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use reqwest;
 use reqwest::header::HeaderMap;
 use reqwest::header::HeaderValue;
-use reqwest::Client;
+use reqwest::{Client, StatusCode};
+
+use std::time::Duration;
 
 use serde::{Deserialize, Deserializer};
 use serde_json;
 use serde_json::Value;
 
+use tokio::time::timeout;
+
 use crate::error::Error;
 
 #[derive(Debug, Deserialize, PartialEq, Clone)]
@@ -87,65 +90,110 @@ where
     }
 }
 
-pub async fn query(client: &Client) -> Result<InstanceMetadata, Error> {
-    let url = "http://169.254.169.254/metadata/instance?api-version=2021-02-01";
+// Set of StatusCodes that should be retried,
+// e.g. 400, 404, 410, 429, 500, 503.
+const RETRY_CODES: &[StatusCode] = &[
+    StatusCode::BAD_REQUEST,
+    StatusCode::NOT_FOUND,
+    StatusCode::GONE,
+    StatusCode::TOO_MANY_REQUESTS,
+    StatusCode::INTERNAL_SERVER_ERROR,
+    StatusCode::SERVICE_UNAVAILABLE,
+];
+
+static DEFAULT_IMDS_URL: &str =
+    "http://169.254.169.254/metadata/instance?api-version=2021-02-01";
+
+pub async fn query(
+    client: &Client,
+    retry_interval: Duration,
+    total_timeout: Duration,
+    input_url: Option<&str>,
+) -> Result<InstanceMetadata, Error> {
     let mut headers = HeaderMap::new();
 
+    let url = match input_url {
+        Some(url) => url,
+        None => DEFAULT_IMDS_URL,
+    };
+
     headers.insert("Metadata", HeaderValue::from_static("true"));
 
-    let request = client.get(url).headers(headers);
-    let response = request.send().await?;
+    let response = timeout(total_timeout, async {
+        loop {
+            if let Ok(response) = client
+                .get(url)
+                .timeout(Duration::from_secs(30))
+                .send()
+                .await
+            {
+                let statuscode = response.status();
 
-    if response.status().is_success() {
-        let imds_body = response.text().await?;
-        let metadata: InstanceMetadata = serde_json::from_str(&imds_body)?;
+                if statuscode.is_success() && statuscode == StatusCode::OK {
+                    return response.error_for_status();
+                }
 
-        Ok(metadata)
-    } else {
-        Err(Error::HttpStatus {
-            endpoint: url.to_owned(),
-            status: response.status(),
-        })
-    }
+                if !RETRY_CODES.contains(&statuscode) {
+                    return response.error_for_status();
+                }
+            }
+
+            tokio::time::sleep(retry_interval).await;
+        }
+    })
+    .await?;
+
+    let imds_body = response?.text().await?;
+
+    let metadata: InstanceMetadata = serde_json::from_str(&imds_body)?;
+
+    Ok(metadata)
 }
 
 #[cfg(test)]
 mod tests {
     use serde_json::json;
 
-    use super::{InstanceMetadata, OsProfile};
+    use super::{query, InstanceMetadata, OsProfile, RETRY_CODES};
+
+    use reqwest::{header, Client, StatusCode};
+    use std::time::Duration;
+    use tokio::io::AsyncWriteExt;
+    use tokio::net::TcpListener;
+    use tokio::time;
+
+    static BODY_CONTENTS: &str = r#"
+{
+  "compute": {
+    "azEnvironment": "cloud_env",
+    "customData": "",
+    "evictionPolicy": "",
+    "isHostCompatibilityLayerVm": "false",
+    "licenseType": "",
+    "location": "eastus",
+    "name": "AzTux-MinProvAgent-Test-0001",
+    "offer": "0001-com-ubuntu-server-focal",
+    "osProfile": {
+      "adminUsername": "MinProvAgentUser",
+      "computerName": "AzTux-MinProvAgent-Test-0001",
+      "disablePasswordAuthentication": "true"
+    },
+    "publicKeys": [
+      {
+        "keyData": "ssh-rsa test_key1",
+        "path": "/path/to/.ssh/authorized_keys"
+      },
+      {
+        "keyData": "ssh-rsa test_key2",
+        "path": "/path/to/.ssh/authorized_keys"
+      }
+    ]
+  }
+}"#;
 
     #[test]
     fn instance_metadata_deserialization() {
-        let file_body = r#"
-        {
-            "compute": {
-              "azEnvironment": "cloud_env",
-              "customData": "",
-              "evictionPolicy": "",
-              "isHostCompatibilityLayerVm": "false",
-              "licenseType": "",
-              "location": "eastus",
-              "name": "AzTux-MinProvAgent-Test-0001",
-              "offer": "0001-com-ubuntu-server-focal",
-              "osProfile": {
-                "adminUsername": "MinProvAgentUser",
-                "computerName": "AzTux-MinProvAgent-Test-0001",
-                "disablePasswordAuthentication": "true"
-              },
-              "publicKeys": [
-                {
-                  "keyData": "ssh-rsa test_key1",
-                  "path": "/path/to/.ssh/authorized_keys"
-                },
-                {
-                  "keyData": "ssh-rsa test_key2",
-                  "path": "/path/to/.ssh/authorized_keys"
-                }
-              ]
-            }
-        }"#
-        .to_string();
+        let file_body = BODY_CONTENTS.to_string();
 
         let metadata: InstanceMetadata =
             serde_json::from_str(&file_body).unwrap();
@@ -206,4 +254,74 @@ mod tests {
             serde_json::from_value(os_profile);
         assert!(os_profile.is_err_and(|err| err.is_data()));
     }
+
+    // Runs a test around sending via imds::query() with a given statuscode.
+    async fn run_imds_query_retry(statuscode: &StatusCode) -> bool {
+        const IMDS_HTTP_TOTAL_TIMEOUT_SEC: u64 = 5 * 60;
+        const IMDS_HTTP_PERCLIENT_TIMEOUT_SEC: u64 = 30;
+        const IMDS_HTTP_RETRY_INTERVAL_SEC: u64 = 2;
+
+        let mut default_headers = header::HeaderMap::new();
+        let user_agent =
+            header::HeaderValue::from_str("azure-init test").unwrap();
+
+        // Run a local test server that replies with simple test data.
+        let serverlistener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = serverlistener.local_addr().unwrap();
+
+        // Reply message includes the whole body in case of OK, otherwise empty data.
+        let ok_body = match statuscode {
+            &StatusCode::OK => format!("HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", statuscode.as_u16(), statuscode.to_string(), BODY_CONTENTS.len(), BODY_CONTENTS.to_string()),
+            _ => {
+                format!("HTTP/1.1 {} {}\r\n\r\n", statuscode.as_u16(), statuscode.to_string())
+            }
+        };
+
+        tokio::spawn(async move {
+            let (mut serverstream, _) = serverlistener.accept().await.unwrap();
+            serverstream.write_all(ok_body.as_bytes()).await.unwrap();
+        });
+
+        // Advance time to 5 minutes later, to prevent tests from being blocked
+        // for long time when retrying on RETRY_CODES.
+        time::pause();
+        time::advance(Duration::from_secs(IMDS_HTTP_TOTAL_TIMEOUT_SEC)).await;
+
+        default_headers.insert(header::USER_AGENT, user_agent);
+        let client = Client::builder()
+            .timeout(std::time::Duration::from_secs(
+                IMDS_HTTP_PERCLIENT_TIMEOUT_SEC,
+            ))
+            .default_headers(default_headers)
+            .build()
+            .unwrap();
+
+        let res = query(
+            &client,
+            Duration::from_secs(IMDS_HTTP_RETRY_INTERVAL_SEC),
+            Duration::from_secs(IMDS_HTTP_TOTAL_TIMEOUT_SEC),
+            Some(format!("http://{:}:{:}/", addr.ip(), addr.port()).as_str()),
+        )
+        .await;
+
+        time::resume();
+
+        res.is_ok()
+    }
+
+    #[tokio::test]
+    async fn imds_query_retry() {
+        // status codes that should succeed.
+        assert!(run_imds_query_retry(&StatusCode::OK).await);
+
+        // status codes that should be retried up to 5 minutes.
+        for rc in RETRY_CODES {
+            assert!(!run_imds_query_retry(rc).await);
+        }
+
+        // status codes that should result into immediate failures.
+        assert!(!run_imds_query_retry(&StatusCode::UNAUTHORIZED).await);
+        assert!(!run_imds_query_retry(&StatusCode::FORBIDDEN).await);
+        assert!(!run_imds_query_retry(&StatusCode::METHOD_NOT_ALLOWED).await);
+    }
 }

src/main.rs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 use std::process::ExitCode;
+use std::time::Duration;
 
 use anyhow::Context;
 use clap::Parser;
@@ -130,12 +131,22 @@ async fn provision() -> Result<(), anyhow::Error> {
         .default_headers(default_headers)
         .build()?;
 
+    let imds_http_timeout_sec: u64 = 5 * 60;
+    let imds_http_retry_interval_sec: u64 = 2;
+
     // Username can be obtained either via fetching instance metadata from IMDS
     // or mounting a local device for OVF environment file. It should not fail
     // immediately in a single failure, instead it should fall back to the other
     // mechanism. So it is not a good idea to use `?` for query() or
     // get_environment().
-    let instance_metadata = imds::query(&client).await.ok();
+    let instance_metadata = imds::query(
+        &client,
+        Duration::from_secs(imds_http_retry_interval_sec),
+        Duration::from_secs(imds_http_timeout_sec),
+        None, // default IMDS URL
+    )
+    .await
+    .ok();
 
     let environment = get_environment().ok();