feat(mmds): Support token generation via PUT in MMDS v1

To enable smoother migration from v1 to v2, make MMDS v1 support PUT
request to /latest/api/token for token generation.

We will not make MMDS v1 deny GET requests with invalid token in order
not to break existing workloads. It does not validate a given toekn at
this point, but it will validate to increment a metric that will be
added to count such invalid tokens.

Signed-off-by: Takahiro Itazuri <[email protected]>

Author: zulinx86

src/vmm/src/builder.rs

@@ -987,7 +987,7 @@ pub(crate) mod tests {
         net_builder.build(net_config).unwrap();
         let net = net_builder.iter().next().unwrap();
         let mut mmds = Mmds::default();
-        mmds.set_version(mmds_version).unwrap();
+        mmds.set_version(mmds_version);
         net.lock().unwrap().configure_mmds_network_stack(
             MmdsNetworkStack::default_ipv4_addr(),
             Arc::new(Mutex::new(mmds)),

src/vmm/src/device_manager/persist.rs

@@ -569,7 +569,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             // If there's at least one network device having an mmds_ns, it means
             // that we are restoring from a version that did not persist the `MmdsVersionState`.
             // Init with the default.
-            constructor_args.vm_resources.mmds_or_default();
+            constructor_args.vm_resources.mmds_or_default()?;
         }
 
         for net_state in &state.net_devices {

src/vmm/src/mmds/data_store.rs

@@ -12,9 +12,9 @@ use crate::mmds::token::{MmdsTokenError as TokenError, TokenAuthority};
 /// The Mmds is the Microvm Metadata Service represented as an untyped json.
 #[derive(Debug)]
 pub struct Mmds {
+    version: MmdsVersion,
     data_store: Value,
-    // None when MMDS V1 is configured, Some for MMDS V2.
-    token_authority: Option<TokenAuthority>,
+    token_authority: TokenAuthority,
     is_initialized: bool,
     data_store_limit: usize,
 }
@@ -65,19 +65,20 @@ pub enum MmdsDatastoreError {
 // Used for ease of use in tests.
 impl Default for Mmds {
     fn default() -> Self {
-        Self::default_with_limit(51200)
+        Self::try_new(51200).unwrap()
     }
 }
 
 impl Mmds {
     /// MMDS default instance with limit `data_store_limit`
-    pub fn default_with_limit(data_store_limit: usize) -> Self {
-        Mmds {
+    pub fn try_new(data_store_limit: usize) -> Result<Self, MmdsDatastoreError> {
+        Ok(Mmds {
+            version: MmdsVersion::V1,
             data_store: Value::default(),
-            token_authority: None,
+            token_authority: TokenAuthority::try_new()?,
             is_initialized: false,
             data_store_limit,
-        }
+        })
     }
 
     /// This method is needed to check if data store is initialized.
@@ -92,52 +93,29 @@ impl Mmds {
     }
 
     /// Set the MMDS version.
-    pub fn set_version(&mut self, version: MmdsVersion) -> Result<(), MmdsDatastoreError> {
-        match version {
-            MmdsVersion::V1 => {
-                self.token_authority = None;
-                Ok(())
-            }
-            MmdsVersion::V2 => {
-                if self.token_authority.is_none() {
-                    self.token_authority = Some(TokenAuthority::try_new()?);
-                }
-                Ok(())
-            }
-        }
+    pub fn set_version(&mut self, version: MmdsVersion) {
+        self.version = version;
     }
 
-    /// Return the MMDS version by checking the token authority field.
+    /// Get the MMDS version.
     pub fn version(&self) -> MmdsVersion {
-        if self.token_authority.is_none() {
-            MmdsVersion::V1
-        } else {
-            MmdsVersion::V2
-        }
+        self.version
     }
 
     /// Sets the Additional Authenticated Data to be used for encryption and
-    /// decryption of the session token when MMDS version 2 is enabled.
+    /// decryption of the session token.
     pub fn set_aad(&mut self, instance_id: &str) {
-        if let Some(ta) = self.token_authority.as_mut() {
-            ta.set_aad(instance_id);
-        }
+        self.token_authority.set_aad(instance_id);
     }
 
     /// Checks if the provided token has not expired.
-    pub fn is_valid_token(&self, token: &str) -> Result<bool, TokenError> {
-        self.token_authority
-            .as_ref()
-            .ok_or(TokenError::InvalidState)
-            .map(|ta| ta.is_valid(token))
+    pub fn is_valid_token(&self, token: &str) -> bool {
+        self.token_authority.is_valid(token)
     }
 
     /// Generate a new Mmds token using the token authority.
     pub fn generate_token(&mut self, ttl_seconds: u32) -> Result<String, TokenError> {
-        self.token_authority
-            .as_mut()
-            .ok_or(TokenError::InvalidState)
-            .and_then(|ta| ta.generate_token_secret(ttl_seconds))
+        self.token_authority.generate_token_secret(ttl_seconds)
     }
 
     /// set MMDS data store limit to `data_store_limit`
@@ -304,11 +282,11 @@ mod tests {
         assert_eq!(mmds.version(), MmdsVersion::V1);
 
         // Test setting MMDS version to v2.
-        mmds.set_version(MmdsVersion::V2).unwrap();
+        mmds.set_version(MmdsVersion::V2);
         assert_eq!(mmds.version(), MmdsVersion::V2);
 
-        // Test setting MMDS version back to default.
-        mmds.set_version(MmdsVersion::V1).unwrap();
+        // Test setting MMDS version back to v1.
+        mmds.set_version(MmdsVersion::V1);
         assert_eq!(mmds.version(), MmdsVersion::V1);
     }
 
@@ -593,37 +571,4 @@ mod tests {
 
         assert_eq!(mmds.get_data_str().len(), 2);
     }
-
-    #[test]
-    fn test_is_valid() {
-        let mut mmds = Mmds::default();
-        // Set MMDS version to V2.
-        mmds.set_version(MmdsVersion::V2).unwrap();
-        assert_eq!(mmds.version(), MmdsVersion::V2);
-
-        assert!(!mmds.is_valid_token("aaa").unwrap());
-
-        mmds.token_authority = None;
-        assert_eq!(
-            mmds.is_valid_token("aaa").unwrap_err().to_string(),
-            TokenError::InvalidState.to_string()
-        )
-    }
-
-    #[test]
-    fn test_generate_token() {
-        let mut mmds = Mmds::default();
-        // Set MMDS version to V2.
-        mmds.set_version(MmdsVersion::V2).unwrap();
-        assert_eq!(mmds.version(), MmdsVersion::V2);
-
-        let token = mmds.generate_token(1).unwrap();
-        assert!(mmds.is_valid_token(&token).unwrap());
-
-        mmds.token_authority = None;
-        assert_eq!(
-            mmds.generate_token(1).err().unwrap().to_string(),
-            TokenError::InvalidState.to_string()
-        );
-    }
 }

src/vmm/src/mmds/mod.rs

@@ -106,6 +106,7 @@ fn sanitize_uri(mut uri: String) -> String {
 
 /// Build a response for `request` and return response based on MMDS version
 pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response {
+    // Check URI is not empty
     let uri = request.uri().get_abs_path();
     if uri.is_empty() {
         return build_response(
@@ -118,16 +119,13 @@ pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response
 
     let mut mmds_guard = mmds.lock().expect("Poisoned lock");
 
-    match mmds_guard.version() {
-        MmdsVersion::V1 => respond_to_request_mmdsv1(&mmds_guard, request),
-        MmdsVersion::V2 => respond_to_request_mmdsv2(&mut mmds_guard, request),
-    }
-}
-
-fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
-    // Allow only GET requests.
+    // Allow only GET and PUT requests
     match request.method() {
-        Method::Get => respond_to_get_request_unchecked(mmds, request),
+        Method::Get => match mmds_guard.version() {
+            MmdsVersion::V1 => respond_to_get_request_v1(&mmds_guard, request),
+            MmdsVersion::V2 => respond_to_get_request_v2(&mmds_guard, request),
+        },
+        Method::Put => respond_to_put_request(&mut mmds_guard, request),
         _ => {
             let mut response = build_response(
                 request.http_version(),
@@ -136,31 +134,18 @@ fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
                 Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
             );
             response.allow_method(Method::Get);
+            response.allow_method(Method::Put);
             response
         }
     }
 }
 
-fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
-    // Allow only GET and PUT requests.
-    match request.method() {
-        Method::Get => respond_to_get_request_checked(mmds, request),
-        Method::Put => respond_to_put_request(mmds, request),
-        _ => {
-            let mut response = build_response(
-                request.http_version(),
-                StatusCode::MethodNotAllowed,
-                MediaType::PlainText,
-                Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
-            );
-            response.allow_method(Method::Get);
-            response.allow_method(Method::Put);
-            response
-        }
-    }
+fn respond_to_get_request_v1(mmds: &Mmds, request: Request) -> Response {
+    // TODO: Increments metrics that will be added in an upcoming commit.
+    respond_to_get_request(mmds, request)
 }
 
-fn respond_to_get_request_checked(mmds: &Mmds, request: Request) -> Response {
+fn respond_to_get_request_v2(mmds: &Mmds, request: Request) -> Response {
     // Check whether a token exists.
     let x_metadata_token = XMetadataToken::from(request.headers.custom_entries());
     let token = match x_metadata_token.0 {
@@ -178,18 +163,17 @@ fn respond_to_get_request_checked(mmds: &Mmds, request: Request) -> Response {
 
     // Validate the token.
     match mmds.is_valid_token(&token) {
-        Ok(true) => respond_to_get_request_unchecked(mmds, request),
-        Ok(false) => build_response(
+        true => respond_to_get_request(mmds, request),
+        false => build_response(
             request.http_version(),
             StatusCode::Unauthorized,
             MediaType::PlainText,
             Body::new(VmmMmdsError::InvalidToken.to_string()),
         ),
-        Err(_) => unreachable!(),
     }
 }
 
-fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
+fn respond_to_get_request(mmds: &Mmds, request: Request) -> Response {
     let uri = request.uri().get_abs_path();
 
     // The data store expects a strict json path, so we need to
@@ -465,8 +449,7 @@ mod tests {
         // Set version to V1.
         mmds.lock()
             .expect("Poisoned lock")
-            .set_version(MmdsVersion::V1)
-            .unwrap();
+            .set_version(MmdsVersion::V1);
         assert_eq!(
             mmds.lock().expect("Poisoned lock").version(),
             MmdsVersion::V1
@@ -494,7 +477,7 @@ mod tests {
         assert_eq!(actual_response, expected_response);
 
         // Test not allowed HTTP Method.
-        let not_allowed_methods = ["PUT", "PATCH"];
+        let not_allowed_methods = ["PATCH"];
         for method in not_allowed_methods.iter() {
             let request = Request::try_from(
                 format!("{method} http://169.254.169.255/ HTTP/1.0\r\n\r\n").as_bytes(),
@@ -506,6 +489,7 @@ mod tests {
             expected_response.set_content_type(MediaType::PlainText);
             expected_response.set_body(Body::new(VmmMmdsError::MethodNotAllowed.to_string()));
             expected_response.allow_method(Method::Get);
+            expected_response.allow_method(Method::Put);
             let actual_response = convert_to_response(mmds.clone(), request);
             assert_eq!(actual_response, expected_response);
         }
@@ -528,12 +512,47 @@ mod tests {
         let actual_response = convert_to_response(mmds.clone(), request);
         assert_eq!(actual_response, expected_response);
 
-        // Test Ok path.
+        // Test valid v1 request.
         let (request, expected_response) = generate_request_and_expected_response(
             b"GET http://169.254.169.254/ HTTP/1.0\r\n\
               Accept: application/json\r\n\r\n",
             MediaType::ApplicationJson,
         );
+        let actual_response = convert_to_response(mmds.clone(), request);
+        assert_eq!(actual_response, expected_response);
+
+        // Test valid v2 request.
+        let request = Request::try_from(
+            b"PUT http://169.254.169.254/latest/api/token HTTP/1.0\r\n\
+              X-metadata-token-ttl-seconds: 60\r\n\r\n",
+            None,
+        )
+        .unwrap();
+        let actual_response = convert_to_response(mmds.clone(), request);
+        assert_eq!(actual_response.status(), StatusCode::OK);
+        assert_eq!(actual_response.content_type(), MediaType::PlainText);
+
+        let valid_token = String::from_utf8(actual_response.body().unwrap().body).unwrap();
+        #[rustfmt::skip]
+        let (request, expected_response) = generate_request_and_expected_response(
+            format!(
+                "GET http://169.254.169.254/ HTTP/1.0\r\n\
+                 Accept: application/json\r\n\
+                 X-metadata-token: {valid_token}\r\n\r\n",
+            )
+            .as_bytes(),
+            MediaType::ApplicationJson,
+        );
+        let actual_response = convert_to_response(mmds.clone(), request);
+        assert_eq!(actual_response, expected_response);
+
+        // Test GET request with invalid token is accepted when v1 is configured.
+        let (request, expected_response) = generate_request_and_expected_response(
+            b"GET http://169.254.169.254/ HTTP/1.0\r\n\
+              Accept: application/json\r\n\
+              X-metadata-token: INVALID_TOKEN\r\n\r\n",
+            MediaType::ApplicationJson,
+        );
         let actual_response = convert_to_response(mmds, request);
         assert_eq!(actual_response, expected_response);
     }
@@ -546,8 +565,7 @@ mod tests {
         // Set version to V2.
         mmds.lock()
             .expect("Poisoned lock")
-            .set_version(MmdsVersion::V2)
-            .unwrap();
+            .set_version(MmdsVersion::V2);
         assert_eq!(
             mmds.lock().expect("Poisoned lock").version(),
             MmdsVersion::V2

src/vmm/src/mmds/token.rs

@@ -59,8 +59,6 @@ pub enum MmdsTokenError {
     EntropyPool(#[from] io::Error),
     /// Failed to extract expiry value from token.
     ExpiryExtraction,
-    /// Invalid token authority state.
-    InvalidState,
     /// Invalid time to live value provided for token: {0}. Please provide a value between {MIN_TOKEN_TTL_SECONDS:} and {MAX_TOKEN_TTL_SECONDS:}.
     InvalidTtlValue(u32),
     /// Bincode serialization failed: {0}.