From 7d7c41b3e3df6902ae68be52b7bfd67123e7512e Mon Sep 17 00:00:00 2001
From: windmgc <windmgc@gmail.com>
Date: Mon, 6 May 2024 14:32:53 +0800
Subject: [PATCH] fix(sts): do not inject region info for sts service with VPC
 endpoint hostname

---
 README.md                       |  2 ++
 spec/01-generic/02-aws_spec.lua | 33 +++++++++++++++++++++++++++++++++
 src/resty/aws/init.lua          |  5 ++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 142fa59..b10819d 100644
--- a/README.md
+++ b/README.md
@@ -182,6 +182,8 @@ Release process:
 - feat: decode AWS api response json body with array metatable
   [114](https://github.com/Kong/lua-resty-aws/pull/114)
 
+- fix: do not inject region info for sts service with VPC endpoint hostname
+  [113](https://github.com/Kong/lua-resty-aws/pull/113)
 
 ### 1.4.1 (19-Apr-2024)
 
diff --git a/spec/01-generic/02-aws_spec.lua b/spec/01-generic/02-aws_spec.lua
index 25cf9a9..33e0986 100644
--- a/spec/01-generic/02-aws_spec.lua
+++ b/spec/01-generic/02-aws_spec.lua
@@ -92,4 +92,37 @@ describe("AWS main instance", function()
     assert.same("https://sts.eu-central-1.amazonaws.com", sts.config.endpoint)
   end)
 
+  it("do not inject sts region info for sts vpc endpoint url", function()
+    local aws = AWS({
+      region = "eu-central-1",
+      stsRegionalEndpoints = "regional",
+    })
+
+    aws.config.credentials = aws:Credentials {
+      accessKeyId = "test_id",
+      secretAccessKey = "test_key",
+    }
+
+    assert.is.table(aws.config)
+
+    local regional_vpc_endpoint_url = "https://vpce-abcdefg-hijklmn-eu-central-1a.sts.eu-central-1.vpce.amazonaws.com"
+
+    local sts, _ = aws:STS({
+      endpoint = regional_vpc_endpoint_url,
+    })
+    local _, _ = sts:assumeRole {
+      RoleArn = "aws:arn::XXXXXXXXXXXXXXXXX:test123",
+      RoleSessionName = "aws-test",
+    }
+
+    assert.same(regional_vpc_endpoint_url, sts.config.endpoint)
+
+    local _, _ = sts:assumeRole {
+      RoleArn = "aws:arn::XXXXXXXXXXXXXXXXX:test123",
+      RoleSessionName = "aws-test",
+    }
+    assert.same(regional_vpc_endpoint_url, sts.config.endpoint)
+  end)
+
+
 end)
diff --git a/src/resty/aws/init.lua b/src/resty/aws/init.lua
index 92e79fb..9bfa816 100644
--- a/src/resty/aws/init.lua
+++ b/src/resty/aws/init.lua
@@ -323,7 +323,10 @@ local function generate_service_methods(service)
         -- https://github.com/aws/aws-sdk-js/blob/307e82673b48577fce4389e4ce03f95064e8fe0d/lib/services/sts.js#L78-L82
         assert(service.config.region, "region is required when using STS regional endpoints")
 
-        if not service.config._regionalEndpointInjected then
+        -- If the endpoint is a VPC endpoint DNS hostname then we don't need to inject the region
+        -- VPC endpoint DNS hostnames always contain region, see
+        -- https://docs.aws.amazon.com/vpc/latest/privatelink/privatelink-access-aws-services.html#interface-endpoint-dns-hostnames
+        if not service.config._regionalEndpointInjected and not service.config.endpoint:match("^(.+)(%.vpce%.amazonaws%.com)$") then
           local pre, post = service.config.endpoint:match("^(.+)(%.amazonaws%.com)$")
           service.config.endpoint = pre .. "." .. service.config.region .. post
           service.config.signingRegion = service.config.region