Merge pull request #48 from mitchya1/aws-sdk-v2-migration

migrate to aws sdk go v2
chilledornaments · Jun 9, 2021 · 11a62cc · 11a62cc
2 parents 235dd13 + 6ccae17
commit 11a62cc
Show file tree

Hide file tree

Showing 7 changed files with 291 additions and 151 deletions.
diff --git a/README.md b/README.md
@@ -159,7 +159,7 @@ You must not change the `containerPath` for the `retriever` container, otherwise
                 "awslogs-stream-prefix": "init"
             }
         },
-        "image": "mitchya1/ecs-ssm-retriever:v0.2.2",
+        "image": "mitchya1/ecs-config-retriever:v0.5.0",
         "memory": 100,
         "memoryReservation": 50,
         "name": "my-container-init",

diff --git a/cmd/retriever/main.go b/cmd/retriever/main.go
@@ -2,14 +2,16 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"flag"
 	"fmt"
+	"log"
 	"os"
 	"time"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/ssm"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
 	"github.com/mitchya1/ecs-config-retriever/pkg/retriever"
 
 	vault "github.com/hashicorp/vault/api"
@@ -42,17 +44,30 @@ type ParameterSetting struct {
 	Path     string `json:"path"`
 }
 
-func ssmHandler(log *logrus.Logger) {
-
-	sess, err := session.NewSession(&aws.Config{Region: aws.String(os.Getenv("AWS_REGION"))})
+func newAWSConfig(region string) aws.Config {
+	cfg, err := config.LoadDefaultConfig(
+		context.TODO(),
+		config.WithRegion(region),
+	)
 
 	if err != nil {
-		log.Fatalf("Error creating AWS Session: %s", err.Error())
+		log.Fatalf("Failed to load SDK configuration, %v", err)
 	}
 
-	ssmClient := ssm.New(sess)
+	return cfg
+}
+
+func newSSMClient(cfg aws.Config) *ssm.Client {
+	return ssm.NewFromConfig(cfg)
+}
+
+func ssmHandler(log *logrus.Logger) {
+
+	awsConfig := newAWSConfig(os.Getenv("AWS_REGION"))
+
+	ssmClient := newSSMClient(awsConfig)
 
-	v, e := retriever.GetParameterFromSSM(ssmClient, log, parameterName, parameterIsEncrypted, parameterIsEncoded)
+	v, e := retriever.GetParameterFromSSM(context.Background(), ssmClient, log, parameterName, parameterIsEncrypted, parameterIsEncoded)
 
 	if e != nil {
 		// GetParameterFromSSM already logs the error
@@ -69,16 +84,12 @@ func ssmHandler(log *logrus.Logger) {
 }
 
 func ssmJSONHandler(log *logrus.Logger, j JSONArgument) {
-	sess, err := session.NewSession(&aws.Config{Region: aws.String(os.Getenv("AWS_REGION"))})
-
-	if err != nil {
-		log.Fatalf("Error creating AWS Session: %s", err.Error())
-	}
+	awsConfig := newAWSConfig(os.Getenv("AWS_REGION"))
 
-	ssmClient := ssm.New(sess)
+	ssmClient := newSSMClient(awsConfig)
 
 	for _, p := range j.Parameters {
-		v, e := retriever.GetParameterFromSSM(ssmClient, log, p.Name, p.Encryped, p.Encoded)
+		v, e := retriever.GetParameterFromSSM(context.Background(), ssmClient, log, p.Name, p.Encryped, p.Encoded)
 
 		if e != nil {
 			// GetParameterFromSSM already logs the error

diff --git a/go.mod b/go.mod
@@ -3,9 +3,12 @@ module github.com/mitchya1/ecs-config-retriever
 go 1.16
 
 require (
-	github.com/aws/aws-sdk-go v1.38.14
-	github.com/hashicorp/vault v1.6.2
+	github.com/aws/aws-sdk-go v1.38.58 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.6.0
+	github.com/aws/aws-sdk-go-v2/config v1.3.0
+	github.com/aws/aws-sdk-go-v2/service/ssm v1.6.2
+	github.com/hashicorp/vault v1.7.2
 	github.com/hashicorp/vault/api v1.1.0
-	github.com/sirupsen/logrus v1.7.0
+	github.com/sirupsen/logrus v1.8.1
 	gotest.tools v2.2.0+incompatible
 )
diff --git a/go.sum b/go.sum
diff --git a/pkg/retriever/retriever.go b/pkg/retriever/retriever.go
@@ -1,26 +1,27 @@
 package retriever
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/service/ssm"
-	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
 	vault "github.com/hashicorp/vault/api"
 	"github.com/sirupsen/logrus"
 )
 
 // GetParameterFromSSM retrieves the parameter from SSM
-func GetParameterFromSSM(c ssmiface.SSMAPI, log *logrus.Logger, name string, encrypted bool, encoded bool) (string, error) {
+func GetParameterFromSSM(ctx context.Context, c SSMClient, log *logrus.Logger, name string, encrypted bool, encoded bool) (string, error) {
 
 	log.Infof("Retrieving parameter '%s'", name)
 
 	input := &ssm.GetParameterInput{
 		Name:           aws.String(name),
-		WithDecryption: aws.Bool(encrypted),
+		WithDecryption: encrypted,
 	}
-	param, err := c.GetParameter(input)
+
+	param, err := c.GetParameter(ctx, input)
 
 	if err != nil {
 		log.Errorf("Error retrieving parameter: %s", err.Error())

diff --git a/pkg/retriever/retriever_test.go b/pkg/retriever/retriever_test.go
@@ -1,32 +1,28 @@
 package retriever
 
 import (
+	"context"
 	"encoding/base64"
 	"fmt"
 	"net"
 	"testing"
 	"time"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/service/ssm"
-	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
+	ssmtypes "github.com/aws/aws-sdk-go-v2/service/ssm/types"
 	"github.com/hashicorp/vault/api"
 	"github.com/hashicorp/vault/http"
 	"github.com/hashicorp/vault/vault"
 	"github.com/sirupsen/logrus"
 	"gotest.tools/assert"
 )
 
-const (
-	secretPath = "kv-v2/test/ci/secret"
-)
-
 type mockSSMClient struct {
-	ssmiface.SSMAPI
 	Encoded bool
 }
 
-func (m *mockSSMClient) GetParameter(input *ssm.GetParameterInput) (*ssm.GetParameterOutput, error) {
+func (m *mockSSMClient) GetParameter(ctx context.Context, params *ssm.GetParameterInput, optFns ...func(*ssm.Options)) (*ssm.GetParameterOutput, error) {
 
 	var v string
 
@@ -37,12 +33,12 @@ func (m *mockSSMClient) GetParameter(input *ssm.GetParameterInput) (*ssm.GetPara
 	}
 
 	r := ssm.GetParameterOutput{
-		Parameter: &ssm.Parameter{
+		Parameter: &ssmtypes.Parameter{
 			ARN:              aws.String(""),
 			DataType:         aws.String("text"),
 			Name:             aws.String("test"),
 			LastModifiedDate: aws.Time(time.Now()),
-			Type:             aws.String("String"),
+			Type:             ssmtypes.ParameterTypeString,
 			Value:            &v,
 		},
 	}
@@ -55,7 +51,7 @@ func TestRetrievePlaintextSSMParameter(t *testing.T) {
 		Encoded: false,
 	}
 
-	v, e := GetParameterFromSSM(&c, logrus.New(), "/ci/test", false, false)
+	v, e := GetParameterFromSSM(context.TODO(), &c, logrus.New(), "/ci/test", false, false)
 
 	assert.Equal(t, e, nil)
 	assert.Equal(t, v, "This is a CI test")
@@ -66,7 +62,7 @@ func TestRetrieveEncodedSSMParameter(t *testing.T) {
 		Encoded: true,
 	}
 
-	v, e := GetParameterFromSSM(&c, logrus.New(), "/ci/example", false, true)
+	v, e := GetParameterFromSSM(context.TODO(), &c, logrus.New(), "/ci/example", false, true)
 
 	assert.Equal(t, e, nil)
 	assert.Equal(t, v, "This is a base64 encoded CI test")

diff --git a/pkg/retriever/types.go b/pkg/retriever/types.go
@@ -0,0 +1,11 @@
+package retriever
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
+)
+
+type SSMClient interface {
+	GetParameter(ctx context.Context, params *ssm.GetParameterInput, optFns ...func(*ssm.Options)) (*ssm.GetParameterOutput, error)
+}