Fix an issue "tsh aws ssm start-session" fails when KMS encryption is…

… enabled (#50402) * Fix an issue "tsh aws ssm start-session" fails when KMS encryption is enabled * remove httputil dump * fix ut * remove unused funcs
gravitational · Jan 6, 2025 · 9abeb8c · 9abeb8c
1 parent 7019692
commit 9abeb8c
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 89 deletions.
diff --git a/lib/srv/alpnproxy/forward_proxy.go b/lib/srv/alpnproxy/forward_proxy.go
@@ -174,7 +174,27 @@ func MatchAllRequests(req *http.Request) bool {
 // MatchAWSRequests is a MatchFunc that returns true if request is an AWS API
 // request.
 func MatchAWSRequests(req *http.Request) bool {
-	return awsapiutils.IsAWSEndpoint(req.Host)
+	return awsapiutils.IsAWSEndpoint(req.Host) &&
+		// Avoid proxying SSM session WebSocket requests and let the forward proxy
+		// send it directly to AWS.
+		//
+		// `aws ssm start-session` first calls ssm.<region>.amazonaws.com to get
+		// a stream URL and a token. Then it makes a wss connection with the
+		// provided token to the provided stream URL. The stream URL looks like:
+		// wss://ssmmessages.region.amazonaws.com/v1/data-channel/session-id?stream=(input|output)
+		//
+		// The wss request currently respects HTTPS_PROXY but does not
+		// respect local CA bundle we provided thus causing a failure. The
+		// request is not signed with SigV4 either.
+		//
+		// Reference:
+		// https://github.com/aws/session-manager-plugin/
+		!isAWSSSMWebsocketRequest(req)
+}
+
+func isAWSSSMWebsocketRequest(req *http.Request) bool {
+	return awsapiutils.IsAWSEndpoint(req.Host) &&
+		strings.HasPrefix(req.Host, "ssmmessages.")
 }
 
 // MatchAzureRequests is a MatchFunc that returns true if request is an Azure API

diff --git a/lib/srv/alpnproxy/forward_proxy_test.go b/lib/srv/alpnproxy/forward_proxy_test.go
@@ -247,3 +247,43 @@ func TestMatchGCPRequests(t *testing.T) {
 		})
 	}
 }
+
+func TestMatchAWSRequests(t *testing.T) {
+	makeRequest := func(url string) *http.Request {
+		// Forward proxy always receives CONNECT requests.
+		request, err := http.NewRequest("CONNECT", url, nil)
+		require.NoError(t, err)
+		return request
+	}
+	tests := []struct {
+		name  string
+		req   *http.Request
+		check require.BoolAssertionFunc
+	}{
+		{
+			name:  "AWS request",
+			req:   makeRequest("http://s3.ca-central-1.amazonaws.com"),
+			check: require.True,
+		},
+		{
+			name:  "non-AWS request",
+			req:   makeRequest("https://registry.terraform.io"),
+			check: require.False,
+		},
+		{
+			name:  "SSM API",
+			req:   makeRequest("https://ssm.ca-central-1.amazonaws.com"),
+			check: require.True,
+		},
+		{
+			name:  "SSM session WebSocket",
+			req:   makeRequest("wss://ssmmessages.ca-central-1.amazonaws.com"),
+			check: require.False,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.check(t, MatchAWSRequests(tt.req))
+		})
+	}
+}
diff --git a/tool/tsh/common/app_aws.go b/tool/tsh/common/app_aws.go
@@ -24,7 +24,6 @@ import (
 	"io"
 	"os"
 	"os/exec"
-	"strings"
 	"sync"
 
 	awsarn "github.com/aws/aws-sdk-go/aws/arn"
@@ -51,11 +50,6 @@ 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(cf.Context)
 	if err != nil {
 		return trace.Wrap(err)
@@ -81,55 +75,6 @@ func onAWS(cf *CLIConf) error {
 	return awsApp.RunCommand(cmd)
 }
 
-func shouldUseAWSEndpointURLMode(cf *CLIConf) bool {
-	inputAWSCommand := strings.Join(removeAWSCommandFlags(cf.AWSCommandArgs), " ")
-	switch inputAWSCommand {
-	// `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/
-	//
-	// "aws ecs execute-command" also start SSM sessions.
-	case "ssm start-session", "ecs execute-command":
-		return true
-	default:
-		return false
-	}
-}
-
-func removeAWSCommandFlags(args []string) (ret []string) {
-	for i := 0; i < len(args); i++ {
-		switch {
-		case isAWSFlag(args, i):
-			// Skip next arg, if next arg is not a flag but a flag value.
-			if !isAWSFlag(args, i+1) {
-				i++
-			}
-			continue
-		default:
-			ret = append(ret, args[i])
-		}
-	}
-	return
-}
-
-func isAWSFlag(args []string, i int) bool {
-	if i >= len(args) {
-		return false
-	}
-	return strings.HasPrefix(args[i], "--")
-}
-
 // awsApp is an AWS app that can start local proxies to serve AWS APIs.
 type awsApp struct {
 	*localProxyApp

diff --git a/tool/tsh/common/app_aws_test.go b/tool/tsh/common/app_aws_test.go
@@ -149,39 +149,6 @@ 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)
-	})
-	t.Run("aws ecs execute-command", 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, 13)
-			require.Equal(t, []string{"aws", "ecs", "execute-command", "--debug", "--cluster", "cluster-name", "--task", "task-name", "--command", "/bin/bash", "--interactive", "--endpoint-url"}, cmd.Args[:12])
-			require.Contains(t, cmd.Args[12], "127.0.0.1:")
-			return nil
-		}
-		err = Run(
-			context.Background(),
-			[]string{"aws", "ecs", "execute-command", "--debug", "--cluster", "cluster-name", "--task", "task-name", "--command", "/bin/bash", "--interactive"},
-			setHomePath(tmpHomePath),
-			setCmdRunner(validateCmd),
-		)
-		require.NoError(t, err)
-	})
 }
 
 // TestAWSConsoleLogins given a AWS console application, execute a app login