Add support for fully qualified SSM ARNs in cross-account parameters #minor (#46)

Author: stevebarreira

.gitignore

@@ -2,3 +2,5 @@
 /coverage.*
 version_wf
 .DS_Store
+
+.idea

Makefile

@@ -1,4 +1,4 @@
-GO_VERSION ?= 1.18
+GO_VERSION ?= 1.24
 GO_CI_VERSION = v1.55.0
 BINARIES = aws-env
 WD ?= $(shell pwd)

awsenv/file_replacer.go

@@ -75,15 +75,17 @@ func (r *FileReplacer) ReplaceAll(ctx context.Context) error {
 		}
 
 		path := strings.FieldsFunc(line[idx+len(r.prefix):], splitPath)[0]
+		plainPath := stripARNPrefix(path)
 
 		// if we haven't seen the path yet, init the slice
-		if _, ok := replacementIndices[path]; !ok {
-			replacementIndices[path] = make([]replacementIndex, 0, 4)
+		if _, ok := replacementIndices[plainPath]; !ok {
+			replacementIndices[plainPath] = make([]replacementIndex, 0, 4)
 		}
 
-		replacementIndices[path] = append(replacementIndices[path], replacementIndex{
-			lineNumber: i,
-			index:      idx,
+		replacementIndices[plainPath] = append(replacementIndices[plainPath], replacementIndex{
+			lineNumber:   i,
+			index:        idx,
+			originalPath: path,
 		})
 		paths = append(paths, path)
 	}
@@ -107,7 +109,7 @@ func (r *FileReplacer) ReplaceAll(ctx context.Context) error {
 
 			ln := replacement.lineNumber
 			idx := replacement.index
-			lines[ln] = fmt.Sprintf("%s%s%s", lines[ln][:idx], value, lines[ln][idx+len(r.prefix)+len(path):])
+			lines[ln] = fmt.Sprintf("%s%s%s", lines[ln][:idx], value, lines[ln][idx+len(r.prefix)+len(replacement.originalPath):])
 		}
 	}
 
@@ -129,12 +131,13 @@ func (r *FileReplacer) MustReplaceAll(ctx context.Context) {
 }
 
 type replacementIndex struct {
-	lineNumber int
-	index      int
+	lineNumber   int
+	index        int
+	originalPath string
 }
 
 // return false if the given rune isn't an acceptable Parameter Store path
 func splitPath(r rune) bool {
 	return !unicode.IsLetter(r) && !unicode.IsDigit(r) &&
-		r != '/' && r != '_' && r != '-' && r != '.'
+		r != '/' && r != '_' && r != '-' && r != '.' && r != ':'
 }

awsenv/file_replacer_test.go

@@ -3,6 +3,7 @@ package awsenv
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"log"
 	"os"
@@ -66,6 +67,24 @@ mysql_users:
 		admin_password = "awsenv:/path/to/the/password",
 	}
  )
