Add sts_endpoint param to kinesis plugin (#39)

aws · Jul 17, 2020 · ae47001 · ae47001
1 parent b6f852c
commit ae47001
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 20 deletions.
diff --git a/fluent-bit-kinesis.go b/fluent-bit-kinesis.go
@@ -67,8 +67,10 @@ func newKinesisOutput(ctx unsafe.Pointer, pluginID int) (*kinesis.OutputPlugin,
 	logrus.Infof("[kinesis %d] plugin parameter partition_key = '%s'", pluginID, partitionKey)
 	roleARN := output.FLBPluginConfigKey(ctx, "role_arn")
 	logrus.Infof("[kinesis %d] plugin parameter role_arn = '%s'", pluginID, roleARN)
-	endpoint := output.FLBPluginConfigKey(ctx, "endpoint")
-	logrus.Infof("[kinesis %d] plugin parameter endpoint = '%s'", pluginID, endpoint)
+	kinesisEndpoint := output.FLBPluginConfigKey(ctx, "endpoint")
+	logrus.Infof("[kinesis %d] plugin parameter endpoint = '%s'", pluginID, kinesisEndpoint)
+	stsEndpoint := output.FLBPluginConfigKey(ctx, "sts_endpoint")
+	logrus.Infof("[kinesis %d] plugin parameter sts_endpoint = '%s'", pluginID, stsEndpoint)
 	appendNewline := output.FLBPluginConfigKey(ctx, "append_newline")
 	logrus.Infof("[kinesis %d] plugin parameter append_newline = %s", pluginID, appendNewline)
 	timeKey := output.FLBPluginConfigKey(ctx, "time_key")
@@ -130,7 +132,7 @@ func newKinesisOutput(ctx unsafe.Pointer, pluginID int) (*kinesis.OutputPlugin,
 		concurrencyRetriesInt = defaultConcurrentRetries
 	}
 
-	return kinesis.NewOutputPlugin(region, stream, dataKeys, partitionKey, roleARN, endpoint, timeKey, timeKeyFmt, concurrencyInt, concurrencyRetriesInt, appendNL, pluginID)
+	return kinesis.NewOutputPlugin(region, stream, dataKeys, partitionKey, roleARN, kinesisEndpoint, stsEndpoint, timeKey, timeKeyFmt, concurrencyInt, concurrencyRetriesInt, appendNL, pluginID)
 }
 
 // The "export" comments have syntactic meaning

diff --git a/kinesis/kinesis.go b/kinesis/kinesis.go
@@ -93,8 +93,8 @@ type OutputPlugin struct {
 }
 
 // NewOutputPlugin creates an OutputPlugin object
-func NewOutputPlugin(region, stream, dataKeys, partitionKey, roleARN, endpoint, timeKey, timeFmt string, concurrency, retryLimit int, appendNewline bool, pluginID int) (*OutputPlugin, error) {
-	client, err := newPutRecordsClient(roleARN, region, endpoint)
+func NewOutputPlugin(region, stream, dataKeys, partitionKey, roleARN, kinesisEndpoint, stsEndpoint, timeKey, timeFmt string, concurrency, retryLimit int, appendNewline bool, pluginID int) (*OutputPlugin, error) {
+	client, err := newPutRecordsClient(roleARN, region, kinesisEndpoint, stsEndpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -144,28 +144,30 @@ func NewOutputPlugin(region, stream, dataKeys, partitionKey, roleARN, endpoint,
 }
 
 // newPutRecordsClient creates the Kinesis client for calling the PutRecords method
-func newPutRecordsClient(roleARN string, awsRegion string, endpoint string) (*kinesis.Kinesis, error) {
+func newPutRecordsClient(roleARN string, awsRegion string, kinesisEndpoint string, stsEndpoint string) (*kinesis.Kinesis, error) {
+	customResolverFn := func(service, region string, optFns ...func(*endpoints.Options)) (endpoints.ResolvedEndpoint, error) {
+		if service == endpoints.KinesisServiceID && kinesisEndpoint != "" {
+			return endpoints.ResolvedEndpoint{
+				URL: kinesisEndpoint,
+			}, nil
+		} else if service == endpoints.StsServiceID && stsEndpoint != "" {
+			return endpoints.ResolvedEndpoint{
+				URL: stsEndpoint,
+			}, nil
+		}
+		return endpoints.DefaultResolver().EndpointFor(service, region, optFns...)
+	}
+
 	sess, err := session.NewSession(&aws.Config{
-		Region: aws.String(awsRegion),
+		Region:                        aws.String(awsRegion),
+		EndpointResolver:              endpoints.ResolverFunc(customResolverFn),
+		CredentialsChainVerboseErrors: aws.Bool(true),
 	})
 	if err != nil {
 		return nil, err
 	}
 
 	svcConfig := &aws.Config{}
-	if endpoint != "" {
-		defaultResolver := endpoints.DefaultResolver()
-		cwCustomResolverFn := func(service, region string, optFns ...func(*endpoints.Options)) (endpoints.ResolvedEndpoint, error) {
-			if service == "kinesis" {
-				return endpoints.ResolvedEndpoint{
-					URL: endpoint,
-				}, nil
-			}
-			return defaultResolver.EndpointFor(service, region, optFns...)
-		}
-		svcConfig.EndpointResolver = endpoints.ResolverFunc(cwCustomResolverFn)
-	}
-
 	if roleARN != "" {
 		creds := stscreds.NewCredentials(sess, roleARN)
 		svcConfig.Credentials = creds