Merge pull request #40 from alvinlin123/regional-sts-by-default

Use STS regional endpoint by default

main.go

@@ -16,29 +16,32 @@
 package main
 
 import (
-    "net/http"
-    "os"
-    "strconv"
-    "time"
-    "crypto/tls"
-
-    "aws-sigv4-proxy/handler"
-    "github.com/aws/aws-sdk-go/aws/credentials"
-    "github.com/aws/aws-sdk-go/aws/credentials/stscreds"
-    "github.com/aws/aws-sdk-go/aws/session"
-    "github.com/aws/aws-sdk-go/aws/signer/v4"
-    log "github.com/sirupsen/logrus"
-    "gopkg.in/alecthomas/kingpin.v2"
+	"crypto/tls"
+	"net/http"
+	"os"
+	"strconv"
+	"time"
+
+	"aws-sigv4-proxy/handler"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
+	"github.com/aws/aws-sdk-go/aws/endpoints"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/aws/signer/v4"
+	log "github.com/sirupsen/logrus"
+	"gopkg.in/alecthomas/kingpin.v2"
 )
 
 var (
-	debug = kingpin.Flag("verbose", "enable additional logging").Short('v').Bool()
-	port  = kingpin.Flag("port", "port to serve http on").Default(":8080").String()
-	strip = kingpin.Flag("strip", "Headers to strip from incoming request").Short('s').Strings()
-	roleArn = kingpin.Flag("role-arn", "Amazon Resource Name (ARN) of the role to assume").String()
-	signingNameOverride = kingpin.Flag("name", "AWS Service to sign for").String();
-	hostOverride = kingpin.Flag("host", "Host to proxy to").String();
-	regionOverride = kingpin.Flag("region", "AWS region to sign for").String();
+	debug                  = kingpin.Flag("verbose", "enable additional logging").Short('v').Bool()
+	port                   = kingpin.Flag("port", "port to serve http on").Default(":8080").String()
+	strip                  = kingpin.Flag("strip", "Headers to strip from incoming request").Short('s').Strings()
+	roleArn                = kingpin.Flag("role-arn", "Amazon Resource Name (ARN) of the role to assume").String()
+	signingNameOverride    = kingpin.Flag("name", "AWS Service to sign for").String()
+	hostOverride           = kingpin.Flag("host", "Host to proxy to").String()
+	regionOverride         = kingpin.Flag("region", "AWS region to sign for").String()
 	disableSSLVerification = kingpin.Flag("no-verify-ssl", "Disable peer SSL certificate validation").Bool()
 )
 
@@ -50,10 +53,25 @@ func main() {
 		log.SetLevel(log.DebugLevel)
 	}
 
-    session, err := session.NewSession()
-    if err != nil {
-        log.Fatal(err)
-    }
+	sessionConfig := aws.Config{}
+	if v := os.Getenv("AWS_STS_REGIONAL_ENDPOINTS"); len(v) == 0 {
+		sessionConfig.STSRegionalEndpoint = endpoints.RegionalSTSEndpoint
+	}
+
+	session, err := session.NewSession(&sessionConfig)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if *regionOverride != "" {
+		session.Config.Region = regionOverride
+	}
+
+	// For STS regional endpoint to be effective config's region must be set.
+	if *session.Config.Region == "" {
+		defaultRegion := "us-east-1"
+		session.Config.Region = &defaultRegion
+	}
 
 	if *disableSSLVerification {
 		log.Warn("Peer SSL Certificate validation is DISABLED")
@@ -62,11 +80,11 @@ func main() {
 
 	var credentials *credentials.Credentials
 	if *roleArn != "" {
-	    credentials = stscreds.NewCredentials(session, *roleArn, func(p *stscreds.AssumeRoleProvider) {
-		    p.RoleSessionName = roleSessionName()
-	    })
+		credentials = stscreds.NewCredentials(session, *roleArn, func(p *stscreds.AssumeRoleProvider) {
+			p.RoleSessionName = roleSessionName()
+		})
 	} else {
-	    credentials = session.Config.Credentials
+		credentials = session.Config.Credentials
 	}
 
 	signer := v4.NewSigner(credentials)
@@ -77,12 +95,12 @@ func main() {
 	log.Fatal(
 		http.ListenAndServe(*port, &handler.Handler{
 			ProxyClient: &handler.ProxyClient{
-				Signer: signer,
-				Client: http.DefaultClient,
+				Signer:              signer,
+				Client:              http.DefaultClient,
 				StripRequestHeaders: *strip,
 				SigningNameOverride: *signingNameOverride,
-				HostOverride: *hostOverride,
-				RegionOverride: *regionOverride,
+				HostOverride:        *hostOverride,
+				RegionOverride:      *regionOverride,
 			},
 		}),
 	)