From 81e18aeccdbb6ed712f567a3d24a76374a16032f Mon Sep 17 00:00:00 2001
From: Gregory Reshetniak <ror6ax@gmail.com>
Date: Mon, 6 Nov 2017 19:31:38 +0100
Subject: [PATCH] added AWS enpoint handling (#3416)

---
 builtin/logical/aws/client.go               | 17 +++++++++++---
 builtin/logical/aws/path_config_root.go     | 26 ++++++++++++++++-----
 website/source/api/secret/aws/index.html.md |  4 ++++
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/builtin/logical/aws/client.go b/builtin/logical/aws/client.go
index f6bbbe2e521a..3702f75aa6c4 100644
--- a/builtin/logical/aws/client.go
+++ b/builtin/logical/aws/client.go
@@ -13,8 +13,9 @@ import (
 	"github.com/hashicorp/vault/logical"
 )
 
-func getRootConfig(s logical.Storage) (*aws.Config, error) {
+func getRootConfig(s logical.Storage, clientType string) (*aws.Config, error) {
 	credsConfig := &awsutil.CredentialsConfig{}
+	var endpoint string
 
 	entry, err := s.Get("config/root")
 	if err != nil {
@@ -29,6 +30,12 @@ func getRootConfig(s logical.Storage) (*aws.Config, error) {
 		credsConfig.AccessKey = config.AccessKey
 		credsConfig.SecretKey = config.SecretKey
 		credsConfig.Region = config.Region
+		switch {
+		case clientType == "iam" && config.IAMEndpoint != "":
+			endpoint = *aws.String(config.IAMEndpoint)
+		case clientType == "sts" && config.STSEndpoint != "":
+			endpoint = *aws.String(config.STSEndpoint)
+		}
 	}
 
 	if credsConfig.Region == "" {
@@ -51,16 +58,19 @@ func getRootConfig(s logical.Storage) (*aws.Config, error) {
 	return &aws.Config{
 		Credentials: creds,
 		Region:      aws.String(credsConfig.Region),
+		Endpoint:    &endpoint,
 		HTTPClient:  cleanhttp.DefaultClient(),
 	}, nil
 }
 
 func clientIAM(s logical.Storage) (*iam.IAM, error) {
-	awsConfig, err := getRootConfig(s)
+	awsConfig, err := getRootConfig(s, "iam")
 	if err != nil {
 		return nil, err
 	}
+
 	client := iam.New(session.New(awsConfig))
+
 	if client == nil {
 		return nil, fmt.Errorf("could not obtain iam client")
 	}
@@ -68,11 +78,12 @@ func clientIAM(s logical.Storage) (*iam.IAM, error) {
 }
 
 func clientSTS(s logical.Storage) (*sts.STS, error) {
-	awsConfig, err := getRootConfig(s)
+	awsConfig, err := getRootConfig(s, "sts")
 	if err != nil {
 		return nil, err
 	}
 	client := sts.New(session.New(awsConfig))
+
 	if client == nil {
 		return nil, fmt.Errorf("could not obtain sts client")
 	}
diff --git a/builtin/logical/aws/path_config_root.go b/builtin/logical/aws/path_config_root.go
index 754e5b2a4318..b85015c7492e 100644
--- a/builtin/logical/aws/path_config_root.go
+++ b/builtin/logical/aws/path_config_root.go
@@ -23,6 +23,14 @@ func pathConfigRoot() *framework.Path {
 				Type:        framework.TypeString,
 				Description: "Region for API calls.",
 			},
+			"iam_endpoint": &framework.FieldSchema{
+				Type:        framework.TypeString,
+				Description: "Endpoint to custom IAM server URL",
+			},
+			"sts_endpoint": &framework.FieldSchema{
+				Type:        framework.TypeString,
+				Description: "Endpoint to custom STS server URL",
+			},
 		},
 
 		Callbacks: map[logical.Operation]framework.OperationFunc{
@@ -37,11 +45,15 @@ func pathConfigRoot() *framework.Path {
 func pathConfigRootWrite(
 	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
 	region := data.Get("region").(string)
+	iamendpoint := data.Get("iam_endpoint").(string)
+	stsendpoint := data.Get("sts_endpoint").(string)
 
 	entry, err := logical.StorageEntryJSON("config/root", rootConfig{
-		AccessKey: data.Get("access_key").(string),
-		SecretKey: data.Get("secret_key").(string),
-		Region:    region,
+		AccessKey:   data.Get("access_key").(string),
+		SecretKey:   data.Get("secret_key").(string),
+		IAMEndpoint: iamendpoint,
+		STSEndpoint: stsendpoint,
+		Region:      region,
 	})
 	if err != nil {
 		return nil, err
@@ -55,9 +67,11 @@ func pathConfigRootWrite(
 }
 
 type rootConfig struct {
-	AccessKey string `json:"access_key"`
-	SecretKey string `json:"secret_key"`
-	Region    string `json:"region"`
+	AccessKey   string `json:"access_key"`
+	SecretKey   string `json:"secret_key"`
+	IAMEndpoint string `json:"iam_endpoint"`
+	STSEndpoint string `json:"sts_endpoint"`
+	Region      string `json:"region"`
 }
 
 const pathConfigRootHelpSyn = `
diff --git a/website/source/api/secret/aws/index.html.md b/website/source/api/secret/aws/index.html.md
index 76cce5747598..65adc8be25f0 100644
--- a/website/source/api/secret/aws/index.html.md
+++ b/website/source/api/secret/aws/index.html.md
@@ -52,6 +52,10 @@ valid AWS credentials with proper permissions.
   will use the `AWS_REGION` env var, `AWS_DEFAULT_REGION` env var, or
   `us-east-1` in that order.
 
+- `iam_endpoint` `(string: <optional>)` – Specifies a custom HTTP IAM endpoint to use.
+
+- `sts_endpoint` `(string: <optional>)` – Specifies a custom HTTP STS endpoint to use.
+
 ### Sample Payload
 
 ```json