From 211625b7e8536b9c76d76349337bdcf72d76d019 Mon Sep 17 00:00:00 2001
From: windmgc <windmgc@gmail.com>
Date: Sun, 5 Jun 2022 01:17:19 +0800
Subject: [PATCH] add tests; fix minor bugs

---
 .../aws-lambda/iam-sts-credentials.lua        | 12 +++-
 .../07-iam-sts-credentials_spec.lua           | 72 +++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 spec/03-plugins/27-aws-lambda/07-iam-sts-credentials_spec.lua

diff --git a/kong/plugins/aws-lambda/iam-sts-credentials.lua b/kong/plugins/aws-lambda/iam-sts-credentials.lua
index 2da0c3b4ced..7f64ff9842c 100644
--- a/kong/plugins/aws-lambda/iam-sts-credentials.lua
+++ b/kong/plugins/aws-lambda/iam-sts-credentials.lua
@@ -8,6 +8,7 @@ local kong = kong
 
 local DEFAULT_SESSION_DURATION_SECONDS = 3600
 local DEFAULT_HTTP_CLINET_TIMEOUT = 60000
+local DEFAULT_ROLE_SESSION_NAME = "kong"
 
 
 local function get_regional_sts_endpoint(aws_region)
@@ -22,6 +23,12 @@ end
 local function fetch_assume_role_credentials(aws_region, assume_role_arn,
                                              role_session_name, access_key,
                                              secret_key, session_token)
+  if not assume_role_arn then
+    return nil, "Missing required parameter 'assume_role_arn' for fetching STS credentials"
+  end
+
+  role_session_name = role_session_name or DEFAULT_ROLE_SESSION_NAME
+
   kong.log.debug('Trying to assume role [', assume_role_arn, ']')
 
   local sts_host = get_regional_sts_endpoint(aws_region)
@@ -42,19 +49,20 @@ local function fetch_assume_role_credentials(aws_region, assume_role_arn,
     RoleSessionName = role_session_name,
   }
 
-  local ar_sign_params = {
+  local assume_role_sign_params = {
     region          = aws_region,
     service         = "sts",
     access_key      = access_key,
     secret_key      = secret_key,
     method          = "GET",
+    host            = sts_host,
     port            = 443,
     headers         = assume_role_request_headers,
     query           = utils.encode_args(assume_role_query_params)
   }
 
   local request, err
-  request, err = aws_v4(ar_sign_params)
+  request, err = aws_v4(assume_role_sign_params)
 
   if err then
     return nil, 'Unable to build signature to assume role ['
diff --git a/spec/03-plugins/27-aws-lambda/07-iam-sts-credentials_spec.lua b/spec/03-plugins/27-aws-lambda/07-iam-sts-credentials_spec.lua
new file mode 100644
index 00000000000..36991890ddf
--- /dev/null
+++ b/spec/03-plugins/27-aws-lambda/07-iam-sts-credentials_spec.lua
@@ -0,0 +1,72 @@
+require "spec.helpers"
+
+describe("[AWS Lambda] iam-sts", function()
+
+  local fetch_sts_assume_role, http_responses
+
+  before_each(function()
+    package.loaded["kong.plugins.aws-lambda.iam-sts-credentials"] = nil
+    package.loaded["resty.http"] = nil
+    local http = require "resty.http"
+    -- mock the http module
+    http.new = function()
+      return {
+        set_timeout = function() end,
+        request_uri = function()
+          local body = http_responses[1]
+          table.remove(http_responses, 1)
+          return {
+            status = 200,
+            body = body,
+          }
+        end,
+      }
+    end
+    fetch_sts_assume_role = require("kong.plugins.aws-lambda.iam-sts-credentials").fetch_assume_role_credentials
+  end)
+
+  after_each(function()
+  end)
+
+  it("should fetch credentials from sts service", function()
+    http_responses = {
+      [[
+{
+  "AssumeRoleResponse": {
+    "AssumeRoleResult": {
+      "SourceIdentity": "kong_session",
+      "AssumedRoleUser": {
+        "Arn": "arn:aws:iam::000000000001:role/temp-role",
+        "AssumedRoleId": "arn:aws:iam::000000000001:role/temp-role"
+      },
+      "Credentials": {
+        "AccessKeyId": "the Access Key",
+        "SecretAccessKey": "the Big Secret",
+        "SessionToken": "the Token of Appreciation",
+        "Expiration": "2019-03-12T20:56:10Z"
+      },
+      "PackedPolicySize": 1000
+    },
+    "ResponseMetadata": {
+      "RequestId": "c6104cbe-af31-11e0-8154-cbc7ccf896c7"
+    }
+  }
+}
+]]
+    }
+
+    local aws_region = "ap-east-1"
+    local assume_role_arn = "arn:aws:iam::000000000001:role/temp-role"
+    local role_session_name = "kong_session"
+    local access_key = "test_access_key"
+    local secret_key = "test_secret_key"
+    local session_token = "test_session_token"
+    local iam_role_credentials, err = fetch_sts_assume_role(aws_region, assume_role_arn, role_session_name, access_key, secret_key, session_token)
+
+    assert.is_nil(err)
+    assert.equal("the Access Key", iam_role_credentials.access_key)
+    assert.equal("the Big Secret", iam_role_credentials.secret_key)
+    assert.equal("the Token of Appreciation", iam_role_credentials.session_token)
+    assert.equal(1552424170, iam_role_credentials.expiration)
+  end)
+end)