Skip to content

Commit b1c208b

Browse files
zulinx86zulinx86
committed
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]>
1 parent 87a89b4 commit b1c208b

File tree

9 files changed

+116
-161
lines changed

9 files changed

+116
-161
lines changed

‎src/vmm/src/builder.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -987,7 +987,7 @@ pub(crate) mod tests {
987987
net_builder.build(net_config).unwrap();
988988
let net = net_builder.iter().next().unwrap();
989989
let mut mmds = Mmds::default();
990-
mmds.set_version(mmds_version).unwrap();
990+
mmds.set_version(mmds_version);
991991
net.lock().unwrap().configure_mmds_network_stack(
992992
MmdsNetworkStack::default_ipv4_addr(),
993993
Arc::new(Mutex::new(mmds)),

‎src/vmm/src/device_manager/persist.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -569,7 +569,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
569569
// If there's at least one network device having an mmds_ns, it means
570570
// that we are restoring from a version that did not persist the `MmdsVersionState`.
571571
// Init with the default.
572-
constructor_args.vm_resources.mmds_or_default();
572+
constructor_args.vm_resources.mmds_or_default()?;
573573
}
574574

575575
for net_state in &state.net_devices {

‎src/vmm/src/mmds/data_store.rs‎

Lines changed: 20 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ use crate::mmds::token::{MmdsTokenError as TokenError, TokenAuthority};
1212
/// The Mmds is the Microvm Metadata Service represented as an untyped json.
1313
#[derive(Debug)]
1414
pub struct Mmds {
15+
version: MmdsVersion,
1516
data_store: Value,
16-
// None when MMDS V1 is configured, Some for MMDS V2.
17-
token_authority: Option<TokenAuthority>,
17+
token_authority: TokenAuthority,
1818
is_initialized: bool,
1919
data_store_limit: usize,
2020
}
@@ -65,19 +65,20 @@ pub enum MmdsDatastoreError {
6565
// Used for ease of use in tests.
6666
impl Default for Mmds {
6767
fn default() -> Self {
68-
Self::default_with_limit(51200)
68+
Self::try_new(51200).unwrap()
6969
}
7070
}
7171

7272
impl Mmds {
7373
/// MMDS default instance with limit `data_store_limit`
74-
pub fn default_with_limit(data_store_limit: usize) -> Self {
75-
Mmds {
74+
pub fn try_new(data_store_limit: usize) -> Result<Self, MmdsDatastoreError> {
75+
Ok(Mmds {
76+
version: MmdsVersion::V1,
7677
data_store: Value::default(),
77-
token_authority: None,
78+
token_authority: TokenAuthority::try_new()?,
7879
is_initialized: false,
7980
data_store_limit,
80-
}
81+
})
8182
}
8283

8384
/// This method is needed to check if data store is initialized.
@@ -92,52 +93,29 @@ impl Mmds {
9293
}
9394

9495
/// Set the MMDS version.
95-
pub fn set_version(&mut self, version: MmdsVersion) -> Result<(), MmdsDatastoreError> {
96-
match version {
97-
MmdsVersion::V1 => {
98-
self.token_authority = None;
99-
Ok(())
100-
}
101-
MmdsVersion::V2 => {
102-
if self.token_authority.is_none() {
103-
self.token_authority = Some(TokenAuthority::try_new()?);
104-
}
105-
Ok(())
106-
}
107-
}
96+
pub fn set_version(&mut self, version: MmdsVersion) {
97+
self.version = version;
10898
}
10999

110-
/// Return the MMDS version by checking the token authority field.
100+
/// Get the MMDS version.
111101
pub fn version(&self) -> MmdsVersion {
112-
if self.token_authority.is_none() {
113-
MmdsVersion::V1
114-
} else {
115-
MmdsVersion::V2
116-
}
102+
self.version
117103
}
118104

119105
/// Sets the Additional Authenticated Data to be used for encryption and
120-
/// decryption of the session token when MMDS version 2 is enabled.
106+
/// decryption of the session token.
121107
pub fn set_aad(&mut self, instance_id: &str) {
122-
if let Some(ta) = self.token_authority.as_mut() {
123-
ta.set_aad(instance_id);
124-
}
108+
self.token_authority.set_aad(instance_id);
125109
}
126110

127111
/// Checks if the provided token has not expired.
128-
pub fn is_valid_token(&self, token: &str) -> Result<bool, TokenError> {
129-
self.token_authority
130-
.as_ref()
131-
.ok_or(TokenError::InvalidState)
132-
.map(|ta| ta.is_valid(token))
112+
pub fn is_valid_token(&self, token: &str) -> bool {
113+
self.token_authority.is_valid(token)
133114
}
134115

135116
/// Generate a new Mmds token using the token authority.
136117
pub fn generate_token(&mut self, ttl_seconds: u32) -> Result<String, TokenError> {
137-
self.token_authority
138-
.as_mut()
139-
.ok_or(TokenError::InvalidState)
140-
.and_then(|ta| ta.generate_token_secret(ttl_seconds))
118+
self.token_authority.generate_token_secret(ttl_seconds)
141119
}
142120

143121
/// set MMDS data store limit to `data_store_limit`
@@ -304,11 +282,11 @@ mod tests {
304282
assert_eq!(mmds.version(), MmdsVersion::V1);
305283

306284
// Test setting MMDS version to v2.
307-
mmds.set_version(MmdsVersion::V2).unwrap();
285+
mmds.set_version(MmdsVersion::V2);
308286
assert_eq!(mmds.version(), MmdsVersion::V2);
309287

310-
// Test setting MMDS version back to default.
311-
mmds.set_version(MmdsVersion::V1).unwrap();
288+
// Test setting MMDS version back to v1.
289+
mmds.set_version(MmdsVersion::V1);
312290
assert_eq!(mmds.version(), MmdsVersion::V1);
313291
}
314292

@@ -593,37 +571,4 @@ mod tests {
593571

594572
assert_eq!(mmds.get_data_str().len(), 2);
595573
}
596-
597-
#[test]
598-
fn test_is_valid() {
599-
let mut mmds = Mmds::default();
600-
// Set MMDS version to V2.
601-
mmds.set_version(MmdsVersion::V2).unwrap();
602-
assert_eq!(mmds.version(), MmdsVersion::V2);
603-
604-
assert!(!mmds.is_valid_token("aaa").unwrap());
605-
606-
mmds.token_authority = None;
607-
assert_eq!(
608-
mmds.is_valid_token("aaa").unwrap_err().to_string(),
609-
TokenError::InvalidState.to_string()
610-
)
611-
}
612-
613-
#[test]
614-
fn test_generate_token() {
615-
let mut mmds = Mmds::default();
616-
// Set MMDS version to V2.
617-
mmds.set_version(MmdsVersion::V2).unwrap();
618-
assert_eq!(mmds.version(), MmdsVersion::V2);
619-
620-
let token = mmds.generate_token(1).unwrap();
621-
assert!(mmds.is_valid_token(&token).unwrap());
622-
623-
mmds.token_authority = None;
624-
assert_eq!(
625-
mmds.generate_token(1).err().unwrap().to_string(),
626-
TokenError::InvalidState.to_string()
627-
);
628-
}
629574
}

‎src/vmm/src/mmds/mod.rs‎

Lines changed: 61 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,9 @@ pub enum VmmMmdsError {
3333
InvalidURI,
3434
/// Not allowed HTTP method.
3535
MethodNotAllowed,
36-
/// No MMDS token provided. Use `X-metadata-token` header to specify the session token.
36+
/// No MMDS token provided. Use `X-metadata-token` or `X-aws-ec2-metadata-token` header to specify the session token.
3737
NoTokenProvided,
38-
/// Token time to live value not found. Use `X-metadata-token-ttl-seconds` header to specify the token's lifetime.
38+
/// Token time to live value not found. Use `X-metadata-token-ttl-seconds` or `X-aws-ec2-metadata-token-ttl-seconds` header to specify the token's lifetime.
3939
NoTtlProvided,
4040
/// Resource not found: {0}.
4141
ResourceNotFound(String),
@@ -106,6 +106,7 @@ fn sanitize_uri(mut uri: String) -> String {
106106

107107
/// Build a response for `request` and return response based on MMDS version
108108
pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response {
109+
// Check URI is not empty
109110
let uri = request.uri().get_abs_path();
110111
if uri.is_empty() {
111112
return build_response(
@@ -118,16 +119,13 @@ pub fn convert_to_response(mmds: Arc<Mutex<Mmds>>, request: Request) -> Response
118119

119120
let mut mmds_guard = mmds.lock().expect("Poisoned lock");
120121

121-
match mmds_guard.version() {
122-
MmdsVersion::V1 => respond_to_request_mmdsv1(&mmds_guard, request),
123-
MmdsVersion::V2 => respond_to_request_mmdsv2(&mut mmds_guard, request),
124-
}
125-
}
126-
127-
fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
128-
// Allow only GET requests.
122+
// Allow only GET and PUT requests
129123
match request.method() {
130-
Method::Get => respond_to_get_request_unchecked(mmds, request),
124+
Method::Get => match mmds_guard.version() {
125+
MmdsVersion::V1 => respond_to_get_request_v1(&mmds_guard, request),
126+
MmdsVersion::V2 => respond_to_get_request_v2(&mmds_guard, request),
127+
},
128+
Method::Put => respond_to_put_request(&mut mmds_guard, request),
131129
_ => {
132130
let mut response = build_response(
133131
request.http_version(),
@@ -136,31 +134,18 @@ fn respond_to_request_mmdsv1(mmds: &Mmds, request: Request) -> Response {
136134
Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
137135
);
138136
response.allow_method(Method::Get);
137+
response.allow_method(Method::Put);
139138
response
140139
}
141140
}
142141
}
143142

144-
fn respond_to_request_mmdsv2(mmds: &mut Mmds, request: Request) -> Response {
145-
// Allow only GET and PUT requests.
146-
match request.method() {
147-
Method::Get => respond_to_get_request_checked(mmds, request),
148-
Method::Put => respond_to_put_request(mmds, request),
149-
_ => {
150-
let mut response = build_response(
151-
request.http_version(),
152-
StatusCode::MethodNotAllowed,
153-
MediaType::PlainText,
154-
Body::new(VmmMmdsError::MethodNotAllowed.to_string()),
155-
);
156-
response.allow_method(Method::Get);
157-
response.allow_method(Method::Put);
158-
response
159-
}
160-
}
143+
fn respond_to_get_request_v1(mmds: &Mmds, request: Request) -> Response {
144+
// TODO: Increments metrics that will be added in an upcoming commit.
145+
respond_to_get_request(mmds, request)
161146
}
162147

163-
fn respond_to_get_request_checked(mmds: &Mmds, request: Request) -> Response {
148+
fn respond_to_get_request_v2(mmds: &Mmds, request: Request) -> Response {
164149
// Check whether a token exists.
165150
let x_metadata_token = XMetadataToken::from(request.headers.custom_entries());
166151
let token = match x_metadata_token.0 {
@@ -178,18 +163,17 @@ fn respond_to_get_request_checked(mmds: &Mmds, request: Request) -> Response {
178163

179164
// Validate the token.
180165
match mmds.is_valid_token(&token) {
181-
Ok(true) => respond_to_get_request_unchecked(mmds, request),
182-
Ok(false) => build_response(
166+
true => respond_to_get_request(mmds, request),
167+
false => build_response(
183168
request.http_version(),
184169
StatusCode::Unauthorized,
185170
MediaType::PlainText,
186171
Body::new(VmmMmdsError::InvalidToken.to_string()),
187172
),
188-
Err(_) => unreachable!(),
189173
}
190174
}
191175

192-
fn respond_to_get_request_unchecked(mmds: &Mmds, request: Request) -> Response {
176+
fn respond_to_get_request(mmds: &Mmds, request: Request) -> Response {
193177
let uri = request.uri().get_abs_path();
194178

195179
// The data store expects a strict json path, so we need to
@@ -465,8 +449,7 @@ mod tests {
465449
// Set version to V1.
466450
mmds.lock()
467451
.expect("Poisoned lock")
468-
.set_version(MmdsVersion::V1)
469-
.unwrap();
452+
.set_version(MmdsVersion::V1);
470453
assert_eq!(
471454
mmds.lock().expect("Poisoned lock").version(),
472455
MmdsVersion::V1
@@ -494,7 +477,7 @@ mod tests {
494477
assert_eq!(actual_response, expected_response);
495478

496479
// Test not allowed HTTP Method.
497-
let not_allowed_methods = ["PUT", "PATCH"];
480+
let not_allowed_methods = ["PATCH"];
498481
for method in not_allowed_methods.iter() {
499482
let request = Request::try_from(
500483
format!("{method} http://169.254.169.255/ HTTP/1.0\r\n\r\n").as_bytes(),
@@ -506,6 +489,7 @@ mod tests {
506489
expected_response.set_content_type(MediaType::PlainText);
507490
expected_response.set_body(Body::new(VmmMmdsError::MethodNotAllowed.to_string()));
508491
expected_response.allow_method(Method::Get);
492+
expected_response.allow_method(Method::Put);
509493
let actual_response = convert_to_response(mmds.clone(), request);
510494
assert_eq!(actual_response, expected_response);
511495
}
@@ -528,12 +512,47 @@ mod tests {
528512
let actual_response = convert_to_response(mmds.clone(), request);
529513
assert_eq!(actual_response, expected_response);
530514

531-
// Test Ok path.
515+
// Test valid v1 request.
532516
let (request, expected_response) = generate_request_and_expected_response(
533517
b"GET http://169.254.169.254/ HTTP/1.0\r\n\
534518
Accept: application/json\r\n\r\n",
535519
MediaType::ApplicationJson,
536520
);
521+
let actual_response = convert_to_response(mmds.clone(), request);
522+
assert_eq!(actual_response, expected_response);
523+
524+
// Test valid v2 request.
525+
let request = Request::try_from(
526+
b"PUT http://169.254.169.254/latest/api/token HTTP/1.0\r\n\
527+
X-metadata-token-ttl-seconds: 60\r\n\r\n",
528+
None,
529+
)
530+
.unwrap();
531+
let actual_response = convert_to_response(mmds.clone(), request);
532+
assert_eq!(actual_response.status(), StatusCode::OK);
533+
assert_eq!(actual_response.content_type(), MediaType::PlainText);
534+
535+
let valid_token = String::from_utf8(actual_response.body().unwrap().body).unwrap();
536+
#[rustfmt::skip]
537+
let (request, expected_response) = generate_request_and_expected_response(
538+
format!(
539+
"GET http://169.254.169.254/ HTTP/1.0\r\n\
540+
Accept: application/json\r\n\
541+
X-metadata-token: {valid_token}\r\n\r\n",
542+
)
543+
.as_bytes(),
544+
MediaType::ApplicationJson,
545+
);
546+
let actual_response = convert_to_response(mmds.clone(), request);
547+
assert_eq!(actual_response, expected_response);
548+
549+
// Test GET request with invalid token is accepted when v1 is configured.
550+
let (request, expected_response) = generate_request_and_expected_response(
551+
b"GET http://169.254.169.254/ HTTP/1.0\r\n\
552+
Accept: application/json\r\n\
553+
X-metadata-token: INVALID_TOKEN\r\n\r\n",
554+
MediaType::ApplicationJson,
555+
);
537556
let actual_response = convert_to_response(mmds, request);
538557
assert_eq!(actual_response, expected_response);
539558
}
@@ -546,8 +565,7 @@ mod tests {
546565
// Set version to V2.
547566
mmds.lock()
548567
.expect("Poisoned lock")
549-
.set_version(MmdsVersion::V2)
550-
.unwrap();
568+
.set_version(MmdsVersion::V2);
551569
assert_eq!(
552570
mmds.lock().expect("Poisoned lock").version(),
553571
MmdsVersion::V2
@@ -844,13 +862,14 @@ mod tests {
844862

845863
assert_eq!(
846864
VmmMmdsError::NoTokenProvided.to_string(),
847-
"No MMDS token provided. Use `X-metadata-token` header to specify the session token."
865+
"No MMDS token provided. Use `X-metadata-token` or `X-aws-ec2-metadata-token` header \
866+
to specify the session token."
848867
);
849868

850869
assert_eq!(
851870
VmmMmdsError::NoTtlProvided.to_string(),
852-
"Token time to live value not found. Use `X-metadata-token-ttl-seconds` header to \
853-
specify the token's lifetime."
871+
"Token time to live value not found. Use `X-metadata-token-ttl-seconds` or \
872+
`X-aws-ec2-metadata-token-ttl-seconds` header to specify the token's lifetime."
854873
);
855874

856875
assert_eq!(

‎src/vmm/src/mmds/token.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,6 @@ pub enum MmdsTokenError {
5959
EntropyPool(#[from] io::Error),
6060
/// Failed to extract expiry value from token.
6161
ExpiryExtraction,
62-
/// Invalid token authority state.
63-
InvalidState,
6462
/// Invalid time to live value provided for token: {0}. Please provide a value between {MIN_TOKEN_TTL_SECONDS:} and {MAX_TOKEN_TTL_SECONDS:}.
6563
InvalidTtlValue(u32),
6664
/// Bincode serialization failed: {0}.

0 commit comments

Comments
 (0)