From 8d6e5706c7d61726c859083e9433c3ef8bd036cb Mon Sep 17 00:00:00 2001
From: Danil Uzlov <36223296+d-uzlov@users.noreply.github.com>
Date: Wed, 26 Mar 2025 10:31:43 +0000
Subject: [PATCH] close ssh connections on cleanup

---
 src/driver/controller-zfs-generic/index.js     |  4 +++-
 src/driver/freenas/ssh.js                      |  4 +++-
 src/driver/zfs-local-ephemeral-inline/index.js |  4 +++-
 src/utils/zfs_ssh_exec_client.js               | 12 ++++++++++++
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/driver/controller-zfs-generic/index.js b/src/driver/controller-zfs-generic/index.js
index 3e4836f..40a821d 100644
--- a/src/driver/controller-zfs-generic/index.js
+++ b/src/driver/controller-zfs-generic/index.js
@@ -16,10 +16,12 @@ class ControllerZfsGenericDriver extends ControllerZfsBaseDriver {
   getExecClient() {
     return this.ctx.registry.get(`${__REGISTRY_NS__}:exec_client`, () => {
       if (this.options.sshConnection) {
-        return new SshClient({
+        const sshClient = new SshClient({
           logger: this.ctx.logger,
           connection: this.options.sshConnection,
         });
+        this.cleanup.push(() => sshClient.finalize());
+        return sshClient;
       } else {
         return new LocalCliExecClient({
           logger: this.ctx.logger,
diff --git a/src/driver/freenas/ssh.js b/src/driver/freenas/ssh.js
index 3f635ac..55f8813 100644
--- a/src/driver/freenas/ssh.js
+++ b/src/driver/freenas/ssh.js
@@ -57,10 +57,12 @@ class FreeNASSshDriver extends ControllerZfsBaseDriver {
 
   getExecClient() {
     return this.ctx.registry.get(`${__REGISTRY_NS__}:exec_client`, () => {
-      return new SshClient({
+      const sshClient = new SshClient({
         logger: this.ctx.logger,
         connection: this.options.sshConnection,
       });
+      this.cleanup.push(() => sshClient.finalize());
+      return sshClient;
     });
   }
 
diff --git a/src/driver/zfs-local-ephemeral-inline/index.js b/src/driver/zfs-local-ephemeral-inline/index.js
index 276b1f4..3cec30e 100644
--- a/src/driver/zfs-local-ephemeral-inline/index.js
+++ b/src/driver/zfs-local-ephemeral-inline/index.js
@@ -125,10 +125,12 @@ class ZfsLocalEphemeralInlineDriver extends CsiBaseDriver {
 
   getSshClient() {
     return this.ctx.registry.get(`${__REGISTRY_NS__}:ssh_client`, () => {
-      return new SshClient({
+      const sshClient = new SshClient({
         logger: this.ctx.logger,
         connection: this.options.sshConnection,
       });
+      this.cleanup.push(() => sshClient.finalize());
+      return sshClient;
     });
   }
 
diff --git a/src/utils/zfs_ssh_exec_client.js b/src/utils/zfs_ssh_exec_client.js
index 4a232fa..8a5ee5d 100644
--- a/src/utils/zfs_ssh_exec_client.js
+++ b/src/utils/zfs_ssh_exec_client.js
@@ -75,8 +75,14 @@ class SshClient {
     const start_error_event_count = this.error_event_count;
     try {
       await this.conn_mutex.runExclusive(async () => {
+        if (this.finalized) {
+          throw 'using finalized ssh client';
+        }
         this.conn.connect(this.options.connection);
         do {
+          if (this.finalized) {
+            throw 'ssh client finalized during connection';
+          }
           if (start_error_event_count != this.error_event_count) {
             throw this.conn_err;
           }
@@ -104,6 +110,12 @@ class SshClient {
     return this._connect();
   }
 
+  finalize() {
+    this.finalized = true;
+    const conn = this.conn;
+    conn.end();
+  }
+
   async exec(command, options = {}, stream_proxy = null) {
     // default is to reuse
     if (process.env.SSH_REUSE_CONNECTION == "0") {