From 0ce2b44b4de05d12f6788e6b87f6c093d3db09a1 Mon Sep 17 00:00:00 2001
From: Miguel Campos <amodestduck@gmail.com>
Date: Fri, 23 Feb 2024 17:16:05 -0800
Subject: [PATCH] SSH Command: add support for local port forwarding

---
 src/commands/ssh.ts          | 18 ++++++++++++++++++
 src/plugins/aws/ssm/index.ts | 26 ++++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/commands/ssh.ts b/src/commands/ssh.ts
index d1237dc..550a4ff 100644
--- a/src/commands/ssh.ts
+++ b/src/commands/ssh.ts
@@ -30,9 +30,13 @@ import yargs from "yargs";
 export type SshCommandArgs = {
   instance: string;
   command?: string;
+  L?: string;
   arguments: string[];
 };
 
+// Matches strings with the pattern "digits:digits" (e.g. 1234:5678)
+const LOCAL_PORT_FORWARD_PATTERN = /^\d+:\d+$/;
+
 /** Maximum amount of time to wait after access is approved to wait for access
  *  to be configured
  */
@@ -57,6 +61,20 @@ export const sshCommand = (yargs: yargs.Argv) =>
           array: true,
           string: true,
           default: [] as string[],
+        })
+        .check((argv: yargs.ArgumentsCamelCase<SshCommandArgs>) => {
+          if (argv.L == null) return true;
+
+          return (
+            argv.L.match(LOCAL_PORT_FORWARD_PATTERN) ||
+            "Local port forward should be in the format `local_port:remote_port`"
+          );
+        })
+        .option("L", {
+          type: "string",
+          describe:
+            // the order of the sockets in the address matches the ssh man page
+            "Forward a local port to the remote host; `local_socket:remote_socket`",
         }),
     guard(ssh)
   );
diff --git a/src/plugins/aws/ssm/index.ts b/src/plugins/aws/ssm/index.ts
index 896cbd8..fdb677a 100644
--- a/src/plugins/aws/ssm/index.ts
+++ b/src/plugins/aws/ssm/index.ts
@@ -37,6 +37,9 @@ const MAX_SSM_RETRIES = 30;
 const INSTANCE_ARN_PATTERN =
   /^arn:aws:ssm:([^:]+):([^:]+):managed-instance\/([^:]+)$/;
 
+/** The name of the SessionManager port forwarding document. This document is managed by AWS.  */
+const LOCAL_PORT_FORWARDING_DOCUMENT_NAME = "AWS-StartPortForwardingSession";
+
 type SsmArgs = {
   instance: string;
   region: string;
@@ -44,6 +47,7 @@ type SsmArgs = {
   documentName: string;
   credential: AwsCredentials;
   command?: string;
+  forwardPortAddress?: string;
 };
 
 /** Checks if access has propagated through AWS to the SSM agent
@@ -88,6 +92,12 @@ const accessPropagationGuard = (
 };
 
 const createSsmCommand = (args: Omit<SsmArgs, "requestId">) => {
+  const hasCommand = args.command && args.command.trim();
+
+  if (hasCommand && args.forwardPortAddress) {
+    throw "Invalid arguments. Specify either a command or port forwarding, not both.";
+  }
+
   const ssmCommand = [
     "aws",
     "ssm",
@@ -97,11 +107,22 @@ const createSsmCommand = (args: Omit<SsmArgs, "requestId">) => {
     "--target",
     args.instance,
     "--document-name",
-    args.documentName,
+    // Port forwarding is a special case that uses an AWS-managed document and
+    // not the user-generated document we use for an SSH session
+    args.forwardPortAddress
+      ? LOCAL_PORT_FORWARDING_DOCUMENT_NAME
+      : args.documentName,
   ];
 
-  if (args.command && args.command.trim()) {
+  if (hasCommand) {
     ssmCommand.push("--parameters", `command='${args.command}'`);
+  } else if (args.forwardPortAddress) {
+    const [localPort, remotePort] = args.forwardPortAddress.split(":");
+
+    ssmCommand.push(
+      "--parameters",
+      `localPortNumber=${localPort},portNumber=${remotePort}`
+    );
   }
 
   return ssmCommand;
@@ -193,6 +214,7 @@ export const ssm = async (
     region: region!,
     documentName: request.generated.documentName,
     requestId: request.id,
+    forwardPortAddress: args.L,
     credential,
     command: commandParameter(args),
   };