Skip to content

Commit

Permalink
Batch requests to SSM API
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@ags
ags committed Dec 5, 2017
1 parent 1d1f817 commit ed17ea7
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 41 deletions.
91 changes: 65 additions & 26 deletions main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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()))
}
Expand Down Expand Up @@ -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) {
Expand All @@ -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)

Expand All @@ -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 {
Expand Down
80 changes: 65 additions & 15 deletions main_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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")
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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")
Expand All @@ -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 {
Expand Down

0 comments on commit ed17ea7

Please sign in to comment.