Try #383:

JuliaCloud · Jul 16, 2021 · 5e6def5 · 5e6def5
2 parents 1bd6737 + ef5067f
commit 5e6def5
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 39 deletions.
diff --git a/src/AWSCredentials.jl b/src/AWSCredentials.jl
@@ -114,7 +114,7 @@ function AWSCredentials(; profile=nothing, throw_cred_error=true)
         () -> dot_aws_config(profile),
         credentials_from_webtoken,
         ecs_instance_credentials,
-        ec2_instance_credentials
+        () -> ec2_instance_credentials(profile),
     ]
 
     # Loop through our search locations until we get credentials back
@@ -235,32 +235,71 @@ ec2_instance_region() = ec2_instance_metadata("/latest/meta-data/placement/regio
 
 
 """
-    ec2_instance_credentials() -> AWSCredentials
+    ec2_instance_credentials(profile::AbstractString) -> AWSCredentials
 
 Parse the EC2 metadata to retrieve AWSCredentials.
 """
-function ec2_instance_credentials()
-    info = ec2_instance_metadata("/latest/meta-data/iam/info")
+function ec2_instance_credentials(profile::AbstractString)
+    path = dot_aws_config_file()
+    ini = Inifile()
+    if isfile(path)
+        ini = read(ini, path)
+    end
 
-    if info === nothing
-        return nothing
+    # Any profile except default must specify the credential_source as Ec2InstanceMetadata.
+    if profile != "default"
+        source = _get_ini_value(ini, profile, "credential_source")
+        source == "Ec2InstanceMetadata" || return nothing
     end
 
+    info = ec2_instance_metadata("/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")
-    new_creds = JSON.parse(creds)
+    parsed = JSON.parse(creds)
+    instance_profile_creds = AWSCredentials(
+        parsed["AccessKeyId"],
+        parsed["SecretAccessKey"],
+        parsed["Token"],
+        info["InstanceProfileArn"];
+        expiry=DateTime(rstrip(parsed["Expiration"], 'Z')),
+        renew=() -> ec2_instance_credentials(profile),
+    )
 
-    expiry = DateTime(rstrip(new_creds["Expiration"], 'Z'))
+    # Look for a role to assume and return instance profile credentials if there is none.
+    role_arn = _get_ini_value(ini, profile, "role_arn")
+    role_arn === nothing && return instance_profile_creds
 
+    # Assume the role.
+    role_session = get(ENV, "AWS_ROLE_SESSION_NAME") do
+        _role_session_name(
+            "AWS.jl-role-",
+            basename(role_arn),
+            "-" * Dates.format(@mock(now(UTC)), dateformat"yyyymmdd\THHMMSS\Z"),
+        )
+    end
+    params = Dict{String, Any}("RoleArn" => role_arn, "RoleSessionName" => role_session)
+    duration = _get_ini_value(ini, profile, "duration_seconds")
+    if duration !== nothing
+        params["DurationSeconds"] = parse(Int, duration)
+    end
+    resp = @mock AWSServices.sts(
+        "AssumeRole",
+        params;
+        aws_config=AWSConfig(creds=instance_profile_creds),
+    )
+    role_creds = resp["AssumeRoleResult"]["Credentials"]
+    role_user = resp["AssumeRoleResult"]["AssumedRoleUser"]
     return AWSCredentials(
-        new_creds["AccessKeyId"],
-        new_creds["SecretAccessKey"],
-        new_creds["Token"],
-        info["InstanceProfileArn"];
-        expiry=expiry,
-        renew=ec2_instance_credentials
+        role_creds["AccessKeyId"],
+        role_creds["SecretAccessKey"],
+        role_creds["SessionToken"],
+        role_user["Arn"],
+        expiry=DateTime(rstrip(role_creds["Expiration"], 'Z')),
+        renew=() -> ec2_instance_credentials(profile),
     )
 end
 
@@ -480,14 +519,14 @@ function credentials_from_webtoken()
         )
     end
 
