impl: identify the PID for the SSH command

fioan89 · fioan89 · commit c5c1850995e3 · 2025-05-09T23:51:09.000+03:00
Toolbox spawns a native process running the SSH client. The ssh client then spawns
another process which associated to the coder proxy command. SSH network metrics
are saved into json file with the name equal to the pid of the ssh command (not
to be confused with the proxy command's name)
diff --git a/src/main/kotlin/com/coder/toolbox/cli/SshCommandProcessHandle.kt b/src/main/kotlin/com/coder/toolbox/cli/SshCommandProcessHandle.kt
@@ -0,0 +1,42 @@
+package com.coder.toolbox.cli
+
+import com.coder.toolbox.CoderToolboxContext
+import com.coder.toolbox.sdk.v2.models.Workspace
+import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
+import kotlin.jvm.optionals.getOrNull
+
+/**
+ * Identifies the PID for the SSH Coder command spawned by Toolbox.
+ */
+class SshCommandProcessHandle(private val ctx: CoderToolboxContext) {
+
+    /**
+     * Finds the PID of a Coder (not the proxy command) ssh cmd associated with the specified workspace and agent.
+     * Null is returned when no ssh command process was found.
+     *
+     * Implementation Notes:
+     * An iterative DFS approach where we start with Toolbox's direct children, grep the command
+     * and if nothing is found we continue with the processes children. Toolbox spawns an ssh command
+     * as a separate command which in turns spawns another child for the proxy command.
+     */
+    fun findByWorkspaceAndAgent(ws: Workspace, agent: WorkspaceAgent): Long? {
+        val stack = ArrayDeque<ProcessHandle>(ProcessHandle.current().children().toList())
+        while (stack.isNotEmpty()) {
+            val processHandle = stack.removeLast()
+            val cmdLine = processHandle.info().commandLine().getOrNull()
+            ctx.logger.debug("SSH command PID: ${processHandle.pid()} Command: $cmdLine")
+            if (cmdLine != null && cmdLine.isSshCommandFor(ws, agent)) {
+                ctx.logger.debug("SSH command with PID: ${processHandle.pid()} and Command: $cmdLine matches ${ws.name}.${agent.name}")
+                return processHandle.pid()
+            } else {
+                stack.addAll(processHandle.children().toList())
+            }
+        }
+        return null
+    }
+
+    private fun String.isSshCommandFor(ws: Workspace, agent: WorkspaceAgent): Boolean {
+        // usage-app is present only in the ProxyCommand
+        return !this.contains("--usage-app=jetbrains") && this.contains("${ws.name}.${agent.name}")
+    }
+}
