From 729320d76fe8ddb79b103c8559652cc1e391279e Mon Sep 17 00:00:00 2001
From: Neil Ramsay <2934552+neilramsay@users.noreply.github.com>
Date: Fri, 5 Jul 2024 13:41:46 +1200
Subject: [PATCH] fix: Prevent global config of API Endpoint

There is a bug where the AWS SDK session is
configured to use a specific endpoint.
This commit moves the endpoint configuration from the SDK Session to
specific SDK Clients (SSM).

The SSM endpoint is passed in to the session-manager-plugin from the
AWS CLI, and I believe the intention is to ensure any Region overrides
from the AWS CLI are respected by the session-manager-plugin.

A problem occurs when the AWS credentials are sourced from AWS SSO /
Identity Center.
When the SDK Session is configured with an endpoint, this overrides all
Service Endpoints, including SDK internal SSO credential providers.
When the current SSO credentials need to be refreshed, the refresh
request is sent to https://ssm.[aws-region].amazonaws.com/token.
The SSM API does not know about the /token endpoint, or CreateToken
operation, and returns an UnknownOperation error.
---
 src/sdkutil/awsconfig.go                           | 10 ++--------
 src/sessionmanagerplugin/session/sessionhandler.go | 14 +++++++-------
 src/ssmclicommands/startsession.go                 | 10 ++++------
 3 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/src/sdkutil/awsconfig.go b/src/sdkutil/awsconfig.go
index 039481ce..c1e6008e 100644
--- a/src/sdkutil/awsconfig.go
+++ b/src/sdkutil/awsconfig.go
@@ -26,14 +26,13 @@ import (
 var defaultRegion string
 var defaultProfile string
 
-// GetNewSessionWithEndpoint creates aws sdk session with given profile, region and endpoint
-func GetNewSessionWithEndpoint(endpoint string) (sess *session.Session, err error) {
+// GetDefaultSession creates aws sdk session with given profile and region
+func GetDefaultSession() (sess *session.Session, err error) {
 	if sess, err = session.NewSessionWithOptions(session.Options{
 		Config: aws.Config{
 			Retryer:    newRetryer(),
 			SleepDelay: sleepDelay,
 			Region:     aws.String(defaultRegion),
-			Endpoint:   aws.String(endpoint),
 		},
 		SharedConfigState: session.SharedConfigEnable,
 		Profile:           defaultProfile,
@@ -43,11 +42,6 @@ func GetNewSessionWithEndpoint(endpoint string) (sess *session.Session, err erro
 	return sess, nil
 }
 
-// GetDefaultSession creates aws sdk session with given profile and region
-func GetDefaultSession() (sess *session.Session, err error) {
-	return GetNewSessionWithEndpoint("")
-}
-
 // Sets the region and profile for default aws sessions
 func SetRegionAndProfile(region string, profile string) {
 	defaultRegion = region
diff --git a/src/sessionmanagerplugin/session/sessionhandler.go b/src/sessionmanagerplugin/session/sessionhandler.go
index cab3cfa7..ee10407f 100644
--- a/src/sessionmanagerplugin/session/sessionhandler.go
+++ b/src/sessionmanagerplugin/session/sessionhandler.go
@@ -19,6 +19,7 @@ import (
 	"math/rand"
 	"os"
 
+	"github.com/aws/aws-sdk-go/aws"
 	sdkSession "github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/session-manager-plugin/src/config"
@@ -99,10 +100,10 @@ func (s *Session) GetResumeSessionParams(log log.T) (string, error) {
 		sdkSession          *sdkSession.Session
 	)
 
-	if sdkSession, err = sdkutil.GetNewSessionWithEndpoint(s.Endpoint); err != nil {
+	if sdkSession, err = sdkutil.GetDefaultSession(); err != nil {
 		return "", err
 	}
-	s.sdk = ssm.New(sdkSession)
+	s.sdk = ssm.New(sdkSession, aws.NewConfig().WithEndpoint(s.Endpoint))
 
 	resumeSessionInput := ssm.ResumeSessionInput{
 		SessionId: &s.SessionId,
@@ -140,15 +141,14 @@ func (s *Session) ResumeSessionHandler(log log.T) (err error) {
 // TerminateSession calls TerminateSession API
 func (s *Session) TerminateSession(log log.T) error {
 	var (
-		err        error
-		newSession *sdkSession.Session
+		err error
 	)
 
-	if newSession, err = sdkutil.GetNewSessionWithEndpoint(s.Endpoint); err != nil {
-		log.Errorf("Terminate Session failed: %v", err)
+	if sdkSession, err := sdkutil.GetDefaultSession(); err != nil {
 		return err
+	} else {
+		s.sdk = ssm.New(sdkSession, aws.NewConfig().WithEndpoint(s.Endpoint))
 	}
-	s.sdk = ssm.New(newSession)
 
 	terminateSessionInput := ssm.TerminateSessionInput{
 		SessionId: &s.SessionId,
diff --git a/src/ssmclicommands/startsession.go b/src/ssmclicommands/startsession.go
index 81ef703a..bd05aa55 100644
--- a/src/ssmclicommands/startsession.go
+++ b/src/ssmclicommands/startsession.go
@@ -21,7 +21,7 @@ import (
 	"html/template"
 	"strings"
 
-	sdkSession "github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/session-manager-plugin/src/datachannel"
 	"github.com/aws/session-manager-plugin/src/jsonutil"
@@ -92,13 +92,11 @@ type StartSessionCommand struct {
 var getSSMClient = func(log log.T, region string, profile string, endpoint string) (*ssm.SSM, error) {
 	sdkutil.SetRegionAndProfile(region, profile)
 
-	var sdkSession *sdkSession.Session
-	sdkSession, err := sdkutil.GetNewSessionWithEndpoint(endpoint)
-	if err != nil {
-		log.Errorf("Get session with endpoint Failed: %v", err)
+	if sdkSession, err := sdkutil.GetDefaultSession(); err != nil {
 		return nil, err
+	} else {
+		return ssm.New(sdkSession, aws.NewConfig().WithEndpoint(endpoint)), nil
 	}
-	return ssm.New(sdkSession), nil
 }
 
 // executeSession to open datachannel