From 7bb16a107d283ef4b9f698f2b6f6819ad51a3240 Mon Sep 17 00:00:00 2001
From: Arun Annamalai <aruanna@amazon.com>
Date: Thu, 14 Dec 2023 10:36:23 -0800
Subject: [PATCH] Bug: fsx windows fileserver SSM arn parsing was incorrect

---
 .../fsxwindowsfileserver_windows.go           |  5 +-
 .../fsxwindowsfileserver_windows_test.go      | 81 +++++++++++++------
 2 files changed, 58 insertions(+), 28 deletions(-)

diff --git a/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows.go b/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows.go
index 6b9e5e87992..af5c5bb9f84 100644
--- a/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows.go
+++ b/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows.go
@@ -20,7 +20,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"os/exec"
-	"path/filepath"
 	"strings"
 	"sync"
 	"time"
@@ -479,7 +478,9 @@ func (fv *FSxWindowsFileServerResource) retrieveSSMCredentials(credentialsParame
 	}
 
 	ssmClient := fv.ssmClientCreator.NewSSMClient(fv.region, iamCredentials)
-	ssmParam := filepath.Base(parsedARN.Resource)
+	// parsedARN.Resource looks like "arn:aws:ssm:us-west-2:123456789012:parameter/sample1/sample2/parameter1"
+	// We split by parameter and get ["", "/sample1/sample2/parameter1"]
+	ssmParam := strings.Split(parsedARN.Resource, "parameter")[1]
 	ssmParams := []string{ssmParam}
 
 	ssmParamMap, err := ssm.GetParametersFromSSM(ssmParams, ssmClient)
diff --git a/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go b/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go
index 20682d960a6..57c11b62af3 100644
--- a/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go
+++ b/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go
@@ -143,7 +143,7 @@ func TestRetrieveCredentials(t *testing.T) {
 		InvalidParameters: []*string{},
 		Parameters: []*ssm.Parameter{
 			&ssm.Parameter{
-				Name:  aws.String("test"),
+				Name:  aws.String("/test"),
 				Value: aws.String(ssmTestData),
 			},
 		},
@@ -167,35 +167,64 @@ func TestRetrieveCredentials(t *testing.T) {
 }
 
 func TestRetrieveSSMCredentials(t *testing.T) {
-	fv, _, ssmClientCreator, _, _, mockSSMClient, _, _ := setup(t)
-	credentialsParameterARN := "arn:aws:ssm:us-west-2:123456789012:parameter/test"
-
-	ssmTestData := "{\n\"username\": \"user\", \n\"password\": \"pass\"\n}"
-	ssmClientOutput := &ssm.GetParametersOutput{
-		InvalidParameters: []*string{},
-		Parameters: []*ssm.Parameter{
-			&ssm.Parameter{
-				Name:  aws.String("test"),
-				Value: aws.String(ssmTestData),
-			},
+	cases := []struct {
+		Name                     string
+		CredentialsParameterARN  string
+		CredentialsParameterName string
+	}{
+		{
+			Name:                     "TestRetrieveSSMCredentialsSimple",
+			CredentialsParameterARN:  "arn:aws:ssm:us-west-2:123456789012:parameter/test",
+			CredentialsParameterName: "/test",
+		},
+		{
+			Name:                     "TestRetrieveSSMCredentialsSimple2",
+			CredentialsParameterARN:  "arn:aws:ssm:us-west-2:123456789012:parameter/hello",
+			CredentialsParameterName: "/hello",
+		},
+		{
+			Name:                     "TestRetrieveSSMCredentialsPath",
+			CredentialsParameterARN:  "arn:aws:ssm:us-west-2:123456789012:parameter/path1/path2/hello",
+			CredentialsParameterName: "/path1/path2/hello",
 		},
 	}
-
-	iamCredentials := credentials.IAMRoleCredentials{
-		CredentialsID: "test-cred-id",
+	for _, tc := range cases {
+		t.Run(tc.Name, func(t *testing.T) {
+			fv, _, ssmClientCreator, _, _, mockSSMClient, _, _ := setup(t)
+			credentialsParameterARN := tc.CredentialsParameterARN
+
+			ssmTestData := "{\n\"username\": \"user\", \n\"password\": \"pass\"\n}"
+			ssmClientOutput := &ssm.GetParametersOutput{
+				InvalidParameters: []*string{},
+				Parameters: []*ssm.Parameter{
+					&ssm.Parameter{
+						Name:  aws.String(tc.CredentialsParameterName),
+						Value: aws.String(ssmTestData),
+					},
+				},
+			}
+
+			iamCredentials := credentials.IAMRoleCredentials{
+				CredentialsID: "test-cred-id",
+			}
+
+			gomock.InOrder(
+				ssmClientCreator.EXPECT().NewSSMClient(gomock.Any(), gomock.Any()).Return(mockSSMClient),
+				mockSSMClient.EXPECT().GetParameters(&ssm.GetParametersInput{
+					Names:          []*string{&tc.CredentialsParameterName},
+					WithDecryption: aws.Bool(false),
+				}).Return(ssmClientOutput, nil).Times(1),
+			)
+
+			err := fv.retrieveSSMCredentials(credentialsParameterARN, iamCredentials)
+			assert.NoError(t, err)
+
+			credentials := fv.Credentials
+			assert.Equal(t, "user", credentials.Username)
+			assert.Equal(t, "pass", credentials.Password)
+		})
 	}
 
-	gomock.InOrder(
-		ssmClientCreator.EXPECT().NewSSMClient(gomock.Any(), gomock.Any()).Return(mockSSMClient),
-		mockSSMClient.EXPECT().GetParameters(gomock.Any()).Return(ssmClientOutput, nil).Times(1),
-	)
-
-	err := fv.retrieveSSMCredentials(credentialsParameterARN, iamCredentials)
-	assert.NoError(t, err)
-
-	credentials := fv.Credentials
-	assert.Equal(t, "user", credentials.Username)
-	assert.Equal(t, "pass", credentials.Password)
 }
 
 func TestRetrieveASMCredentials(t *testing.T) {