Improve ec2_instance_metadata and create ec2_instance_region

Project.toml

@@ -23,7 +23,7 @@ XMLDict = "228000da-037f-5747-90a9-8195ccbf91a5"
 [compat]
 Compat = "3"
 GitHub = "5"
-HTTP = "0.8, 0.9"
+HTTP = "0.9.6"
 IniFile = "0.5"
 JSON = "0.18, 0.19, 0.20, 0.21"
 MbedTLS = "0.6, 0.7, 1"

src/AWS.jl

@@ -15,6 +15,7 @@ import URIs
 
 export @service
 export _merge, AbstractAWSConfig, AWSConfig, AWSExceptions, AWSServices, Request
+export ec2_instance_metadata, ec2_instance_region
 export global_aws_config, generate_service_url, set_user_agent, sign!, sign_aws2!, sign_aws4!
 export JSONService, RestJSONService, RestXMLService, QueryService
 

src/AWSCredentials.jl

@@ -291,48 +291,57 @@ end
 
 
 """
-    _ec2_metadata(metadata_endpoint::String) -> Union{String, Nothing}
+    ec2_instance_metadata(path::AbstractString) -> Union{String, Nothing}
 
-Retrieve the EC2 meta data from the local AWS endpoint. Return the EC2 metadata request
-body, or `nothing` if not running on an EC2 instance.
+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
-- `metadata_endpoint::String`: AWS internal meta data endpoint to hit
-    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#instance-metadata-ex-1
-
-# Throws
-- `StatusError`: If the response status is >= 300, except for 404
-- `ParsingError`: Invalid HTTP request target
+- `path`: The URI path to used to specify that metadata to return
 """
-function _ec2_metadata(metadata_endpoint::String)
-    try
-        request = @mock HTTP.request("GET", "http://169.254.169.254/latest/meta-data/$metadata_endpoint")
-
-        return request === nothing ? nothing : String(request.body)
+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
-        e isa HTTP.IOError || e isa HTTP.StatusError && e.status == 404 || rethrow(e)
+        if e isa HTTP.ConnectionPool.ConnectTimeout
+            nothing
+        else
+            rethrow()
+        end
     end
 
-    return nothing
+    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() = ec2_instance_metadata("/latest/meta-data/placement/region")
+
+
 """
     ec2_instance_credentials() -> AWSCredentials
 
 Parse the EC2 metadata to retrieve AWSCredentials.
 """
 function ec2_instance_credentials()
-    info = _ec2_metadata("iam/info")
+    info = ec2_instance_metadata("/latest/meta-data/iam/info")
 
     if info === nothing
         return nothing
     end
 
     info = JSON.parse(info)
 
-    name = _ec2_metadata("iam/security-credentials/")
-    creds = _ec2_metadata("iam/security-credentials/$name")
+    name = ec2_instance_metadata("/latest/meta-data/iam/security-credentials/")
+    creds = ec2_instance_metadata("/latest/meta-data/iam/security-credentials/$name")
     new_creds = JSON.parse(creds)
 
     expiry = DateTime(rstrip(new_creds["Expiration"], 'Z'))

test/AWSCredentials.jl

@@ -44,6 +44,32 @@ end
     @test AWS._role_session_name("a" ^ 22, "b" ^ 22, "c" ^ 22) == "a" ^ 22 * "b" ^ 20 * "c" ^ 22
 end
 
+@testset "ec2_instance_metadata" begin
+    @testset "connect timeout" begin
+        apply(Patches._instance_metadata_timeout_patch) do
+            @test AWS.ec2_instance_metadata("/latest/meta-data") === nothing
+        end
+    end
+end
+
+@testset "ec2_instance_region" begin
+    @testset "available" begin
+        patch = @patch function HTTP.request(method::String, url; kwargs...)
+            return HTTP.Response("ap-atlantis-1")  # Made up region
+        end
+
+        apply(patch) do
+            @test AWS.ec2_instance_region() == "ap-atlantis-1"
+        end
+    end
+
+    @testset "unavailable" begin
+        apply(Patches._instance_metadata_timeout_patch) do
+            @test AWS.ec2_instance_region() === nothing
+        end
+    end
+end
+
 @testset "AWSCredentials" begin
     @testset "Defaults" begin
         creds = AWSCredentials("access_key_id" ,"secret_key")
@@ -328,9 +354,10 @@ end
         "Security-Credentials" => "Test-Security-Credentials"
     )
 
-    _http_request_patch = @patch function HTTP.request(method::String, url::String)
+    _http_request_patch = @patch function HTTP.request(method::String, url; kwargs...)
         security_credentials = test_values["Security-Credentials"]
         uri = test_values["URI"]
+        url = string(url)
 
         if url == "http://169.254.169.254/latest/meta-data/iam/info"
             instance_profile_arn = test_values["InstanceProfileArn"]
@@ -533,31 +560,19 @@ end
     end
 
     @testset "Credentials Not Found" begin
-        _http_request_patch = @patch function HTTP.request(method::String, url::String)
-            return nothing
-        end
-
-        ACCESS_KEY = "AWS_ACCESS_KEY_ID"
-        CONTAINER_URI = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
-
-        old_aws_access_key_id = get(ENV, ACCESS_KEY, "")
-        old_container_uri = get(ENV, CONTAINER_URI, "")
-
-        try
-            delete!(ENV, "AWS_ACCESS_KEY_ID")
-            delete!(ENV, "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")
+        patches = [
+            @patch HTTP.request(method::String, url; kwargs...) = nothing
+            Patches._cred_file_patch
+            Patches._config_file_patch
+        ]
 
-            apply([_http_request_patch, Patches._cred_file_patch, Patches._config_file_patch]) do
+        withenv(
+            "AWS_ACCESS_KEY_ID" => nothing,
+            "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" => nothing,
+        ) do
+            apply(patches) do
                 @test_throws NoCredentials AWSConfig()
             end
-        finally
-            if !isempty(old_aws_access_key_id)
-                ENV[ACCESS_KEY] = old_aws_access_key_id
-            end
-
-            if !isempty(old_container_uri)
-                ENV[CONTAINER_URI] = old_container_uri
-            end
         end
     end
 

test/patch.jl

@@ -101,4 +101,8 @@ _github_tree_patch = @patch function tree(repo, tree_obj; kwargs...)
     end
 end
 
+_instance_metadata_timeout_patch = @patch function HTTP.request(method::String, url; kwargs...)
+    throw(HTTP.ConnectionPool.ConnectTimeout("169.254.169.254", "80"))
+end
+
 end