From c7b3bd9bc93490fdaaa96f42fc3f56ce6cc71ac1 Mon Sep 17 00:00:00 2001
From: Umputun <umputun@gmail.com>
Date: Thu, 19 Dec 2024 15:12:01 -0600
Subject: [PATCH] add missing cleanup for local temp script file

this file is uploaded to remote temp and removed properly as a part of returned teardown(). However the local file created first and the removal was missing
---
 pkg/executor/remote.go      |  2 +-
 pkg/runner/commands.go      |  9 ++++++++-
 pkg/runner/commands_test.go | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/pkg/executor/remote.go b/pkg/executor/remote.go
index 463c3e85..77f291b0 100644
--- a/pkg/executor/remote.go
+++ b/pkg/executor/remote.go
@@ -285,7 +285,7 @@ func (ex *Remote) Delete(ctx context.Context, remoteFile string, opts *DeleteOpt
 			}
 		}
 
-		log.Printf("[INFO] deleted recursively %s", remoteFile)
+		log.Printf("[INFO] deleted recursively %s from %s", remoteFile, ex.hostName)
 	}
 
 	if fileInfo.IsDir() && !recursive {
diff --git a/pkg/runner/commands.go b/pkg/runner/commands.go
index 783cdeab..2e3f614d 100644
--- a/pkg/runner/commands.go
+++ b/pkg/runner/commands.go
@@ -451,7 +451,12 @@ func (ec *execCmd) prepScript(ctx context.Context, s string, r io.Reader) (cmd,
 	if err = tmp.Close(); err != nil {
 		return "", "", nil, ec.errorFmt("can't close temporary file: %w", err)
 	}
-
+	defer func() {
+		// remove local copy of the script after upload or in case of error
+		if err := os.Remove(tmp.Name()); err != nil {
+			log.Printf("[WARN] can't remove local temp script %s: %v", tmp.Name(), err)
+		}
+	}()
 	// make the script executable locally, upload preserves the permissions
 	if err = os.Chmod(tmp.Name(), 0o700); err != nil { // nolint
 		return "", "", nil, ec.errorFmt("can't chmod temporary file: %w", err)
@@ -471,6 +476,8 @@ func (ec *execCmd) prepScript(ctx context.Context, s string, r io.Reader) (cmd,
 	cmd = fmt.Sprintf("%s -c %s", ec.shell(), dst)
 
 	teardown = func() error {
+		log.Printf("[DEBUG] removed local temp script %s", tmp.Name())
+
 		// remove the temp dir with the script from the remote hostAddr,
 		// should be invoked by the caller after the command is executed
 		if err := ec.exec.Delete(ctx, tmpRemoteDir, &executor.DeleteOpts{Recursive: true}); err != nil {
diff --git a/pkg/runner/commands_test.go b/pkg/runner/commands_test.go
index 3d096881..b779f9aa 100644
--- a/pkg/runner/commands_test.go
+++ b/pkg/runner/commands_test.go
@@ -483,6 +483,37 @@ func Test_execCmd(t *testing.T) {
 		assert.Equal(t, " {copy: testdata/inventory.yml -> /tmp/inventory.txt}", resp.details)
 	})
 
+	t.Run("script temp files cleanup", func(t *testing.T) {
+		wr := bytes.NewBuffer(nil)
+		log.SetOutput(io.MultiWriter(wr, os.Stdout))
+
+		// run a multi-line script that will create temp files
+		ec := execCmd{exec: sess, tsk: &config.Task{Name: "test"}, cmd: config.Cmd{
+			Script: "echo 'line1'\necho 'line2'\necho 'line3'",
+		}}
+		resp, err := ec.Script(ctx)
+		require.NoError(t, err)
+		t.Logf("resp: %+v", resp)
+
+		// check that the temp file was cleaned up
+		_, err = sess.Run(ctx, "ls -la /tmp/spot*", nil)
+		require.Error(t, err, "temp script file should be cleaned up")
+		_, err = sess.Run(ctx, "ls -la /tmp/.spot-*", nil)
+		require.Error(t, err, "temp dir should be cleaned up")
+
+		// also test cleanup on failure
+		ec = execCmd{exec: sess, tsk: &config.Task{Name: "test"}, cmd: config.Cmd{
+			Script: "echo 'line1'\nexit 1\necho 'line3'",
+		}}
+		_, err = ec.Script(ctx)
+		require.Error(t, err)
+
+		// check cleanup after failure
+		_, err = sess.Run(ctx, "ls -la /tmp/spot*", nil)
+		require.Error(t, err, "temp script file should be cleaned up even after failure")
+		_, err = sess.Run(ctx, "ls -la /tmp/.spot-*", nil)
+		require.Error(t, err, "temp dir should be cleaned up even after failure")
+	})
 }
 
 func Test_execCmdWithTmp(t *testing.T) {
@@ -501,7 +532,6 @@ func Test_execCmdWithTmp(t *testing.T) {
 		re := regexp.MustCompile(pattern)
 		match := re.FindStringSubmatch(log)
 		if len(match) > 1 {
-
 			return strings.ReplaceAll(match[1], "localhost:", "")
 		}
 		return ""