+`
+	sampleCnfFile5 = `
+mysql_users:
+ (
+	{
+		username = "awsenv:arn:aws:ssm:us-east-1:123456789012:parameter/remote/username",
+		password = "awsenv:arn:aws:ssm:us-east-1:123456789012:parameter/remote/password",
+	}
+ )
+`
+	sampleCnfFile6 = `
+mysql_users:
+ (
+	{
+		username = "awsenv:/path/to/the/username",
+		password = "awsenv:arn:aws:ssm:us-east-1:123456789012:parameter/remote/password",
+	}
+ )
 `
 )
 
@@ -214,6 +233,101 @@ mysql_users:
 	require.Equal(t, expectedContent, string(f))
 }
 
+func TestFileReplacer_ReplaceAll_CrossAccountARN(t *testing.T) {
+
+	fileName, cleanup := writeTempFile(sampleCnfFile5)
+	defer cleanup()
+
+	oldContent, err := ioutil.ReadFile(fileName) //nolint: gosec
+	require.NoError(t, err)
+	require.Equal(t, sampleCnfFile5, string(oldContent))
+
+	// Use mockParamsGetter to simulate SSM's behavior of stripping ARN prefixes from result keys
+	getter := mockParamsGetter(func(_ context.Context, paths []string) (map[string]string, error) {
+		store := map[string]string{
+			"/remote/username": "remote_user",
+			"/remote/password": "remote_pass",
+		}
+		result := make(map[string]string, len(paths))
+		for _, p := range paths {
+			plain := stripARNPrefix(p)
+			val, ok := store[plain]
+			if !ok {
+				return nil, fmt.Errorf("not found: %s", p)
+			}
+			result[plain] = val
+		}
+		return result, nil
+	})
+
+	r := NewFileReplacer(DefaultPrefix, fileName, getter)
+
+	ctx := context.Background()
+	err = r.ReplaceAll(ctx)
+	require.NoError(t, err, "expected no error")
+
+	expectedContent := `
+mysql_users:
+ (
+	{
+		username = "remote_user",
+		password = "remote_pass",
+	}
+ )
+`
+	f, err := ioutil.ReadFile(fileName) //nolint: gosec
+	require.NoError(t, err)
+
+	require.Equal(t, expectedContent, string(f))
+}
+
+func TestFileReplacer_ReplaceAll_MixedLocalAndCrossAccount(t *testing.T) {
+
+	fileName, cleanup := writeTempFile(sampleCnfFile6)
+	defer cleanup()
+
+	oldContent, err := ioutil.ReadFile(fileName) //nolint: gosec
+	require.NoError(t, err)
+	require.Equal(t, sampleCnfFile6, string(oldContent))
+
+	getter := mockParamsGetter(func(_ context.Context, paths []string) (map[string]string, error) {
+		store := map[string]string{
+			"/path/to/the/username": "local_user",
+			"/remote/password":      "remote_pass",
+		}
+		result := make(map[string]string, len(paths))
+		for _, p := range paths {
+			plain := stripARNPrefix(p)
+			val, ok := store[plain]
+			if !ok {
+				return nil, fmt.Errorf("not found: %s", p)
+			}
+			result[plain] = val
+		}
+		return result, nil
+	})
+
+	r := NewFileReplacer(DefaultPrefix, fileName, getter)
+
+	ctx := context.Background()
+	err = r.ReplaceAll(ctx)
+	require.NoError(t, err, "expected no error")
+
+	expectedContent := `
+mysql_users:
+ (
+	{
+		username = "local_user",
+		password = "remote_pass",
+	}
+ )
+`
+	f, err := ioutil.ReadFile(fileName) //nolint: gosec
+	require.NoError(t, err)
+
+	require.Equal(t, expectedContent, string(f))
+}
+
 func writeTempFile(contents string) (string, func()) {
 
 	uid, err := uuid.NewV4()
@@ -234,5 +348,5 @@ func writeTempFile(contents string) (string, func()) {
 		log.Fatal(err)
 	}
 
-	return tmpfile.Name(), func() { os.Remove(fName) } //nolint: errcheck,gosec
+	return tmpfile.Name(), func() { os.Remove(tmpfile.Name()) } //nolint: errcheck,gosec
 }

awsenv/helpers.go

@@ -1,5 +1,19 @@
 package awsenv
 
+import "regexp"
+
+// ssmARNPrefix matches the fully qualified SSM parameter ARN prefix used for cross-account parameters.
+// note: this will not cover AWS GovCloud ARNs
+//
+//	example: `arn:aws:ssm:<region>:<account_id>:parameter<parameter_path>`
+var ssmARNPrefix = regexp.MustCompile(`arn:aws:ssm:[^:]+:[^:]+:parameter`)
+
+// stripARNPrefix removes the SSM ARN prefix from a parameter path, returning the plain path.
+// If the path does not contain an ARN prefix, it is returned unchanged.
+func stripARNPrefix(path string) string {
+	return ssmARNPrefix.ReplaceAllString(path, "")
+}
+
 func min(x, y int) int {
 	if x < y {
 		return x

awsenv/helpers_test.go

@@ -83,6 +83,42 @@ func TestMerge(t *testing.T) {
 	}
 }
 
+func TestStripARNPrefix(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{
+			name:  "plain_path",
+			input: "/my/param/path",
+			want:  "/my/param/path",
+		},
+		{
+			name:  "full_arn",
+			input: "arn:aws:ssm:us-east-1:123456789012:parameter/my/param/path",
+			want:  "/my/param/path",
+		},
+		{
+			name:  "empty_string",
+			input: "",
+			want:  "",
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			got := stripARNPrefix(test.input)
+			if got != test.want {
+				t.Errorf("stripARNPrefix(%q) = %q, want %q", test.input, got, test.want)
+			}
+		})
+	}
+}
+
 func TestChunk(t *testing.T) {
 	tests := []struct {
 		size  int

awsenv/replacer.go

@@ -128,6 +128,8 @@ func (r *Replacer) applyParamPathValues(srcEnv map[string]string, replaceWithVal
 		}
 
 		lookupValue := strings.TrimPrefix(value, r.prefix)
+		// values from the env will still include the fully qualified prefix, but the replacement will not
+		lookupValue = stripARNPrefix(lookupValue)
 		if val, ok := replaceWithValues[lookupValue]; ok {
 			srcEnv[name] = val
 		}
@@ -171,6 +173,9 @@ func fetch(ctx context.Context, ssm ParamsGetter, paths []string) (map[string]st
 	merge(dest, results)
 
 	for _, path := range paths {
+		// results from GetParams include only the path (not fully qualified), but the original paths include the fully
+		// qualified name.
+		path = stripARNPrefix(path)
 		_, ok := dest[path]
 		if !ok {
 			return dest, errors.Errorf("awsenv: param not found: %q", path)