Repository: firecracker-microvm/firecracker
File: src/vmm/src/mmds/token.rs (Rust)

Code before:
}

impl Token {
    /// Create a new token struct.
    fn new(iv: [u8; IV_LEN], payload: [u8; PAYLOAD_LEN], tag: [u8; TAG_LEN]) -> Self {
        Token { iv, payload, tag }
    }

    /// Encode token structure into a string using base64 encoding.
    fn base64_encode(&self) -> Result<String, MmdsTokenError> {
        let token_bytes: Vec<u8> = bitcode::serialize(self)?;

        // Encode token structure bytes into base64.
        Ok(base64::engine::general_purpose::STANDARD.encode(token_bytes))
    }

    /// Decode token structure from base64 string.
    fn base64_decode(encoded_token: &str) -> Result<Self, MmdsTokenError> {
        let token_bytes = base64::engine::general_purpose::STANDARD
            .decode(encoded_token)
            .map_err(|_| MmdsTokenError::ExpiryExtraction)?;

        // Check size limit to prevent DOS attacks
        if token_bytes.len() > TOKEN_DESERIALIZATION_BYTES_LIMIT {
            return Err(MmdsTokenError::ExpiryExtraction);
        }

        let token: Token =
            bitcode::deserialize(&token_bytes).map_err(|_| MmdsTokenError::ExpiryExtraction)?;
        Ok(token)
    }
}

#[cfg(test)]
mod tests {
    use std::thread::sleep;
    use std::time::Duration;

    use super::*;

    #[test]
    fn test_check_tll() {
        // Test invalid time to live values.
        assert!(!TokenAuthority::check_ttl(MIN_TOKEN_TTL_SECONDS - 1));
        assert!(!TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS + 1));

        // Test time to live value within bounds.
        assert!(TokenAuthority::check_ttl(MIN_TOKEN_TTL_SECONDS));
        assert!(TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS / 2));
        assert!(TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS));
    }

Review comment (bug):
(Just as a note for myself and people who look at this PR)

As the comment says "to prevent DoS attacks", at first glance the size limit check looks to me like something we shouldn't delete. From a quick look at the history, it seems the original intent was to prevent bincode from allocating more memory than necessary.

Since a malicious guest can submit a token of arbitrary length, it does make sense to validate the size as early as possible to avoid potential memory exhaustion. In this case, the token is 36 bytes, which becomes 48 bytes when base64-encoded. Ideally, the size check should be performed as early as possible.

Looking at the current implementation, the size is already checked in `is_valid()` against `TOKEN_LENGTH_LIMIT` before `base64_decode()` is called. That suggests we don't need to repeat the same check inside `base64_decode()` again unless the serializer/deserializer allocates large memory area to perform serialization/deserialization. I guess zerocopy doesn't do such unnecessary memory allocation (please correct me if I'm wrong, but I think that's why it is named "zerocopy").

Currently `TOKEN_LENGTH_LIMIT` is set to 70, but we may be able to change it to 48 as we no longer do extra serialization/deserialization.

Code after:
    }

    /// Encode token structure into a string using base64 encoding.
    fn base64_encode(&self) -> String {
        base64::engine::general_purpose::STANDARD.encode(self.as_bytes())
    }

    /// Decode token structure from base64 string.
    fn base64_decode(encoded_token: &str) -> Result<Self, MmdsTokenError> {
        let bytes = base64::engine::general_purpose::STANDARD
            .decode(encoded_token)
            .map_err(|_| MmdsTokenError::ExpiryExtraction)?;
        Self::read_from_bytes(&bytes).map_err(|_| MmdsTokenError::ExpiryExtraction)
    }
}

#[cfg(test)]
mod tests {
    use std::thread::sleep;
    use std::time::Duration;

    use super::*;

    #[test]
    fn test_check_tll() {
        // Test invalid time to live values.
        assert!(!TokenAuthority::check_ttl(MIN_TOKEN_TTL_SECONDS - 1));
        assert!(!TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS + 1));

        // Test time to live value within bounds.
        assert!(TokenAuthority::check_ttl(MIN_TOKEN_TTL_SECONDS));
        assert!(TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS / 2));
        assert!(TokenAuthority::check_ttl(MAX_TOKEN_TTL_SECONDS));
    }

    #[test]
    fn test_set_aad() {
        let mut token_authority = TokenAuthority::try_new().unwrap();
        assert_eq!(token_authority.aad, "".to_string());

        token_authority.set_aad("foo");
        assert_eq!(token_authority.aad, "microvmid=foo".to_string());
    }

    #[test]
    fn test_create_token() {
        let mut token_authority = TokenAuthority::try_new().unwrap();

        // Test invalid time to live value.
        assert_eq!(
            token_authority.create_token(0).unwrap_err().to_string(),