From 7d52709e94dfadb6686993cd4e6d98a34184ae47 Mon Sep 17 00:00:00 2001
From: Miguel Campos <amodestduck@gmail.com>
Date: Fri, 23 Feb 2024 17:33:10 -0800
Subject: [PATCH] remove coerce

---
 src/commands/ssh.ts      | 26 +++++++++-----------------
 src/plugins/aws/ssm.ts   | 32 ++++++++++++++++++++++----------
 src/plugins/aws/types.ts |  5 +----
 3 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/src/commands/ssh.ts b/src/commands/ssh.ts
index 88949e2..349d974 100644
--- a/src/commands/ssh.ts
+++ b/src/commands/ssh.ts
@@ -12,7 +12,7 @@ import { authenticate } from "../drivers/auth";
 import { doc, guard } from "../drivers/firestore";
 import { print2 } from "../drivers/stdio";
 import { ssm } from "../plugins/aws/ssm";
-import { AwsSsh, SshCommandArgs } from "../plugins/aws/types";
+import { AwsSsh } from "../plugins/aws/types";
 import { SshConfig } from "../plugins/ssh/types";
 import { Authn } from "../types/identity";
 import {
@@ -32,6 +32,13 @@ import yargs from "yargs";
  */
 const GRANT_TIMEOUT_MILLIS = 60e3;
 
+type SshCommandArgs = {
+  instance: string;
+  command?: string;
+  L?: string;
+  arguments: string[];
+};
+
 export const sshCommand = (yargs: yargs.Argv) =>
   yargs.command<SshCommandArgs>(
     "ssh <instance> [command [arguments..]]",
@@ -57,21 +64,6 @@ export const sshCommand = (yargs: yargs.Argv) =>
           describe:
             // the order of the sockets in the address matche the ssh man page
             'Forward a local port to the remote host ["local_socket:remote_socket"]',
-          coerce: (arg) => {
-            if (!arg) return undefined;
-            const [localPort, remotePort] = arg.split(":").map(Number);
-            if (
-              !localPort ||
-              isNaN(localPort) ||
-              !remotePort ||
-              isNaN(remotePort)
-            ) {
-              throw new Error(
-                "Invalid port forwarding address specified. Please use format <local_socket>:<remote_socket>"
-              );
-            }
-            return arg;
-          },
         }),
     guard(ssh)
   );
@@ -163,7 +155,7 @@ const ssh = async (args: yargs.ArgumentsCamelCase<SshCommandArgs>) => {
   await ssm(authn, {
     ...requestData,
     id,
-    forwardPorts: args.forwardPorts,
+    forwardPortAddress: args.L,
     command: args.command
       ? `${args.command} ${args.arguments
           .map(
diff --git a/src/plugins/aws/ssm.ts b/src/plugins/aws/ssm.ts
index d1267dd..0384319 100644
--- a/src/plugins/aws/ssm.ts
+++ b/src/plugins/aws/ssm.ts
@@ -12,7 +12,7 @@ import { print2 } from "../../drivers/stdio";
 import { Authn } from "../../types/identity";
 import { Request } from "../../types/request";
 import { assumeRoleWithOktaSaml } from "../okta/aws";
-import { AwsCredentials, AwsSsh, SshCommandArgs } from "./types";
+import { AwsCredentials, AwsSsh } from "./types";
 import { ChildProcessByStdio, spawn } from "node:child_process";
 import { Readable } from "node:stream";
 
@@ -44,7 +44,9 @@ type SsmArgs = {
   requestId: string;
   documentName: string;
   credential: AwsCredentials;
-} & Pick<SshCommandArgs, "command" | "forwardPorts" | "instance">;
+  command?: string;
+  forwardPortAddress?: string;
+};
 
 /** Checks if access has propagated through AWS to the SSM agent
  *
@@ -99,15 +101,24 @@ const createSsmCommand = (args: Omit<SsmArgs, "requestId">) => {
     "--document-name",
     // Port forwarding is a special case that uses an AWS-managed document and
     // not the user-generated document we use for our other SSH sessions
-    args.forwardPorts ? LOCAL_PORT_FORWARDING_DOCUMENT_NAME : args.documentName,
+    args.forwardPortAddress
+      ? LOCAL_PORT_FORWARDING_DOCUMENT_NAME
+      : args.documentName,
   ];
 
   if (args.command && args.command.trim()) {
     ssmCommand.push("--parameters", `command='${args.command}'`);
-  } else if (args.forwardPorts) {
+  } else if (args.forwardPortAddress) {
+    const [localPort, remotePort] = args.forwardPortAddress
+      .split(":")
+      .map(Number);
+    if (!localPort || isNaN(localPort) || !remotePort || isNaN(remotePort)) {
+      throw "Invalid port forwarding address specified. Please use format <local_socket>:<remote_socket>";
+    }
+
     ssmCommand.push(
       "--parameters",
-      `localPortNumber=${args.forwardPorts.local},portNumber=${args.forwardPorts.remote}`
+      `localPortNumber=${localPort},portNumber=${remotePort}`
     );
   }
 
@@ -165,10 +176,11 @@ const spawnSsmNode = async (
 /** Connect to an SSH backend using AWS Systems Manager (SSM) */
 export const ssm = async (
   authn: Authn,
-  request: Request<AwsSsh> &
-    Pick<SshCommandArgs, "command" | "forwardPorts"> & {
-      id: string;
-    }
+  request: Request<AwsSsh> & {
+    command?: string;
+    forwardPortAddress?: string;
+    id: string;
+  }
 ) => {
   const match = request.permission.spec.arn.match(INSTANCE_ARN_PATTERN);
   if (!match) throw "Did not receive a properly formatted instance identifier";
@@ -183,7 +195,7 @@ export const ssm = async (
     region: region!,
     documentName: request.generated.documentName,
     requestId: request.id,
-    forwardPorts: request.forwardPorts,
+    forwardPortAddress: request.forwardPortAddress,
     credential,
     command: request.command,
   };
diff --git a/src/plugins/aws/types.ts b/src/plugins/aws/types.ts
index f7cde34..d463980 100644
--- a/src/plugins/aws/types.ts
+++ b/src/plugins/aws/types.ts
@@ -58,9 +58,6 @@ export type AwsSsh = {
 export type SshCommandArgs = {
   instance: string;
   command?: string;
-  forwardPorts?: {
-    local: string;
-    remote: string;
-  };
+  forwardPortAddress?: string;
   arguments: string[];
 };