Fix regression in STS endpoint for SOPS decryption with AWS KMS

Signed-off-by: Matheus Pimenta <matheuscscp@gmail.com>

Author: matheuscscp

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]
 }

internal/sops/awskms/region_test.go

@@ -25,10 +25,30 @@ 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 _, tc := 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",
+		},
+	} {
+		t.Run(tc.arn, func(t *testing.T) {
+			g := NewWithT(t)
+			g.Expect(awskms.GetRegionFromKMSARN(tc.arn)).To(Equal(tc.expected))
+		})
+	}
 }