From 44aebc255964c98adb9ed3b40ddc10016ec334b0 Mon Sep 17 00:00:00 2001
From: Miguel Campos <miguel.campos@p0.dev>
Date: Mon, 26 Feb 2024 15:01:00 -0800
Subject: [PATCH] update error handling

---
 src/commands/ssh.ts    | 12 +++++++++++-
 src/plugins/aws/ssm.ts | 15 ++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/commands/ssh.ts b/src/commands/ssh.ts
index f662d33..5cf9bba 100644
--- a/src/commands/ssh.ts
+++ b/src/commands/ssh.ts
@@ -27,6 +27,8 @@ import { getDoc, onSnapshot } from "firebase/firestore";
 import { pick } from "lodash";
 import yargs from "yargs";
 
+const LOCAL_PORT_FORWARD_PATTERN = /^\d+:\d+$/;
+
 type SshCommandArgs = {
   instance: string;
   command?: string;
@@ -59,11 +61,19 @@ export const sshCommand = (yargs: yargs.Argv) =>
           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"]',
+            "Forward a local port to the remote host; `local_socket:remote_socket`",
         }),
     guard(ssh)
   );
diff --git a/src/plugins/aws/ssm.ts b/src/plugins/aws/ssm.ts
index 3ab4316..f5973d4 100644
--- a/src/plugins/aws/ssm.ts
+++ b/src/plugins/aws/ssm.ts
@@ -90,6 +90,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",
@@ -106,15 +112,10 @@ const createSsmCommand = (args: Omit<SsmArgs, "requestId">) => {
       : 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(":")
-      .map(Number);
-    if (!localPort || isNaN(localPort) || !remotePort || isNaN(remotePort)) {
-      throw "Invalid port forwarding address specified. Please use format <local_socket>:<remote_socket>";
-    }
+    const [localPort, remotePort] = args.forwardPortAddress.split(":");
 
     ssmCommand.push(
       "--parameters",