634 634 | } else {
|
635 635 | // This is the default behavior.
|
636 636 | // Don't retry timeouts for IMDS, or else it will take ~30 seconds for the default
|
637 637 | // credentials provider chain to fail to provide credentials.
|
638 638 | // Also don't retry non-responses.
|
639 639 | RetryAction::NoActionIndicated
|
640 640 | }
|
641 641 | }
|
642 642 | }
|
643 643 |
|
644 - | #[cfg(test)]
|
645 - | pub(crate) mod test {
|
646 - | use crate::imds::client::{Client, EndpointMode, ImdsResponseRetryClassifier};
|
644 + | /// Utilities for testing IMDS clients
|
645 + | #[cfg(feature = "test-util")]
|
646 + | pub mod test_util {
|
647 647 | use crate::provider_config::ProviderConfig;
|
648 - | use aws_smithy_async::rt::sleep::TokioSleep;
|
649 - | use aws_smithy_async::test_util::{instant_time_and_sleep, InstantSleep};
|
650 - | use aws_smithy_http_client::test_util::{capture_request, ReplayEvent, StaticReplayClient};
|
651 - | use aws_smithy_runtime::test_util::capture_test_logs::capture_test_logs;
|
652 - | use aws_smithy_runtime_api::client::interceptors::context::{
|
653 - | Input, InterceptorContext, Output,
|
654 - | };
|
655 - | use aws_smithy_runtime_api::client::orchestrator::{
|
656 - | HttpRequest, HttpResponse, OrchestratorError,
|
657 - | };
|
658 - | use aws_smithy_runtime_api::client::result::ConnectorError;
|
659 - | use aws_smithy_runtime_api::client::retries::classifiers::{
|
660 - | ClassifyRetry, RetryAction, SharedRetryClassifier,
|
661 - | };
|
648 + | use aws_smithy_async::test_util::InstantSleep;
|
649 + | use aws_smithy_http_client::test_util::StaticReplayClient;
|
650 + | use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse};
|
662 651 | use aws_smithy_types::body::SdkBody;
|
663 - | use aws_smithy_types::error::display::DisplayErrorContext;
|
664 - | use aws_types::os_shim_internal::{Env, Fs};
|
665 - | use http::header::USER_AGENT;
|
666 652 | use http::Uri;
|
667 - | use serde::Deserialize;
|
668 - | use std::collections::HashMap;
|
669 - | use std::error::Error;
|
670 - | use std::io;
|
671 - | use std::time::SystemTime;
|
672 - | use std::time::{Duration, UNIX_EPOCH};
|
673 - | use tracing_test::traced_test;
|
674 - |
|
675 - | macro_rules! assert_full_error_contains {
|
676 - | ($err:expr, $contains:expr) => {
|
677 - | let err = $err;
|
678 - | let message = format!(
|
679 - | "{}",
|
680 - | aws_smithy_types::error::display::DisplayErrorContext(&err)
|
681 - | );
|
682 - | assert!(
|
683 - | message.contains($contains),
|
684 - | "Error message '{message}' didn't contain text '{}'",
|
685 - | $contains
|
686 - | );
|
687 - | };
|
688 - | }
|
689 653 |
|
690 - | const TOKEN_A: &str = "AQAEAFTNrA4eEGx0AQgJ1arIq_Cc-t4tWt3fB0Hd8RKhXlKc5ccvhg==";
|
691 - | const TOKEN_B: &str = "alternatetoken==";
|
692 - |
|
693 - | pub(crate) fn token_request(base: &str, ttl: u32) -> HttpRequest {
|
654 + | /// Create a simple token request
|
655 + | pub fn token_request(base: &str, ttl: u32) -> HttpRequest {
|
694 656 | http::Request::builder()
|
695 657 | .uri(format!("{}/latest/api/token", base))
|
696 658 | .header("x-aws-ec2-metadata-token-ttl-seconds", ttl)
|
697 659 | .method("PUT")
|
698 660 | .body(SdkBody::empty())
|
699 661 | .unwrap()
|
700 662 | .try_into()
|
701 663 | .unwrap()
|
702 664 | }
|
703 665 |
|
704 - | pub(crate) fn token_response(ttl: u32, token: &'static str) -> HttpResponse {
|
666 + | /// Create a simple token response
|
667 + | pub fn token_response(ttl: u32, token: &'static str) -> HttpResponse {
|
705 668 | HttpResponse::try_from(
|
706 669 | http::Response::builder()
|
707 670 | .status(200)
|
708 671 | .header("X-aws-ec2-metadata-token-ttl-seconds", ttl)
|
709 672 | .body(SdkBody::from(token))
|
710 673 | .unwrap(),
|
711 674 | )
|
712 675 | .unwrap()
|
713 676 | }
|
714 677 |
|
715 - | pub(crate) fn imds_request(path: &'static str, token: &str) -> HttpRequest {
|
678 + | /// Create a simple IMDS request
|
679 + | pub fn imds_request(path: &'static str, token: &str) -> HttpRequest {
|
716 680 | http::Request::builder()
|
717 681 | .uri(Uri::from_static(path))
|
718 682 | .method("GET")
|
719 683 | .header("x-aws-ec2-metadata-token", token)
|
720 684 | .body(SdkBody::empty())
|
721 685 | .unwrap()
|
722 686 | .try_into()
|
723 687 | .unwrap()
|
724 688 | }
|
725 689 |
|
726 - | pub(crate) fn imds_response(body: &'static str) -> HttpResponse {
|
690 + | /// Create a simple IMDS response
|
691 + | pub fn imds_response(body: &'static str) -> HttpResponse {
|
727 692 | HttpResponse::try_from(
|
728 693 | http::Response::builder()
|
729 694 | .status(200)
|
730 695 | .body(SdkBody::from(body))
|
731 696 | .unwrap(),
|
732 697 | )
|
733 698 | .unwrap()
|
734 699 | }
|
735 700 |
|
736 - | pub(crate) fn make_imds_client(http_client: &StaticReplayClient) -> super::Client {
|
701 + | /// Create an IMDS client with an underlying [StaticReplayClient]
|
702 + | pub fn make_imds_client(http_client: &StaticReplayClient) -> super::Client {
|
737 703 | tokio::time::pause();
|
738 704 | super::Client::builder()
|
739 705 | .configure(
|
740 706 | &ProviderConfig::no_configuration()
|
741 707 | .with_sleep_impl(InstantSleep::unlogged())
|
742 708 | .with_http_client(http_client.clone()),
|
743 709 | )
|
744 710 | .build()
|
745 711 | }
|
712 + | }
|
713 + |
|
714 + | #[cfg(all(test, feature = "test-util"))]
|
715 + | pub(crate) mod test {
|
716 + | use crate::imds::client::{Client, EndpointMode, ImdsResponseRetryClassifier};
|
717 + | use crate::provider_config::ProviderConfig;
|
718 + | use aws_smithy_async::rt::sleep::TokioSleep;
|
719 + | use aws_smithy_async::test_util::{instant_time_and_sleep, InstantSleep};
|
720 + | use aws_smithy_http_client::test_util::{capture_request, ReplayEvent, StaticReplayClient};
|
721 + | use aws_smithy_runtime::test_util::capture_test_logs::capture_test_logs;
|
722 + | use aws_smithy_runtime_api::client::interceptors::context::{
|
723 + | Input, InterceptorContext, Output,
|
724 + | };
|
725 + | use aws_smithy_runtime_api::client::orchestrator::OrchestratorError;
|
726 + | use aws_smithy_runtime_api::client::result::ConnectorError;
|
727 + | use aws_smithy_runtime_api::client::retries::classifiers::{
|
728 + | ClassifyRetry, RetryAction, SharedRetryClassifier,
|
729 + | };
|
730 + | use aws_smithy_types::body::SdkBody;
|
731 + | use aws_smithy_types::error::display::DisplayErrorContext;
|
732 + | use aws_types::os_shim_internal::{Env, Fs};
|
733 + | use http::header::USER_AGENT;
|
734 + |
|
735 + | use super::test_util::*;
|
736 + | use serde::Deserialize;
|
737 + | use std::collections::HashMap;
|
738 + | use std::error::Error;
|
739 + | use std::io;
|
740 + | use std::time::SystemTime;
|
741 + | use std::time::{Duration, UNIX_EPOCH};
|
742 + | use tracing_test::traced_test;
|
743 + |
|
744 + | macro_rules! assert_full_error_contains {
|
745 + | ($err:expr, $contains:expr) => {
|
746 + | let err = $err;
|
747 + | let message = format!(
|
748 + | "{}",
|
749 + | aws_smithy_types::error::display::DisplayErrorContext(&err)
|
750 + | );
|
751 + | assert!(
|
752 + | message.contains($contains),
|
753 + | "Error message '{message}' didn't contain text '{}'",
|
754 + | $contains
|
755 + | );
|
756 + | };
|
757 + | }
|
758 + |
|
759 + | const TOKEN_A: &str = "AQAEAFTNrA4eEGx0AQgJ1arIq_Cc-t4tWt3fB0Hd8RKhXlKc5ccvhg==";
|
760 + | const TOKEN_B: &str = "alternatetoken==";
|
746 761 |
|
747 762 | fn mock_imds_client(events: Vec<ReplayEvent>) -> (Client, StaticReplayClient) {
|
748 763 | let http_client = StaticReplayClient::new(events);
|
749 764 | let client = make_imds_client(&http_client);
|
750 765 | (client, http_client)
|
751 766 | }
|
752 767 |
|
753 768 | #[tokio::test]
|
754 769 | async fn client_caches_token() {
|
755 770 | let (client, http_client) = mock_imds_client(vec![
|