From ceedfaf8c3943240ab4a3bc6ade99667a212c0c1 Mon Sep 17 00:00:00 2001
From: Jan Rochel <jrochel@besport.com>
Date: Thu, 4 Mar 2021 11:13:58 +0100
Subject: [PATCH] support for authentification using temporary session tokens

according to https://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
---
 async/runtime.ml                    |  2 ++
 async/runtime.mli                   |  1 +
 lib/aws.ml                          | 13 ++++++++-----
 lib/aws.mli                         |  1 +
 libraries/s3/lib_test/test_async.ml |  3 ++-
 libraries/s3/lib_test/test_lwt.ml   |  3 ++-
 lwt/runtime.ml                      |  2 ++
 lwt/runtime.mli                     |  1 +
 8 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/async/runtime.ml b/async/runtime.ml
index 6d30c6d39..2f595a62d 100644
--- a/async/runtime.ml
+++ b/async/runtime.ml
@@ -44,6 +44,7 @@ let run_request
     ~region
     ~access_key
     ~secret_key
+    ?session_token
     (module M : Aws.Call
       with type input = input
        and type output = output
@@ -53,6 +54,7 @@ let run_request
     Aws.Signing.sign_request
       ~access_key
       ~secret_key
+      ?session_token
       ~service:M.service
       ~region
       (M.to_http M.service region inp)
diff --git a/async/runtime.mli b/async/runtime.mli
index 579d4c9dc..744137945 100644
--- a/async/runtime.mli
+++ b/async/runtime.mli
@@ -35,6 +35,7 @@ val run_request :
      region:string
   -> access_key:string
   -> secret_key:string
+  -> ?session_token:string
   -> ('input, 'output, 'error) Aws.call
   -> 'input
   -> [ `Ok of 'output | `Error of 'error Aws.Error.t ] Async.Deferred.t
diff --git a/lib/aws.ml b/lib/aws.ml
index 8c9790c9f..e51bd75fd 100644
--- a/lib/aws.ml
+++ b/lib/aws.ml
@@ -501,7 +501,7 @@ module Signing = struct
   (* NOTE(dbp 2015-01-13): This is a direct translation of reference implementation at:
    * http://docs.aws.amazon.com/general/latest/gr/sigv4-signed-request-examples.html
    *)
-  let sign_request ~access_key ~secret_key ~service ~region (meth, uri, headers) =
+  let sign_request ~access_key ~secret_key ?session_token ~service ~region (meth, uri, headers) =
     let host = Util.of_option_exn (Endpoints.endpoint_of service region) in
     let params = encode_query (Uri.query uri) in
     let sign key msg = Hash.sha256 ~key msg in
@@ -519,6 +519,10 @@ module Signing = struct
       ; "x-amz-content-sha256", payload_hash
       ; "x-amz-date", amzdate
       ]
+      @
+      match session_token with
+      | None -> []
+      | Some token -> ["x-amz-security-token", token]
     in
     let signed_headers = String.concat ";" (List.map fst canonical_headers) in
     let canonical_headers_str =
@@ -571,10 +575,9 @@ module Signing = struct
         ]
     in
     let headers =
-      ("x-amz-date", amzdate)
-      :: ("x-amz-content-sha256", payload_hash)
-      :: ("Authorization", authorization_header)
-      :: headers
+      canonical_headers
+      @ ["Authorization", authorization_header]
+      @ headers
     in
     meth, uri, headers
 end
diff --git a/lib/aws.mli b/lib/aws.mli
index 22f7153e8..2957d56e9 100644
--- a/lib/aws.mli
+++ b/lib/aws.mli
@@ -281,6 +281,7 @@ module Signing : sig
   val sign_request :
        access_key:string
     -> secret_key:string
+    -> ?session_token:string
     -> service:string
     -> region:string
     -> Request.t
diff --git a/libraries/s3/lib_test/test_async.ml b/libraries/s3/lib_test/test_async.ml
index 5445b7c0e..3f6a73a63 100644
--- a/libraries/s3/lib_test/test_async.ml
+++ b/libraries/s3/lib_test/test_async.ml
@@ -7,7 +7,8 @@ module T = TestSuite (struct
 
   let secret_key = Unix.getenv "AWS_SECRET_KEY"
 
-  let run_request = Aws_async.Runtime.run_request ~access_key ~secret_key
+  let run_request ~region call input =
+    Aws_async.Runtime.run_request ~region ~access_key ~secret_key call input
 
   let un_m v = Async.Thread_safe.block_on_async_exn (fun () -> v)
 end)
diff --git a/libraries/s3/lib_test/test_lwt.ml b/libraries/s3/lib_test/test_lwt.ml
index 1ea21bc02..6b388cb59 100644
--- a/libraries/s3/lib_test/test_lwt.ml
+++ b/libraries/s3/lib_test/test_lwt.ml
@@ -7,7 +7,8 @@ module T = TestSuite (struct
 
   let secret_key = Unix.getenv "AWS_SECRET_KEY"
 
-  let run_request = Aws_lwt.Runtime.run_request ~access_key ~secret_key
+  let run_request ~region call input =
+    Aws_lwt.Runtime.run_request ~region ~access_key ~secret_key call input
 
   let un_m = Lwt_main.run
 end)
diff --git a/lwt/runtime.ml b/lwt/runtime.ml
index dc6bf4634..99ff75349 100644
--- a/lwt/runtime.ml
+++ b/lwt/runtime.ml
@@ -38,6 +38,7 @@ let run_request
     ~region
     ~access_key
     ~secret_key
+    ?session_token
     (module M : Aws.Call
       with type input = input
        and type output = output
@@ -47,6 +48,7 @@ let run_request
     Aws.Signing.sign_request
       ~access_key
       ~secret_key
+      ?session_token
       ~service:M.service
       ~region
       (M.to_http M.service region inp)
diff --git a/lwt/runtime.mli b/lwt/runtime.mli
index e72eb344f..c38d1b6ab 100644
--- a/lwt/runtime.mli
+++ b/lwt/runtime.mli
@@ -37,6 +37,7 @@ val run_request :
      region:string
   -> access_key:string
   -> secret_key:string
+  -> ?session_token:string
   -> ('input, 'output, 'error) Aws.call
   -> 'input
   -> [ `Ok of 'output | `Error of 'error Aws.Error.t ] Lwt.t