From 645caabc4fc728ebfede9df16e7ed6739d63a9b0 Mon Sep 17 00:00:00 2001
From: Dustin Strobel <ngvs_cacophonous@aleeas.com>
Date: Sun, 15 Oct 2023 15:44:03 +0200
Subject: [PATCH] feat: add context to ssh execution

---
 connection/ssh.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/connection/ssh.go b/connection/ssh.go
index 53c8572..74b590c 100644
--- a/connection/ssh.go
+++ b/connection/ssh.go
@@ -80,8 +80,14 @@ func (c *Connection) runSSH(ctx context.Context, cmd string) (string, string, er
 		return "", "", err
 	}
 
-	// Wait for the command to complete
-	err = s.Wait()
+	// Wait for the command to complete with context support
+	select {
+	case <-ctx.Done():
+		_ = s.Signal(ssh.SIGINT)
+		return "", "", ctx.Err()
+	default:
+		err = s.Wait()
+	}
 
 	// Return the error if stderr has no value
 	if err != nil && stderrBytes == nil {