Support data retrevial from Instance Metadata Service Version 2 (IMDS…

…v2) (#647)

Project.toml

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

docs/make.jl

@@ -71,6 +71,7 @@ makedocs(;
         "Home" => "index.md",
         "Backends" => "backends.md",
         "AWS" => "aws.md",
+        "IMDS" => "imds.md",
         "Services" => _generate_high_level_services_docs(),
     ],
     strict=true,

docs/src/imds.md

@@ -0,0 +1,16 @@
+# IMDS
+
+```@meta
+CurrentModule = AWS
+```
+
+Provides a Julia interface for accessing AWS instance metadata.
+
+### Documentation
+
+```@docs
+AWS.IMDS
+AWS.IMDS.Session
+AWS.IMDS.get
+AWS.IMDS.region
+```

src/AWS.jl

@@ -18,7 +18,7 @@ using XMLDict
 export @service
 export _merge
 export AbstractAWSConfig, AWSConfig, AWSExceptions, AWSServices, Request
-export ec2_instance_metadata, ec2_instance_region
+export IMDS
 export assume_role, generate_service_url, global_aws_config, set_user_agent
 export sign!, sign_aws2!, sign_aws4!
 export JSONService, RestJSONService, RestXMLService, QueryService, set_features
@@ -31,6 +31,7 @@ include("AWSExceptions.jl")
 include("AWSCredentials.jl")
 include("AWSConfig.jl")
 include("AWSMetadata.jl")
+include("IMDS.jl")
 
 include(joinpath("utilities", "request.jl"))
 include(joinpath("utilities", "response.jl"))

src/AWSCredentials.jl

@@ -212,45 +212,6 @@ function check_credentials(aws_creds::AWSCredentials; force_refresh::Bool=false)
 end
 check_credentials(aws_creds::Nothing) = aws_creds
 
-"""
-    ec2_instance_metadata(path::AbstractString) -> Union{String, Nothing}
-
-Retrieve the AWS EC2 instance metadata as a string using the provided `path`. If no instance
-metadata is available (typically due to not running within an EC2 instance) then `nothing`
-will be returned. See the AWS documentation for details on what metadata is available:
-https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
-
-# Arguments
-- `path`: The URI path to used to specify that metadata to return
-"""
-function ec2_instance_metadata(path::AbstractString)
-    uri = HTTP.URI(; scheme="http", host="169.254.169.254", path=path)
-    request = try
-        @mock HTTP.request("GET", uri; connect_timeout=1)
-    catch e
-        if e isa HTTP.ConnectError
-            nothing
-        else
-            rethrow()
-        end
-    end
-
-    return request !== nothing ? String(request.body) : nothing
-end
-
-"""
-    ec2_instance_region() -> Union{String, Nothing}
-
-Determine the AWS region of the machine executing this code if running inside of an EC2
-instance, otherwise `nothing` is returned.
-"""
-ec2_instance_region() =
-    try
-        ec2_instance_metadata("/latest/meta-data/placement/region")
-    catch
-        nothing
-    end
-
 """
     ec2_instance_credentials(profile::AbstractString) -> AWSCredentials
 
@@ -269,13 +230,13 @@ function ec2_instance_credentials(profile::AbstractString)
         source == "Ec2InstanceMetadata" || return nothing
     end
 
-    info = ec2_instance_metadata("/latest/meta-data/iam/info")
+    info = IMDS.get("/latest/meta-data/iam/info")
     info === nothing && return nothing
     info = JSON.parse(info)
 
     # Get credentials for the role associated to the instance via instance profile.
-    name = ec2_instance_metadata("/latest/meta-data/iam/security-credentials/")
-    creds = ec2_instance_metadata("/latest/meta-data/iam/security-credentials/$name")
+    name = IMDS.get("/latest/meta-data/iam/security-credentials/")
+    creds = IMDS.get("/latest/meta-data/iam/security-credentials/$name")
     parsed = JSON.parse(creds)
     instance_profile_creds = AWSCredentials(
         parsed["AccessKeyId"],
@@ -695,7 +656,7 @@ function aws_get_region(; profile=nothing, config=nothing, default=DEFAULT_REGIO
     @something(
         get(ENV, "AWS_DEFAULT_REGION", nothing),
         get(_aws_profile_config(config, profile), "region", nothing),
-        ec2_instance_region(),
+        @mock(IMDS.region()),
         Some(default),
     )
 end

src/AWSExceptions.jl

@@ -5,7 +5,15 @@ using JSON
 using XMLDict
 using XMLDict: XMLDictElement
 
-export AWSException, ProtocolNotDefined, InvalidFileName, NoCredentials
+export AWSException, IMDSUnavailable, ProtocolNotDefined, InvalidFileName, NoCredentials
+
+struct IMDSUnavailable <: Exception end
+
+function Base.show(io::IO, e::IMDSUnavailable)
+    msg = "$IMDSUnavailable: The Instance Metadata Service is unavailable on the host"
+    println(io, msg)
+    return nothing
+end
 
 struct ProtocolNotDefined <: Exception
     message::String

src/IMDS.jl

@@ -0,0 +1,162 @@
+"""
+    IMDS
+
+Front-end for retrieving AWS instance metadata via the Instance Metadata Service (IMDS). For
+details on available metadata see the official AWS documentation on:
+["Instance metadata and user data"](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html).
+
+The IMDS module supports instances using either IMDSv1 or IMDSv2 (preferring IMDSv2 for
+security reasons).
+"""
+module IMDS
+
+using ..AWSExceptions: IMDSUnavailable
+
+using HTTP: HTTP
+using HTTP.Exceptions: ConnectError, StatusError
+using Mocking
+
+const IPv4_ADDRESS = "169.254.169.254"
+const DEFAULT_DURATION = 600  # 5 minutes, in seconds
+
+mutable struct Session
+    token::String
+    duration::Int16
+    expiration::Int64
+end
+
+const _SESSION = Ref{Session}()
+
+function __init__()
+    _SESSION[] = Session()
+    return nothing
+end
+
+"""
+    Session(; duration=$DEFAULT_DURATION)
+
+An IMDS `Session` which retains the IMDSv2 token over multiple requests. When IMDSv2 is
+unavailable the session switches to IMDSv1 mode and avoids future requests for IMDSv2
+tokens.
+
+# Keywords
+- `duration` (optional): Requested session duration, in seconds, for the IMDSv2 token. Can
+  be a minimum of one second and a maximum of six hours (21600).
+"""
+Session(; duration=DEFAULT_DURATION) = Session("", duration, 0)
+
+token_expired(session::Session; drift=10) = time() - session.expiration - drift > 0
+
+function refresh_token!(session::Session, duration::Integer=session.duration)
+    t = floor(Int64, time())
+    headers = ["X-aws-ec2-metadata-token-ttl-seconds" => string(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)
+
+    # 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`.
+    if r.status == 200
+        session.token = String(r.body)
+        session.duration = duration
+        session.expiration = t + duration
+    elseif r.status == 404
+        session.duration = 0
+        session.expiration = typemax(Int64)  # Use IMDSv1 indefinitely
+    else
+        # Could also populate the `StatusError` via `r.request.method` and
+        # `r.request.target` however `r.request` may not be populated under test scenarios.
+        throw(StatusError(r.status, "PUT", uri.path, r))
+    end
+
+    return session
+end
+
+function request(session::Session, method::AbstractString, path::AbstractString; kwargs...)
+    # Attempt to generate token for use with IMDSv2. If we're unable to generate a token
+    # we'll fall back on using IMDSv1. We prefer using IMDSv2 as instances can be configured
+    # to disable IMDSv1 access: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances
+    token_expired(session) && refresh_token!(session)
+    headers = Pair{String,String}[]
+    !isempty(session.token) && push!(headers, "X-aws-ec2-metadata-token" => session.token)
+
+    # 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)
+    return _http_request(method, uri, headers; kwargs...)
+end
+
+function _http_request(args...; status_exception=true, kwargs...)
+    response = try
+        # Always throw status exceptions so we can determine if the IMDS service is available
+        @mock HTTP.request(
+            args...; connect_timeout=1, retry=false, kwargs..., status_exception=true
+        )
+    catch e
+        # When running outside of an EC2 instance the link-local address will be unavailable
+        # and connections will fail. On EC2 instances where IMDS is disabled a HTTP 403 is
+        # returned.
+        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#instance-metadata-returns
+        if e isa ConnectError || e isa StatusError && e.status == 403
+            throw(IMDSUnavailable())
+
+        #! format: off
+        # Return the status exception when `status_exception=false`. We must always cause
+        # `HTTP.request` to throw status errors for our `IMDSUnavailable` check.
+        #! format: on
+        elseif !status_exception && e isa StatusError
+            e.response
+        else
+            rethrow()
+        end
+    end
+
+    return response
+end
+
+"""
+    get([session::Session], path::AbstractString) -> Union{String, Nothing}
+
+Retrieve the AWS instance metadata from the provided HTTP `path`. If no instance metadata is
+available (due to the instance metadata service being disabled or not being run from within
+an EC2 instance) then `nothing` will be returned. For details on available metadata see the
+official AWS documentation on:
+["Instance metadata and user data"](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html).
+
+# Arguments
+- `session` (optional): The IMDS `Session` used to store the IMDSv2 token.
+- `path`: The HTTP path to used to specify the metadata to return.
+"""
+function get(session::Session, path::AbstractString)
+    response = try
+        request(session, "GET", path)
+    catch e
+        if e isa IMDSUnavailable
+            nothing
+        else
+            rethrow()
+        end
+    end
+
+    return !isnothing(response) ? String(response.body) : nothing
+end
+
+get(path::AbstractString) = get(_SESSION[], path)
+
+"""
+    region([session::Session]) -> Union{String, Nothing}
+
+Determine the AWS region of the machine executing this code if running inside of an EC2
+instance, otherwise `nothing` is returned.
+
+# Arguments
+- `session` (optional): The IMDS `Session` used to store the IMDSv2 token.
+"""
+region(session::Session) = get(session, "/latest/meta-data/placement/region")
+region() = region(_SESSION[])
+
+end

src/deprecated.jl

@@ -207,3 +207,6 @@ function legacy_response(
         return (return_headers ? (body, response.headers) : body)
     end
 end
+
+@deprecate ec2_instance_metadata(path::AbstractString) IMDS.get(path)
+@deprecate ec2_instance_region() IMDS.region()