Fallback to IMDSv1 when IMDSv2 token request reaches hop limit (#655)

* Use URI directly

* Fallback to IMDSv1 when token hop limit reached

* Add warning about hop limit

* Update tests to work on Julia 1.6

* Set project version to 1.90.3

Author: omus

Diff for: Project.toml

@@ -1,7 +1,7 @@
 name = "AWS"
 uuid = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
 license = "MIT"
-version = "1.90.2"
+version = "1.90.3"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

Diff for: src/IMDS.jl

@@ -15,6 +15,7 @@ using ..AWSExceptions: IMDSUnavailable
 using HTTP: HTTP
 using HTTP.Exceptions: ConnectError, StatusError
 using Mocking
+using URIs: URI
 
 # Local-link address (https://en.wikipedia.org/wiki/Link-local_address)
 const IPv4_ADDRESS = "169.254.169.254"
@@ -55,8 +56,28 @@ function refresh_token!(session::Session, duration::Integer=session.duration)
     # For IMDSv2, you must use `/latest/api/token` when retrieving the token instead of a
     # version specific path.
     # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#imds-considerations
-    uri = HTTP.URI(; scheme="http", host=IPv4_ADDRESS, path="/latest/api/token")
-    r = _http_request("PUT", uri, headers; status_exception=false)
+    uri = URI(; scheme="http", host=IPv4_ADDRESS, path="/latest/api/token")
+    r = try
+        _http_request("PUT", uri, headers; status_exception=false)
+    catch e
+        # The IMDSv2 uses a default Time To Live (TTL) of 1 (also known as the hop limit) at
+        # the IP layer to ensure token requests occur on the instance. When this occurs we
+        # need to fall back to using IMDSv1. Users may wish to increase the hop limit to
+        # allow for IMDSv2 use in container based environments:
+        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#imds-considerations
+        if is_ttl_expired_exception(e)
+            @warn "IMDSv2 token request rejected due to reaching hop limit. Consider " *
+                "increasing the hop limit to avoid delays upon initial use:\n" *
+                "https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/" *
+                "instancedata-data-retrieval.html#imds-considerations"
+
+            session.duration = 0
+            session.expiration = typemax(Int64)  # Use IMDSv1 indefinitely
+            return session
+        else
+            rethrow()
+        end
+    end
 
     # Store the session token when we receive an HTTP 200. If we receive an HTTP 404 assume
     # that the server is only supports IMDSv1. Otherwise "rethrow" the `StatusError`.
@@ -87,7 +108,7 @@ function request(session::Session, method::AbstractString, path::AbstractString;
     # Only using the IPv4 endpoint as the IPv6 endpoint has to be explicitly enabled and
     # does not disable IPv4 support.
     # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances-ipv4-ipv6-endpoints
-    uri = HTTP.URI(; scheme="http", host=IPv4_ADDRESS, path)
+    uri = URI(; scheme="http", host=IPv4_ADDRESS, path)
     return _http_request(method, uri, headers; kwargs...)
 end
 
@@ -122,10 +143,12 @@ end
 is_connection_exception(e::ConnectError) = true
 is_connection_exception(e::Exception) = false
 
+# https://github.com/JuliaCloud/AWS.jl/issues/654
 # https://github.com/JuliaCloud/AWS.jl/issues/649
-function is_connection_exception(e::HTTP.Exceptions.RequestError)
+function is_ttl_expired_exception(e::HTTP.Exceptions.RequestError)
     return e.error == Base.IOError("read: connection timed out (ETIMEDOUT)", -110)
 end
+is_ttl_expired_exception(e::Exception) = false
 
 """
     get([session::Session], path::AbstractString) -> Union{String, Nothing}

Diff for: test/IMDS.jl

@@ -82,19 +82,22 @@ function _imds_patch(router::HTTP.Router=HTTP.Router(); listening=true, enabled=
 end
 
 @testset "IMDS" begin
-    @testset "is_connection_exception" begin
+    @testset "is_connection_exception / is_ttl_expired_exception" begin
         url = "http://169.254.169.254/latest/api/token"
         connect_timeout = HTTP.ConnectionPool.ConnectTimeout("169.254.169.254", 80)
         e = HTTP.Exceptions.ConnectError(url, connect_timeout)
         @test IMDS.is_connection_exception(e)
+        @test !IMDS.is_ttl_expired_exception(e)
 
         request = HTTP.Request("PUT", "/latest/api/token", [], HTTP.nobody)
         io_error = Base.IOError("read: connection timed out (ETIMEDOUT)", -110)
         e = HTTP.Exceptions.RequestError(request, io_error)
-        @test IMDS.is_connection_exception(e)
+        @test !IMDS.is_connection_exception(e)
+        @test IMDS.is_ttl_expired_exception(e)
 
         e = ErrorException("non-connection error")
         @test !IMDS.is_connection_exception(e)
+        @test !IMDS.is_ttl_expired_exception(e)
     end
 
     @testset "refresh_token!" begin
@@ -225,6 +228,30 @@ end
             @test r isa HTTP.Response
             @test r.status == 401
         end
+
+        # When running in a container running on an EC2 instance and the hop limit is 1 the
+        # IMDSv2 token retrieval will fail so we should fall back to using IMDSv1.
+        # https://github.com/JuliaCloud/AWS.jl/issues/654
+        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#imds-considerations
+        connection_timeout = function (req::HTTP.Request)
+            io_error = Base.IOError("read: connection timed out (ETIMEDOUT)", -110)
+            throw(HTTP.Exceptions.RequestError(request, io_error))
+        end
+        router = Router([
+            Route("PUT", "/latest/api/token", connection_timeout),
+            response_route("GET", path, HTTP.Response(instance_id)),
+        ])
+        apply(_imds_patch(router)) do
+            session = IMDS.Session()
+            msg_regex = r"IMDSv2 token request rejected due to reaching hop limit"
+            r = @test_logs (:warn, msg_regex) begin
+                IMDS.request(session, "GET", path; status_exception=false)
+            end
+            @test r isa HTTP.Response
+            @test r.status == 200
+            @test String(r.body) == instance_id
+            @test isempty(session.token)
+        end
     end
 
     @testset "get" begin