Fixed an issue with tsh aws ssm start-session

tool/tsh/common/app_aws.go

@@ -21,6 +21,7 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"strings"
 	"sync"
 
 	awsarn "github.com/aws/aws-sdk-go/aws/arn"
@@ -46,6 +47,11 @@ func onAWS(cf *CLIConf) error {
 		return trace.Wrap(err)
 	}
 
+	if shouldUseAWSEndpointURLMode(cf) {
+		log.Debugf("Forcing endpoint URL mode for AWS command %q.", cf.AWSCommandArgs)
+		cf.AWSEndpointURLMode = true
+	}
+
 	err = awsApp.StartLocalProxies()
 	if err != nil {
 		return trace.Wrap(err)
@@ -71,6 +77,42 @@ func onAWS(cf *CLIConf) error {
 	return awsApp.RunCommand(cmd)
 }
 
+func shouldUseAWSEndpointURLMode(cf *CLIConf) bool {
+	// `aws ssm start-session` first calls ssm.<region>.amazonaws.com to get an
+	// stream URL and an token. Then it makes a wss connection with the
+	// provided token to the provided stream URL. The wss request currently
+	// respects HTTPS_PROXY but does not respect local CA bundle we provided
+	// thus causing a failure. Even if this is resolved one day, the wss send
+	// the token through websocket data channel for authentication, instead of
+	// sigv4, which likely we won't support.
+	//
+	// When using the endpoint URL mode, only the first request goes through
+	// Teleport Proxy. The wss connection does not respect the endpoint URL and
+	// goes to AWS directly (thus working fine).
+	//
+	// Reference:
+	// https://github.com/aws/session-manager-plugin/
+	return isAWSCommand(cf, "ssm start-session")
+}
+
+func isAWSCommand(cf *CLIConf, wantCommand string) bool {
+	return strings.Join(removeAWSCommandFlags(cf.AWSCommandArgs), " ") == wantCommand
+}
+
+func removeAWSCommandFlags(args []string) (ret []string) {
+	for i := 0; i < len(args); i++ {
+		arg := args[i]
+		switch {
+		case strings.HasPrefix(arg, "--"):
+			i++
+			continue
+		default:
+			ret = append(ret, arg)
+		}
+	}
+	return
+}
+
 // awsApp is an AWS app that can start local proxies to serve AWS APIs.
 type awsApp struct {
 	cf      *CLIConf

tool/tsh/common/app_aws_test.go

@@ -142,6 +142,23 @@ func TestAWS(t *testing.T) {
 		setCmdRunner(validateCmd),
 	)
 	require.NoError(t, err)
+
+	t.Run("aws ssm start-session", func(t *testing.T) {
+		// Validate --endpoint-url 127.0.0.1:<port> is added to the command.
+		validateCmd := func(cmd *exec.Cmd) error {
+			require.Len(t, cmd.Args, 9)
+			require.Equal(t, []string{"aws", "ssm", "--region", "us-west-1", "start-session", "--target", "target-id", "--endpoint-url"}, cmd.Args[:8])
+			require.Contains(t, cmd.Args[8], "127.0.0.1:")
+			return nil
+		}
+		err = Run(
+			context.Background(),
+			[]string{"aws", "ssm", "--region", "us-west-1", "start-session", "--target", "target-id"},
+			setHomePath(tmpHomePath),
+			setCmdRunner(validateCmd),
+		)
+		require.NoError(t, err)
+	})
 }
 
 func makeUserWithAWSRole(t *testing.T) (types.User, types.Role) {