From ed17ea7d8e398cb3c05ecbd1640b06b7edbdbe05 Mon Sep 17 00:00:00 2001 From: Alex Smith Date: Fri, 3 Nov 2017 20:46:31 +1100 Subject: [PATCH] Batch requests to SSM API As reported in https://github.com/remind101/ssm-env/issues/10, the SSM GetParameters API is limited to a maximum of ten parameter names in a single request. Attempting to fetch more than that fails. This splits up the parameter names into batches of ten to avoid that limit. --- main.go | 91 +++++++++++++++++++++++++++++++++++++--------------- main_test.go | 80 ++++++++++++++++++++++++++++++++++++--------- 2 files changed, 130 insertions(+), 41 deletions(-) diff --git a/main.go b/main.go index 641eb0e..a96a712 100644 --- a/main.go +++ b/main.go @@ -15,9 +15,15 @@ import ( "github.com/aws/aws-sdk-go/service/ssm" ) -// DefaultTemplate is the default template used to determine what the SSM -// parameter name is for an environment variable. -const DefaultTemplate = `{{ if hasPrefix .Value "ssm://" }}{{ trimPrefix .Value "ssm://" }}{{ end }}` +const ( + // DefaultTemplate is the default template used to determine what the SSM + // parameter name is for an environment variable. + DefaultTemplate = `{{ if hasPrefix .Value "ssm://" }}{{ trimPrefix .Value "ssm://" }}{{ end }}` + + // defaultBatchSize is the default number of parameters to fetch at once. + // The SSM API limits this to a maximum of 10 at the time of writing. + defaultBatchSize = 10 +) // TemplateFuncs are helper functions provided to the template. var TemplateFuncs = template.FuncMap{ @@ -56,7 +62,12 @@ func main() { t, err := parseTemplate(*template) must(err) - e := &expander{t: t, ssm: ssm.New(session.New()), os: os} + e := &expander{ + batchSize: defaultBatchSize, + t: t, + ssm: ssm.New(session.New()), + os: os, + } must(e.expandEnviron(*decrypt)) must(syscall.Exec(path, args[0:], os.Environ())) } @@ -90,9 +101,10 @@ type ssmVar struct { } type expander struct { - t *template.Template - ssm ssmClient - os environ + t *template.Template + ssm ssmClient + os environ + batchSize int } func (e *expander) parameter(k, v string) (*string, error) { @@ -112,11 +124,7 @@ func (e *expander) expandEnviron(decrypt bool) error { // Environment variables that point to some SSM parameters. var ssmVars []ssmVar - input := &ssm.GetParametersInput{ - WithDecryption: aws.Bool(decrypt), - } - - names := make(map[string]bool) + uniqNames := make(map[string]bool) for _, envvar := range e.os.Environ() { k, v := splitVar(envvar) @@ -126,39 +134,70 @@ func (e *expander) expandEnviron(decrypt bool) error { } if parameter != nil { - names[*parameter] = true + uniqNames[*parameter] = true ssmVars = append(ssmVars, ssmVar{k, *parameter}) } } - for k := range names { - input.Names = append(input.Names, aws.String(k)) - } - - if len(input.Names) == 0 { + if len(uniqNames) == 0 { // Nothing to do, no SSM parameters. return nil } + names := make([]string, len(uniqNames)) + i := 0 + for k := range uniqNames { + names[i] = k + i++ + } + + for i := 0; i < len(names); i += e.batchSize { + j := i + e.batchSize + if j > len(names) { + j = len(names) + } + + values, err := e.getParameters(names[i:j], decrypt) + if err != nil { + return err + } + + for _, v := range ssmVars { + val, ok := values[v.parameter] + if ok { + e.os.Setenv(v.envvar, val) + } + } + } + + return nil +} + +func (e *expander) getParameters(names []string, decrypt bool) (map[string]string, error) { + values := make(map[string]string) + + input := &ssm.GetParametersInput{ + WithDecryption: aws.Bool(decrypt), + } + + for _, n := range names { + input.Names = append(input.Names, aws.String(n)) + } + resp, err := e.ssm.GetParameters(input) if err != nil { - return err + return values, err } if len(resp.InvalidParameters) > 0 { - return newInvalidParametersError(resp) + return values, newInvalidParametersError(resp) } - values := make(map[string]string) for _, p := range resp.Parameters { values[*p.Name] = *p.Value } - for _, v := range ssmVars { - e.os.Setenv(v.envvar, values[v.parameter]) - } - - return nil + return values, nil } type invalidParametersError struct { diff --git a/main_test.go b/main_test.go index 86b6810..e700b70 100644 --- a/main_test.go +++ b/main_test.go @@ -16,9 +16,10 @@ func TestExpandEnviron_NoSSMParameters(t *testing.T) { os := newFakeEnviron() c := new(mockSSM) e := expander{ - t: template.Must(parseTemplate(DefaultTemplate)), - os: os, - ssm: c, + t: template.Must(parseTemplate(DefaultTemplate)), + os: os, + ssm: c, + batchSize: defaultBatchSize, } decrypt := false @@ -37,9 +38,10 @@ func TestExpandEnviron_SimpleSSMParameter(t *testing.T) { os := newFakeEnviron() c := new(mockSSM) e := expander{ - t: template.Must(parseTemplate(DefaultTemplate)), - os: os, - ssm: c, + t: template.Must(parseTemplate(DefaultTemplate)), + os: os, + ssm: c, + batchSize: defaultBatchSize, } os.Setenv("SUPER_SECRET", "ssm://secret") @@ -70,9 +72,10 @@ func TestExpandEnviron_CustomTemplate(t *testing.T) { os := newFakeEnviron() c := new(mockSSM) e := expander{ - t: template.Must(parseTemplate(`{{ if eq .Name "SUPER_SECRET" }}secret{{end}}`)), - os: os, - ssm: c, + t: template.Must(parseTemplate(`{{ if eq .Name "SUPER_SECRET" }}secret{{end}}`)), + os: os, + ssm: c, + batchSize: defaultBatchSize, } os.Setenv("SUPER_SECRET", "ssm://secret") @@ -103,9 +106,10 @@ func TestExpandEnviron_DuplicateSSMParameter(t *testing.T) { os := newFakeEnviron() c := new(mockSSM) e := expander{ - t: template.Must(parseTemplate(DefaultTemplate)), - os: os, - ssm: c, + t: template.Must(parseTemplate(DefaultTemplate)), + os: os, + ssm: c, + batchSize: defaultBatchSize, } os.Setenv("SUPER_SECRET_A", "ssm://secret") @@ -138,9 +142,10 @@ func TestExpandEnviron_InvalidParameters(t *testing.T) { os := newFakeEnviron() c := new(mockSSM) e := expander{ - t: template.Must(parseTemplate(DefaultTemplate)), - os: os, - ssm: c, + t: template.Must(parseTemplate(DefaultTemplate)), + os: os, + ssm: c, + batchSize: defaultBatchSize, } os.Setenv("SUPER_SECRET", "ssm://secret") @@ -159,6 +164,51 @@ func TestExpandEnviron_InvalidParameters(t *testing.T) { c.AssertExpectations(t) } +func TestExpandEnviron_BatchParameters(t *testing.T) { + os := newFakeEnviron() + c := new(mockSSM) + e := expander{ + t: template.Must(parseTemplate(DefaultTemplate)), + os: os, + ssm: c, + batchSize: 1, + } + + os.Setenv("SUPER_SECRET_A", "ssm://secret-a") + os.Setenv("SUPER_SECRET_B", "ssm://secret-b") + + c.On("GetParameters", &ssm.GetParametersInput{ + Names: []*string{aws.String("secret-a")}, + WithDecryption: aws.Bool(false), + }).Return(&ssm.GetParametersOutput{ + Parameters: []*ssm.Parameter{ + {Name: aws.String("secret-a"), Value: aws.String("val-a")}, + }, + }, nil) + + c.On("GetParameters", &ssm.GetParametersInput{ + Names: []*string{aws.String("secret-b")}, + WithDecryption: aws.Bool(false), + }).Return(&ssm.GetParametersOutput{ + Parameters: []*ssm.Parameter{ + {Name: aws.String("secret-b"), Value: aws.String("val-b")}, + }, + }, nil) + + decrypt := false + err := e.expandEnviron(decrypt) + assert.NoError(t, err) + + assert.Equal(t, []string{ + "SHELL=/bin/bash", + "SUPER_SECRET_A=val-a", + "SUPER_SECRET_B=val-b", + "TERM=screen-256color", + }, os.Environ()) + + c.AssertExpectations(t) +} + type fakeEnviron map[string]string func newFakeEnviron() fakeEnviron {