From a36ded21cfd54cc7768fefc1c809d274b1e5653b 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      | 83 +++++++++++++------
 2 files changed, 61 insertions(+), 27 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..5f567522761 100644
--- a/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go
+++ b/agent/taskresource/fsxwindowsfileserver/fsxwindowsfileserver_windows_test.go
@@ -167,35 +167,68 @@ 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
+		CredentialsParameterArgument string
+	}{
+		{
+			Name:                         "TestRetrieveSSMCredentialsSimple",
+			CredentialsParameterARN:      "arn:aws:ssm:us-west-2:123456789012:parameter/test",
+			CredentialsParameterArgument: "/test",
+		},
+		{
+			Name:                         "TestRetrieveSSMCredentialsSimple2",
+			CredentialsParameterARN:      "arn:aws:ssm:us-west-2:123456789012:parameter/hello",
+			CredentialsParameterArgument: "/hello",
+		},
+		{
+			Name:                         "TestRetrieveSSMCredentialsPath",
+			CredentialsParameterARN:      "arn:aws:ssm:us-west-2:123456789012:parameter/path1/path2/hello",
+			CredentialsParameterArgument: "/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.CredentialsParameterArgument),
+						Value: aws.String(ssmTestData),
+					},
+				},
+			}
+
+			iamCredentials := credentials.IAMRoleCredentials{
+				CredentialsID: "test-cred-id",
+			}
+
+			//&ssm.GetParametersInput{
+			//	Names:          []*string{&tc.CredentialsParameterArgument},
+			//	WithDecryption: aws.Bool(false),
+			//}
+			gomock.InOrder(
+				ssmClientCreator.EXPECT().NewSSMClient(gomock.Any(), gomock.Any()).Return(mockSSMClient),
+				mockSSMClient.EXPECT().GetParameters(&ssm.GetParametersInput{
+					Names:          []*string{&tc.CredentialsParameterArgument},
+					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) {