[PR] china support (#116)

* basic china support * add test for saml/extractArns, fix test case
allcloud-io · Jul 28, 2021 · 5056a13 · 5056a13
1 parent e4d3614
commit 5056a13
Show file tree

Hide file tree

Showing 6 changed files with 96 additions and 4 deletions.
diff --git a/Makefile b/Makefile
@@ -40,7 +40,7 @@ all: darwin-amd64 linux-386 linux-amd64 windows-386 windows-amd64
 .PHONY: sign
 sign: darwin-amd64
 	# sign
-	gon -log-level=debug -log-json ./gon.json
+	gon -log-level=info ./gon.json
 
 .PHONY: zip
 zip: all sign

diff --git a/aws/sts.go b/aws/sts.go
@@ -2,6 +2,7 @@ package aws
 
 import (
 	"errors"
+	"regexp"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
@@ -58,7 +59,14 @@ func assumeSAMLRole(PrincipalArn, RoleArn, SAMLAssertion string, duration int64)
 	}
 
 	sess := session.Must(session.NewSession())
-	svc := sts.New(sess)
+
+	config := aws.NewConfig()
+	// If we request credentials for China we need to provide a Chinese region
+	idp := regexp.MustCompile(`^arn:aws-cn:iam::\d+:saml-provider\/\S+$`)
+	if idp.MatchString(PrincipalArn) {
+		config = config.WithRegion("cn-north-1")
+	}
+	svc := sts.New(sess, config)
 
 	aResp, err := svc.AssumeRoleWithSAML(&input)
 	if err != nil {

diff --git a/go.mod b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/PuerkitoBio/goquery v1.5.0
 	github.com/aws/aws-sdk-go v1.37.8
 	github.com/briandowns/spinner v1.12.0
+	github.com/coreos/etcd v3.3.13+incompatible
 	github.com/edaniels/go-saml v0.0.0-20160724042625-8c877c3ab101
 	github.com/fatih/color v1.10.0
 	github.com/go-ini/ini v1.62.0

diff --git a/go.sum b/go.sum
@@ -35,6 +35,7 @@ github.com/briandowns/spinner v1.12.0/go.mod h1:QOuQk7x+EaDASo80FEXwlwiA+j/PPIcX
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=
+github.com/coreos/etcd v3.3.13+incompatible h1:8F3hqu9fGYLBifCmRCJsicFqDx/D68Rt3q1JMazcgBQ=
 github.com/coreos/etcd v3.3.13+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
@@ -131,6 +132,7 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxv
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pty v1.1.1 h1:VkoXIwSboBpnk99O/KFauAEILuNHv5DVFKZMBN/gUgw=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=

diff --git a/saml/saml.go b/saml/saml.go
@@ -78,6 +78,10 @@ func extractArns(attrs []saml.Attribute, pArn string) (arns []ARN) {
 					continue
 				}
 
+				// people like to put spaces in there, AWS accepts them, let's remove them on our end too.
+				components[0] = strings.TrimSpace(components[0])
+				components[1] = strings.TrimSpace(components[1])
+
 				arn := ARN{}
 
 				// Logic here for "preferred arn" for the desired account.
@@ -93,8 +97,8 @@ func extractArns(attrs []saml.Attribute, pArn string) (arns []ARN) {
 					}
 				} else {
 					// Prepare patterns
-					role := regexp.MustCompile(`^arn:aws:iam::(?P<Id>\d+):(?P<Name>role\/\S+)$`)
-					idp := regexp.MustCompile(`^arn:aws:iam::\d+:saml-provider\/\S+$`)
+					role := regexp.MustCompile(`^arn:(aws|aws-cn):iam::(?P<Id>\d+):(?P<Name>role\/\S+)$`)
+					idp := regexp.MustCompile(`^arn:(aws|aws-cn):iam::\d+:saml-provider\/\S+$`)
 
 					if role.MatchString(components[0]) && idp.MatchString(components[1]) {
 						// First component is role

diff --git a/saml/saml_test.go b/saml/saml_test.go
@@ -3,8 +3,85 @@ package saml
 import (
 	"io/ioutil"
 	"testing"
+
+	"github.com/edaniels/go-saml"
 )
 
+func TestExtractArns(t *testing.T) {
+	for _, test := range []struct {
+		name string
+		attrs []saml.Attribute
+		expectARN []ARN
+	}{
+		{"aws-normal", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: "arn:aws:iam::1234567890:role/MyRole,arn:aws:iam::1234567890:saml-provider/MyProvider"}}},
+		},
+		[]ARN{{"arn:aws:iam::1234567890:role/MyRole", "arn:aws:iam::1234567890:saml-provider/MyProvider", ""}}},
+		{"aws-reversed", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: "arn:aws:iam::1234567890:saml-provider/MyProvider,arn:aws:iam::1234567890:role/MyRole"}}},
+		},
+		[]ARN{{"arn:aws:iam::1234567890:role/MyRole", "arn:aws:iam::1234567890:saml-provider/MyProvider", ""}}},
+		{"aws-spaced-front", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: " arn:aws:iam::1234567890:saml-provider/MyProvider,arn:aws:iam::1234567890:role/MyRole"}}},
+		},
+		[]ARN{{"arn:aws:iam::1234567890:role/MyRole", "arn:aws:iam::1234567890:saml-provider/MyProvider", ""}}},
+		{"aws-spaced-end", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: "arn:aws:iam::1234567890:saml-provider/MyProvider,arn:aws:iam::1234567890:role/MyRole "}}},
+		},
+		[]ARN{{"arn:aws:iam::1234567890:role/MyRole", "arn:aws:iam::1234567890:saml-provider/MyProvider", ""}}},
+		{"aws-spaced-between", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: "arn:aws:iam::1234567890:saml-provider/MyProvider, arn:aws:iam::1234567890:role/MyRole"}}},
+		},
+		[]ARN{{"arn:aws:iam::1234567890:role/MyRole", "arn:aws:iam::1234567890:saml-provider/MyProvider", ""}}},
+		{"aws-cn", []saml.Attribute{
+			{FriendlyName: "Role",
+			Name: "https://aws.amazon.com/SAML/Attributes/Role",
+			NameFormat: "string",
+			Values: []saml.AttributeValue{{Value: "arn:aws-cn:iam::1234567890:saml-provider/MyProvider,arn:aws-cn:iam::1234567890:role/MyRole"}}},
+		},
+		[]ARN{{"arn:aws-cn:iam::1234567890:role/MyRole", "arn:aws-cn:iam::1234567890:saml-provider/MyProvider", ""}}},
+	}{
+		t.Run(test.name, func(t *testing.T){
+			arn := extractArns(test.attrs, "")
+
+			if len(test.expectARN) > 0 && len(arn) == 0 {
+				t.Fatalf("expected %d arns, received nothing", len(test.expectARN))
+			}
+			if len(test.expectARN) != len(arn) {
+				t.Errorf("expected %d arns, received %d arns", len(test.expectARN), len(arn))
+			}
+			for i := 0; i < len(test.expectARN); i++ {
+				if test.expectARN[i].Name != arn[i].Name {
+					t.Errorf("expected %q, received %q", test.expectARN[i].Name, arn[i].Name)
+				}
+
+				if test.expectARN[i].Provider != arn[i].Provider {
+					t.Errorf("expected %q, received %q", test.expectARN[i].Provider, arn[i].Provider)
+				}
+
+				if test.expectARN[i].Role != arn[i].Role {
+					t.Errorf("expected %q, received %q", test.expectARN[i].Role, arn[i].Role)
+				}
+			}
+		})
+	}
+}
+
 func TestDecode(t *testing.T) {
 	for _, test := range []struct {
 		name        string