From a3eee286198ef020b209f0c523d13ceda8826ea5 Mon Sep 17 00:00:00 2001
From: Tim Lee <timlee@fleetdm.com>
Date: Tue, 15 Oct 2024 15:48:56 -0600
Subject: [PATCH] Use WaitDelay in non_windows script timeout (#22912)

---
 changes/20248-non-windows-script-timeout  |  1 +
 orbit/pkg/scripts/exec_nonwindows.go      |  9 ++-
 orbit/pkg/scripts/exec_nonwindows_test.go | 75 +++++++++++++++++++++++
 3 files changed, 83 insertions(+), 2 deletions(-)
 create mode 100644 changes/20248-non-windows-script-timeout

diff --git a/changes/20248-non-windows-script-timeout b/changes/20248-non-windows-script-timeout
new file mode 100644
index 0000000000..743694b656
--- /dev/null
+++ b/changes/20248-non-windows-script-timeout
@@ -0,0 +1 @@
+- fixed issue where macOS and Linux scripts failed to timeout on long running commands
\ No newline at end of file
diff --git a/orbit/pkg/scripts/exec_nonwindows.go b/orbit/pkg/scripts/exec_nonwindows.go
index 069e2d676b..bd4a143c46 100644
--- a/orbit/pkg/scripts/exec_nonwindows.go
+++ b/orbit/pkg/scripts/exec_nonwindows.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"time"
 
 	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
 	"github.com/fleetdm/fleet/v4/server/fleet"
@@ -28,7 +29,7 @@ func ExecCmd(ctx context.Context, scriptPath string, env []string) (output []byt
 	cmd := exec.CommandContext(ctx, "/bin/sh", scriptPath)
 
 	if directExecute {
-		err = os.Chmod(scriptPath, 0700)
+		err = os.Chmod(scriptPath, 0o700)
 		if err != nil {
 			return nil, -1, ctxerr.Wrapf(ctx, err, "marking script as executable %s", scriptPath)
 		}
@@ -40,8 +41,12 @@ func ExecCmd(ctx context.Context, scriptPath string, env []string) (output []byt
 	}
 
 	cmd.Dir = filepath.Dir(scriptPath)
+
+	// WaitDelay is necessary to ensure that the process is killed when the
+	// context is cancelled
+	cmd.WaitDelay = time.Second
 	output, err = cmd.CombinedOutput()
-	if cmd.ProcessState != nil {
+	if cmd.ProcessState != nil && ctx.Err() == nil {
 		exitCode = cmd.ProcessState.ExitCode()
 	}
 	return output, exitCode, err
diff --git a/orbit/pkg/scripts/exec_nonwindows_test.go b/orbit/pkg/scripts/exec_nonwindows_test.go
index 891273c4e0..360efea855 100644
--- a/orbit/pkg/scripts/exec_nonwindows_test.go
+++ b/orbit/pkg/scripts/exec_nonwindows_test.go
@@ -4,12 +4,14 @@ package scripts
 
 import (
 	"context"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	"github.com/stretchr/testify/require"
@@ -78,3 +80,76 @@ func TestExecCmdNonWindows(t *testing.T) {
 		})
 	}
 }
+
+func writeTestScript(content string) (string, error) {
+	tmpfile, err := ioutil.TempFile("", "testscript*.sh")
+	if err != nil {
+		return "", err
+	}
+
+	if _, err := tmpfile.Write([]byte(content)); err != nil {
+		tmpfile.Close()
+		return "", err
+	}
+	if err := tmpfile.Close(); err != nil {
+		return "", err
+	}
+
+	err = os.Chmod(tmpfile.Name(), 0o700)
+	if err != nil {
+		return "", err
+	}
+
+	return tmpfile.Name(), nil
+}
+
+func TestExecCmdTimeout(t *testing.T) {
+	scriptContent := `#!/bin/sh
+	sleep 5
+	echo "Finished"`
+	scriptPath, err := writeTestScript(scriptContent)
+	if err != nil {
+		t.Fatalf("Failed to write test script: %v", err)
+	}
+	defer os.Remove(scriptPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	start := time.Now()
+	output, exitCode, err := ExecCmd(ctx, scriptPath, nil)
+	require.NotNil(t, err)
+	require.Contains(t, err.Error(), "signal: killed")
+	if exitCode != -1 {
+		t.Fatalf("Expected exit code -1, got: %d", exitCode)
+	}
+	if len(output) != 0 {
+		t.Fatalf("Expected no output, got: %s", output)
+	}
+	require.True(t, time.Since(start) <= 5*time.Second)
+}
+
+func TestExecCmdSuccess(t *testing.T) {
+	scriptContent := `#!/bin/sh
+	echo "Hello, World!"`
+	scriptPath, err := writeTestScript(scriptContent)
+	if err != nil {
+		t.Fatalf("Failed to write test script: %v", err)
+	}
+	defer os.Remove(scriptPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	output, exitCode, err := ExecCmd(ctx, scriptPath, nil)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+	if exitCode != 0 {
+		t.Fatalf("Expected exit code 0, got: %d", exitCode)
+	}
+	expectedOutput := "Hello, World!\n"
+	if string(output) != expectedOutput {
+		t.Fatalf("Expected output %q, got: %q", expectedOutput, output)
+	}
+}