diff --git a/bastion/ami.go b/bastion/ami.go index 625505a..0943c1d 100644 --- a/bastion/ami.go +++ b/bastion/ami.go @@ -34,17 +34,32 @@ func GetAndValidateAmi(sess *session.Session, input string, instance_type string return ami, nil } - input, err := GetArchitecture(sess, instance_type) + parameter, err := amiParameterForAlias(input, func() (string, error) { + return GetArchitecture(sess, instance_type) + }) if err != nil { return "", err } - if parameter, ok := amis[input]; ok { - ami, err := GetAmiFromParameter(sess, parameter) + return GetAmiFromParameter(sess, parameter) +} + +// amiParameterForAlias returns the SSM parameter path for a named AMI alias. +// Amazon Linux aliases (and unrecognised names) are resolved to the image +// matching the instance type's architecture via detectArch. Other named +// aliases, such as "windows", map directly to their own parameter. +func amiParameterForAlias(input string, detectArch func() (string, error)) (string, error) { + alias := input + if _, known := amis[input]; !known || strings.HasPrefix(input, "amazon-linux") { + arch, err := detectArch() if err != nil { return "", err } - return ami, nil + alias = arch + } + + if parameter, ok := amis[alias]; ok { + return parameter, nil } return "", errors.New("unable to find ami") } diff --git a/bastion/ami_test.go b/bastion/ami_test.go new file mode 100644 index 0000000..91df59c --- /dev/null +++ b/bastion/ami_test.go @@ -0,0 +1,51 @@ +package bastion + +import ( + "errors" + "testing" +) + +func TestAmiParameterForAlias(t *testing.T) { + cases := []struct { + name string + input string + arch string + want string + wantDetect bool + }{ + {"windows alias is not replaced by architecture", "windows", "amazon-linux", amis["windows"], false}, + {"amazon-linux on x86_64", "amazon-linux", "amazon-linux", amis["amazon-linux"], true}, + {"amazon-linux on arm64", "amazon-linux", "amazon-linux-arm64", amis["amazon-linux-arm64"], true}, + {"amazon-linux-arm64 follows instance architecture", "amazon-linux-arm64", "amazon-linux", amis["amazon-linux"], true}, + {"unknown alias falls back to architecture", "something-else", "amazon-linux-arm64", amis["amazon-linux-arm64"], true}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + detected := false + got, err := amiParameterForAlias(tc.input, func() (string, error) { + detected = true + return tc.arch, nil + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tc.want { + t.Errorf("got %q, want %q", got, tc.want) + } + if detected != tc.wantDetect { + t.Errorf("architecture detection called = %v, want %v", detected, tc.wantDetect) + } + }) + } +} + +func TestAmiParameterForAliasDetectError(t *testing.T) { + wantErr := errors.New("describe failed") + _, err := amiParameterForAlias("amazon-linux", func() (string, error) { + return "", wantErr + }) + if !errors.Is(err, wantErr) { + t.Fatalf("got error %v, want %v", err, wantErr) + } +}