From b19875914e92e85496f16e82cd5ac74a3cf8e93d Mon Sep 17 00:00:00 2001
From: Miguel Campos <amodestduck@gmail.com>
Date: Fri, 23 Feb 2024 09:23:12 -0800
Subject: [PATCH] only use one ssm document

---
 src/commands/ssh.ts      | 14 ++++++------
 src/plugins/aws/ssm.ts   | 46 +++++++++++++++-------------------------
 src/plugins/aws/types.ts |  5 +----
 3 files changed, 25 insertions(+), 40 deletions(-)

diff --git a/src/commands/ssh.ts b/src/commands/ssh.ts
index 11d07e24..f6e9451c 100644
--- a/src/commands/ssh.ts
+++ b/src/commands/ssh.ts
@@ -135,14 +135,14 @@ const ssh = async (args: yargs.ArgumentsCamelCase<SshCommandArgs>) => {
   await ssm(authn, {
     ...requestData,
     id,
-    // the command to run on the remote machine, if any
     command: args.command
-      ? `${args.command} ${args.arguments.map(shellEscapeArgument).join(" ")}`.trim()
+      ? `${args.command} ${args.arguments
+          .map(
+            (argument) =>
+              // escape all double quotes (") in commands such as `p0 ssh <instance>> echo 'hello; "world"'`
+              `"${argument.replace(/"/g, '\\"')}"`
+          )
+          .join(" ")}`.trim()
       : undefined,
   });
 };
-
-// Helper function to support double quotes (") in commands such as `p0 ssh <instance>> echo 'hi; "mom"'`
-const shellEscapeArgument = (argument: string) => {
-  return `"${argument.replace('"', '\\"')}"`;
-};
diff --git a/src/plugins/aws/ssm.ts b/src/plugins/aws/ssm.ts
index cb03c743..b7a30379 100644
--- a/src/plugins/aws/ssm.ts
+++ b/src/plugins/aws/ssm.ts
@@ -29,10 +29,7 @@ type SsmArgs = {
   instance: string;
   region: string;
   requestId: string;
-  documentNames: {
-    session: string;
-    command: string;
-  };
+  documentName: string;
   credential: AwsCredentials;
   command?: string;
 };
@@ -78,29 +75,6 @@ const accessPropagationGuard = (
   };
 };
 
-/**
- * Creates an SSM command to start a session or run a command on an instance.
- * Selects the appropriate SSM document based on the presence of a command.
- */
-const buildSsmCommand = (args: Omit<SsmArgs, "requestId">) => {
-  const ssmCommand = [
-    "aws",
-    "ssm",
-    "start-session",
-    "--region",
-    args.region,
-    "--target",
-    args.instance,
-  ];
-  if (args.command) {
-    ssmCommand.push("--document-name", args.documentNames.command);
-    ssmCommand.push("--parameters", `command='${args.command}'`);
-  } else {
-    ssmCommand.push("--document-name", args.documentNames.session);
-  }
-  return ssmCommand;
-};
-
 /** Starts an SSM session in the terminal by spawning `aws ssm` as a subprocess
  *
  * Requires `aws ssm` to be installed on the client machine.
@@ -110,7 +84,21 @@ const spawnSsmNode = async (
   options?: { attemptsRemaining?: number }
 ): Promise<number | null> =>
   new Promise((resolve, reject) => {
-    const ssmCommand = buildSsmCommand(args);
+    const ssmCommand = [
+      "aws",
+      "ssm",
+      "start-session",
+      "--region",
+      args.region,
+      "--target",
+      args.instance,
+      "--document-name",
+      args.documentName,
+      "--parameters",
+      // The empty string avoids the validation in the SSM document that doesn't allow zero-length
+      // parameter values when running the session document.
+      `command='${args.command || " "}'`,
+    ];
     const child = spawn("/usr/bin/env", ssmCommand, {
       env: {
         ...process.env,
@@ -165,7 +153,7 @@ export const ssm = async (
   const args = {
     instance: instance!,
     region: region!,
-    documentNames: request.generated.documentNames,
+    documentName: request.generated.documentName,
     requestId: request.id,
     credential,
     command: request.command,
diff --git a/src/plugins/aws/types.ts b/src/plugins/aws/types.ts
index 6c6e56c3..e5e2fe71 100644
--- a/src/plugins/aws/types.ts
+++ b/src/plugins/aws/types.ts
@@ -41,9 +41,6 @@ export type AwsSsh = {
     type: "session";
   };
   generated: {
-    documentNames: {
-      session: string;
-      command: string;
-    };
+    documentName: string;
   };
 };