From a5debfe10e49fb3b096d6bbad77933853bb390c9 Mon Sep 17 00:00:00 2001
From: Wojtek Mach <wojtek@wojtekmach.pl>
Date: Fri, 28 Jun 2024 17:26:04 +0200
Subject: [PATCH] `put_aws_sigv4`: Add `:session_token` option

---
 lib/req/steps.ex |  5 ++++-
 lib/req/utils.ex | 26 +++++++++++++++-----------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/lib/req/steps.ex b/lib/req/steps.ex
index f9ae0b0..f9d6cea 100644
--- a/lib/req/steps.ex
+++ b/lib/req/steps.ex
@@ -1102,9 +1102,11 @@ defmodule Req.Steps do
 
         * `:secret_access_key` - the AWS secret access key.
 
+        * `:token` - if set, the AWS session token, returned from AWS STS.
+
         * `:service` - the AWS service. Defaults to `:s3`.
 
-        * `:region` - if set, AWS region. Defaults to `"us-east-1"`.
+        * `:region` - the AWS region. Defaults to `"us-east-1"`.
 
         * `:datetime` - the request datetime, defaults to `DateTime.utc_now(:second)`.
 
@@ -1161,6 +1163,7 @@ defmodule Req.Steps do
       Req.Request.validate_options(aws_options, [
         :access_key_id,
         :secret_access_key,
+        :session_token,
         :service,
         :region,
         :datetime,
diff --git a/lib/req/utils.ex b/lib/req/utils.ex
index 417c419..fcb13ab 100644
--- a/lib/req/utils.ex
+++ b/lib/req/utils.ex
@@ -23,6 +23,7 @@ defmodule Req.Utils do
   def aws_sigv4_headers(options) do
     {access_key_id, options} = Keyword.pop!(options, :access_key_id)
     {secret_access_key, options} = Keyword.pop!(options, :secret_access_key)
+    {session_token, options} = Keyword.pop(options, :session_token)
     {region, options} = Keyword.pop!(options, :region)
     {service, options} = Keyword.pop!(options, :service)
     {datetime, options} = Keyword.pop!(options, :datetime)
@@ -41,12 +42,19 @@ defmodule Req.Utils do
 
     method = method |> Atom.to_string() |> String.upcase()
 
-    canonical_headers =
-      headers ++
-        [
-          {"x-amz-content-sha256", body_digest},
-          {"x-amz-date", datetime_string}
-        ]
+    aws_headers = [
+      {"x-amz-content-sha256", body_digest},
+      {"x-amz-date", datetime_string}
+    ]
+
+    aws_headers =
+      if session_token do
+        aws_headers ++ [{"x-amz-security-token", session_token}]
+      else
+        aws_headers
+      end
+
+    canonical_headers = headers ++ aws_headers
 
     ## canonical_headers needs to be sorted for canonical_request construction
     canonical_headers = Enum.sort(canonical_headers)
@@ -95,11 +103,7 @@ defmodule Req.Utils do
     authorization =
       "AWS4-HMAC-SHA256 Credential=#{credential},SignedHeaders=#{signed_headers},Signature=#{signature}"
 
-    [
-      {"authorization", authorization},
-      {"x-amz-content-sha256", body_digest},
-      {"x-amz-date", datetime_string}
-    ] ++ headers
+    [{"authorization", authorization}] ++ aws_headers ++ headers
   end
 
   @doc """