From 127d696d336b2b1dfdfe7177eb3d2033c1d8ba91 Mon Sep 17 00:00:00 2001
From: Matheus Pimenta <matheuscscp@gmail.com>
Date: Thu, 3 Jul 2025 14:46:02 +0100
Subject: [PATCH] Fix regression in STS endpoint for SOPS decryption with AWS
 KMS in US Gov partition

Signed-off-by: Matheus Pimenta <matheuscscp@gmail.com>
(cherry picked from commit 4623a389895f64f7e20e191a66363fcb86c73059)
---
 internal/sops/awskms/region.go      | 22 +++++++--
 internal/sops/awskms/region_test.go | 72 ++++++++++++++++++++++++++---
 2 files changed, 85 insertions(+), 9 deletions(-)

diff --git a/internal/sops/awskms/region.go b/internal/sops/awskms/region.go
index 686b39bb..8de8e4a9 100644
--- a/internal/sops/awskms/region.go
+++ b/internal/sops/awskms/region.go
@@ -17,11 +17,27 @@ limitations under the License.
 package awskms
 
 import (
-	"strings"
+	"regexp"
 )
 
+// arnRegex matches an AWS ARN, for example:
+// "arn:aws:kms:us-west-2:107501996527:key/612d5f0p-p1l3-45e6-aca6-a5b005693a48".
+// The regex matches both KMS keys and aliases, and supports different AWS partition names (aws, aws-cn, aws-us-gov).
+//
+// Copied from SOPS:
+// https://github.com/getsops/sops/blob/b2edaade23453c8774fc28ec491ddbe2b9a4c994/kms/keysource.go#L30-L32
+//
+// ref:
+// https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html
+const arnPattern = `^arn:aws[\w-]*:kms:(.+):[0-9]+:(key|alias)/.+$`
+
+var arnRegex = regexp.MustCompile(arnPattern)
+
 // GetRegionFromKMSARN extracts the region from a KMS ARN.
 func GetRegionFromKMSARN(arn string) string {
-	arn = strings.TrimPrefix(arn, "arn:aws:kms:")
-	return strings.SplitN(arn, ":", 2)[0]
+	m := arnRegex.FindStringSubmatch(arn)
+	if m == nil {
+		return ""
+	}
+	return m[1]
 }
diff --git a/internal/sops/awskms/region_test.go b/internal/sops/awskms/region_test.go
index 98a587d6..40727813 100644
--- a/internal/sops/awskms/region_test.go
+++ b/internal/sops/awskms/region_test.go
@@ -25,10 +25,70 @@ import (
 )
 
 func TestGetRegionFromKMSARN(t *testing.T) {
-	g := NewWithT(t)
-
-	arn := "arn:aws:kms:us-east-1:211125720409:key/mrk-3179bb7e88bc42ffb1a27d5038ceea25"
-
-	region := awskms.GetRegionFromKMSARN(arn)
-	g.Expect(region).To(Equal("us-east-1"))
+	for _, tt := range []struct {
+		arn      string
+		expected string
+	}{
+		{
+			arn:      "arn:aws:kms:us-west-2:107501996527:key/612d5f0p-p1l3-45e6-aca6-a5b005693a48",
+			expected: "us-west-2",
+		},
+		{
+			arn:      "arn:aws-cn:kms:cn-north-1:123456789012:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+			expected: "cn-north-1",
+		},
+		{
+			arn:      "arn:aws-us-gov:kms:us-gov-west-1:123456789012:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+			expected: "us-gov-west-1",
+		},
+		{
+			arn:      "arn:aws:kms:us-west-2:107501996527:alias/my-key-alias",
+			expected: "us-west-2",
+		},
+		{
+			arn:      "arn:aws:kms:us-west-2:107501996527:key/",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:kms:us-west-2:107501996527:alias/",
+			expected: "",
+		},
+		{
+			arn:      "not-an-arn",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:s3:::my-bucket",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:ec2:us-west-2:123456789012:instance/i-1234567890abcdef0",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:iam::123456789012:user/David",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:lambda:us-west-2:123456789012:function:my-function",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:dynamodb:us-west-2:123456789012:table/my-table",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:rds:us-west-2:123456789012:db:my-database",
+			expected: "",
+		},
+		{
+			arn:      "arn:aws:sns:us-west-2:123456789012:my-topic",
+			expected: "",
+		},
+	} {
+		t.Run(tt.arn, func(t *testing.T) {
+			g := NewWithT(t)
+			g.Expect(awskms.GetRegionFromKMSARN(tt.arn)).To(Equal(tt.expected))
+		})
+	}
 }