-    resp = AWSServices.sts(
+    resp = @mock AWSServices.sts(
         "AssumeRoleWithWebIdentity",
         Dict(
             "RoleArn" => role_arn,
             "RoleSessionName" => role_session,  # Required by AssumeRoleWithWebIdentity
             "WebIdentityToken" => web_identity,
         );
-        aws_config=AWSConfig(creds=nothing)
+        aws_config=AWSConfig(creds=nothing),
     )
 
     role_creds = resp["AssumeRoleWithWebIdentityResult"]["Credentials"]

diff --git a/test/AWSCredentials.jl b/test/AWSCredentials.jl
@@ -469,14 +469,50 @@ end
     end
 
     @testset "Instance - EC2" begin
-        apply(_http_request_patch) do
-            result = ec2_instance_credentials()
+        role_name = "foobar"
+        role_arn = "arn:aws:sts::1234:assumed-role/$role_name"
+        access_key = "access-key-$(randstring(6))"
+        secret_key = "secret-key-$(randstring(6))"
+        session_token = "session-token-$(randstring(6))"
+        session_name = "$role_name-session"
+        patch = Patches._assume_role_patch(
+            "AssumeRole";
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+            role_arn=role_arn,
+        )
+
+        apply([patch, _http_request_patch]) do
+            result = ec2_instance_credentials("default")
             @test result.access_key_id == test_values["AccessKeyId"]
             @test result.secret_key == test_values["SecretAccessKey"]
             @test result.token == test_values["Token"]
             @test result.user_arn == test_values["InstanceProfileArn"]
             @test result.expiry == test_values["Expiration"]
-            @test result.renew == ec2_instance_credentials
+            @test result.renew !== nothing
+
+            result = mktemp() do config_file, config_io
+                write(
+                    config_io,
+                    """
+                    [profile $role_name]
+                    credential_source = Ec2InstanceMetadata
+                    role_arn = $role_arn
+                    """
+                )
+                close(config_io)
+
+                withenv("AWS_CONFIG_FILE" => config_file, "AWS_ROLE_SESSION_NAME" => session_name) do
+                    ec2_instance_credentials(role_name)
+                end
+            end
+
+            @test result.access_key_id == access_key
+            @test result.secret_key == secret_key
+            @test result.token == session_token
+            @test result.user_arn == "$(role_arn)/$(session_name)"
+            @test result.renew !== nothing
         end
     end
 
@@ -507,7 +543,8 @@ end
             session_token = "session-token-$(randstring(6))"
             role_arn = "arn:aws:sts::1234:assumed-role/foobar"
 
-            patch = Patches._web_identity_patch(;
+            patch = Patches._assume_role_patch(
+                "AssumeRoleWithWebIdentity";
                 access_key=access_key,
                 secret_key=secret_key,
                 session_token=session_token,

diff --git a/test/patch.jl b/test/patch.jl
@@ -67,31 +67,27 @@ _config_file_patch = @patch function dot_aws_config_file()
     return ""
 end
 
-_web_identity_patch = function (;
-    access_key="web_identity_access_key",
-    secret_key="web_identity_secret_key",
-    session_token="web_session_token",
+_assume_role_patch = function (
+    op;
+    access_key="access_key",
+    secret_key="secret_key",
+    session_token="token",
     role_arn="arn:aws:sts:::assumed-role/role-name",
 )
-    @patch function AWS._http_request(request)
-        params = Dict(split.(split(request.content, '&'), '='))
-        creds = Dict(
-            "AccessKeyId" => access_key,
-            "SecretAccessKey" => secret_key,
-            "SessionToken" => session_token,
-            "Expiration" => string(now(UTC)),
-        )
-
-        result = Dict(
-            "AssumeRoleWithWebIdentityResult" => Dict(
-                "Credentials" => creds,
+    @patch function AWSServices.sts(op, params; aws_config)
+        return Dict(
+            "$(op)Result" => Dict(
+                "Credentials" => Dict(
+                    "AccessKeyId" => access_key,
+                    "SecretAccessKey" => secret_key,
+                    "SessionToken" => session_token,
+                    "Expiration" => string(now(UTC)),
+                ),
                 "AssumedRoleUser" => Dict(
-                    "Arn" => role_arn * "/" * params["RoleSessionName"],
+                    "Arn" => "$(role_arn)/$(params["RoleSessionName"])",
                 ),
             ),
         )
-
-        return HTTP.Response(200, ["Content-Type" => "text/json", "charset" => "utf-8"], body=json(result))
     end
 end