Script Timeout Agent Options Part 1 of 2 (#20266)

cmd/fleetctl/scripts_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/fleetdm/fleet/v4/pkg/scripts"
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	"github.com/fleetdm/fleet/v4/server/ptr"
 	"github.com/fleetdm/fleet/v4/server/service"
@@ -270,10 +271,10 @@ Output:
 			scriptResult: &fleet.HostScriptResult{
 				ExitCode: ptr.Int64(-1),
 				Output:   "Oh no!",
-				Message:  fleet.RunScriptScriptTimeoutErrMsg,
+				Message:  fleet.HostScriptTimeoutMessage(ptr.Int(int(scripts.MaxHostExecutionTime.Seconds()))),
 			},
 			expectOutput: `
-Error: Timeout. Fleet stopped the script after 5 minutes to protect host performance.
+Error: Timeout. Fleet stopped the script after 300 seconds to protect host performance.
 
 Output before timeout:
 

orbit/pkg/scripts/scripts.go

@@ -13,7 +13,6 @@ import (
 	"unicode/utf8"
 
 	"github.com/fleetdm/fleet/v4/orbit/pkg/constant"
-	"github.com/fleetdm/fleet/v4/pkg/scripts"
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	"github.com/rs/zerolog/log"
 )
@@ -31,6 +30,7 @@ type Client interface {
 type Runner struct {
 	Client                 Client
 	ScriptExecutionEnabled bool
+	ScriptExecutionTimeout time.Duration
 
 	// tempDirFn is the function to call to get the temporary directory to use,
 	// inside of which the script-specific subdirectories will be created. If nil,
@@ -114,7 +114,7 @@ func (r *Runner) runOne(script *fleet.HostScriptResult) (finalErr error) {
 		return fmt.Errorf("write script file: %w", err)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), scripts.MaxHostExecutionTime)
+	ctx, cancel := context.WithTimeout(context.Background(), r.ScriptExecutionTimeout)
 	defer cancel()
 
 	execCmdFn := r.execCmdFn
@@ -122,7 +122,7 @@ func (r *Runner) runOne(script *fleet.HostScriptResult) (finalErr error) {
 		execCmdFn = ExecCmd
 	}
 	start := time.Now()
-	log.Debug().Msgf("starting script execution of %v", script.ExecutionID)
+	log.Debug().Msgf("starting script execution of %v with timeout of %v", script.ExecutionID, r.ScriptExecutionTimeout)
 	output, exitCode, execErr := execCmdFn(ctx, scriptFile, nil)
 	log.Debug().Msgf("after script execution of %v", script.ExecutionID)
 	duration := time.Since(start)
@@ -144,6 +144,7 @@ func (r *Runner) runOne(script *fleet.HostScriptResult) (finalErr error) {
 		Output:      string(output),
 		Runtime:     int(duration.Seconds()),
 		ExitCode:    exitCode,
+		Timeout:     int(r.ScriptExecutionTimeout.Seconds()),
 	})
 	if err != nil {
 		return fmt.Errorf("save script result: %w", err)

orbit/pkg/update/notifications.go

@@ -10,6 +10,7 @@ import (
 	"github.com/fleetdm/fleet/v4/orbit/pkg/bitlocker"
 	"github.com/fleetdm/fleet/v4/orbit/pkg/profiles"
 	"github.com/fleetdm/fleet/v4/orbit/pkg/scripts"
+	fleetscripts "github.com/fleetdm/fleet/v4/pkg/scripts"
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	"github.com/rs/zerolog/log"
 )
@@ -346,6 +347,11 @@ func (h *runScriptsConfigReceiver) runDynamicScriptsEnabledCheck() {
 // server sent a list of scripts to execute, starts a goroutine to execute
 // them.
 func (h *runScriptsConfigReceiver) Run(cfg *fleet.OrbitConfig) error {
+	timeout := fleetscripts.MaxHostExecutionTime
+	if cfg.ScriptExeTimeout > 0 {
+		timeout = time.Duration(cfg.ScriptExeTimeout) * time.Second
+	}
+
 	if len(cfg.Notifications.PendingScriptExecutionIDs) > 0 {
 		if h.mu.TryLock() {
 			log.Debug().Msgf("received request to run scripts %v", cfg.Notifications.PendingScriptExecutionIDs)
@@ -353,6 +359,7 @@ func (h *runScriptsConfigReceiver) Run(cfg *fleet.OrbitConfig) error {
 			runner := &scripts.Runner{
 				ScriptExecutionEnabled: h.scriptsEnabled(),
 				Client:                 h.ScriptsClient,
+				ScriptExecutionTimeout: timeout,
 			}
 			fn := runner.Run
 			if h.runScriptsFn != nil {

server/datastore/mysql/migrations/tables/20240709183940_AddTimeoutColumnToScriptResults.go

@@ -0,0 +1,38 @@
+package tables
+
+import (
+	"database/sql"
+)
+
+func init() {
+	MigrationClient.AddMigration(Up_20240709183940, Down_20240709183940)
+}
+
+// At the time of this migration, all script timeouts are
+// hardcoded to 300 seconds. This migration adds a timeout
+// column to the host_script_results table to allow for
+// custom timeouts.
+func Up_20240709183940(tx *sql.Tx) error {
+	stmt := `
+	ALTER TABLE host_script_results
+	ADD COLUMN timeout INT DEFAULT NULL;
+	`
+	if _, err := tx.Exec(stmt); err != nil {
+		return err
+	}
+
+	stmt = `
+	UPDATE host_script_results
+	SET timeout = 300
+	WHERE timeout IS NULL;
+	`
+	if _, err := tx.Exec(stmt); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func Down_20240709183940(tx *sql.Tx) error {
+	return nil
+}

server/datastore/mysql/migrations/tables/20240709183940_AddTimeoutColumnToScriptResults_test.go

@@ -0,0 +1,34 @@
+package tables
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUp_20240709183940(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	insertStmt := `
+		INSERT INTO host_script_results 
+		(host_id, execution_id, output)
+		VALUES (?, ?, ?)
+		`
+	_, err := db.Exec(insertStmt, 1, 1, "output")
+	require.NoError(t, err)
+
+	applyNext(t, db)
+
+	selectStmt := `
+		SELECT timeout FROM host_script_results
+		WHERE host_id = ?
+		`
+	var timeout int
+	err = db.QueryRow(selectStmt, 1).Scan(&timeout)
+	require.NoError(t, err)
+	require.Equal(t, 300, timeout)
+
+	// inserting no timeout succeeds
+	_, err = db.Exec(insertStmt, 2, 2, "output")
+	require.NoError(t, err)
+}

server/datastore/mysql/scripts.go

@@ -96,7 +96,8 @@ func (ds *Datastore) SetHostScriptExecutionResult(ctx context.Context, result *f
   UPDATE host_script_results SET
     output = ?,
     runtime = ?,
-    exit_code = ?
+    exit_code = ?,
+    timeout = ?
   WHERE
     host_id = ? AND
     execution_id = ?`
@@ -138,6 +139,7 @@ func (ds *Datastore) SetHostScriptExecutionResult(ctx context.Context, result *f
 			// it to a 32-bit signed integer.
 			// See /orbit/pkg/scripts/exec_windows.go
 			int32(result.ExitCode),
+			result.Timeout,
 			result.HostID,
 			result.ExecutionID,
 		)
@@ -236,6 +238,7 @@ func (ds *Datastore) getHostScriptExecutionResultDB(ctx context.Context, q sqlx.
     hsr.output,
     hsr.runtime,
     hsr.exit_code,
+    hsr.timeout,
     hsr.created_at,
     hsr.user_id,
     hsr.sync_request,

server/datastore/mysql/scripts_test.go

@@ -93,6 +93,7 @@ func testHostScriptResult(t *testing.T, ds *Datastore) {
 		Output:      "foo",
 		Runtime:     2,
 		ExitCode:    0,
+		Timeout:     300,
 	})
 	require.NoError(t, err)
 
@@ -103,6 +104,7 @@ func testHostScriptResult(t *testing.T, ds *Datastore) {
 		Output:      "foobarbaz",
 		Runtime:     22,
 		ExitCode:    1,
+		Timeout:     360,
 	})
 	require.NoError(t, err)
 	require.Nil(t, hsr)
@@ -119,6 +121,7 @@ func testHostScriptResult(t *testing.T, ds *Datastore) {
 	expectScript.Output = "foo"
 	expectScript.Runtime = 2
 	expectScript.ExitCode = ptr.Int64(0)
+	expectScript.Timeout = ptr.Int(300)
 	require.Equal(t, &expectScript, script)
 
 	// create another script execution request (null user id this time)
@@ -166,6 +169,7 @@ func testHostScriptResult(t *testing.T, ds *Datastore) {
 		Output:      largeOutput,
 		Runtime:     10,
 		ExitCode:    1,
+		Timeout:     300,
 	})
 	require.NoError(t, err)
 
@@ -240,6 +244,7 @@ func testHostScriptResult(t *testing.T, ds *Datastore) {
 		Output:      "foo",
 		Runtime:     1,
 		ExitCode:    math.MaxUint32,
+		Timeout:     300,
 	})
 	require.NoError(t, err)
 	require.EqualValues(t, -1, *unsignedScriptResult.ExitCode)

server/fleet/agent_options.go

@@ -11,7 +11,11 @@ import (
 
 //go:generate go run ../../tools/osquery-agent-options agent_options_generated.go
 
+const maxAgentScriptExecutionTimeout = 3600
+
 type AgentOptions struct {
+	// ScriptExecutionTimeout is the maximum time in seconds that a script can run.
+	ScriptExecutionTimeout int `json:"script_execution_timeout,omitempty"`
 	// Config is the base config options.
 	Config json.RawMessage `json:"config"`
 	// Overrides includes any platform-based overrides.
@@ -49,6 +53,10 @@ func ValidateJSONAgentOptions(ctx context.Context, ds Datastore, rawJSON json.Ra
 		return err
 	}
 
+	if opts.ScriptExecutionTimeout > maxAgentScriptExecutionTimeout {
+		return fmt.Errorf("'script_execution_timeout' value exceeds limit. Maximum value is %d", maxAgentScriptExecutionTimeout)
+	}
+
 	if len(opts.CommandLineStartUpFlags) > 0 {
 		var flags osqueryCommandLineFlags
 		if err := JSONStrictDecode(bytes.NewReader(opts.CommandLineStartUpFlags), &flags); err != nil {

server/fleet/agent_options_test.go

@@ -28,6 +28,10 @@ func TestValidateAgentOptions(t *testing.T) {
 			}
 		}}`, true, `unknown field "foo"`},
 
+		{"valid script timeout", `{"script_execution_timeout": 600}`, true, ""},
+
+		{"invalid script timeout", `{"script_execution_timeout": 3601}`, true, `script_execution_timeout' value exceeds limit. Maximum value is 3600`},
+
 		{"overrides.platform is null", `{"overrides": {
 			"platforms": {
 				"darwin": null

server/fleet/errors.go

@@ -550,7 +550,6 @@ const (
 	RunScriptScriptsDisabledGloballyErrMsg = "Running scripts is disabled in organization settings."
 	RunScriptDisabledErrMsg                = "Scripts are disabled for this host. To run scripts, deploy the fleetd agent with scripts enabled."
 	RunScriptsOrbitDisabledErrMsg          = "Couldn't run script. To run a script, deploy the fleetd agent with --enable-scripts."
-	RunScriptScriptTimeoutErrMsg           = "Timeout. Fleet stopped the script after 5 minutes to protect host performance."
 	RunScriptAsyncScriptEnqueuedErrMsg     = "Script is running or will run when the host comes online."
 	RunScripSavedMaxLenErrMsg              = "Script is too large. It's limited to 500,000 characters (approximately 10,000 lines)."
 	RunScripUnsavedMaxLenErrMsg            = "Script is too large. It's limited to 10,000 characters (approximately 125 lines)."

server/fleet/orbit.go

@@ -39,10 +39,11 @@ type OrbitConfigNotifications struct {
 }
 
 type OrbitConfig struct {
-	Flags         json.RawMessage          `json:"command_line_startup_flags,omitempty"`
-	Extensions    json.RawMessage          `json:"extensions,omitempty"`
-	NudgeConfig   *NudgeConfig             `json:"nudge_config,omitempty"`
-	Notifications OrbitConfigNotifications `json:"notifications,omitempty"`
+	ScriptExeTimeout int                      `json:"script_execution_timeout,omitempty"`
+	Flags            json.RawMessage          `json:"command_line_startup_flags,omitempty"`
+	Extensions       json.RawMessage          `json:"extensions,omitempty"`
+	NudgeConfig      *NudgeConfig             `json:"nudge_config,omitempty"`
+	Notifications    OrbitConfigNotifications `json:"notifications,omitempty"`
 	// UpdateChannels contains the TUF channels to use on fleetd components.
 	//
 	// If UpdateChannels is nil it means the server isn't using/setting this feature.

server/fleet/scripts.go

@@ -3,6 +3,7 @@ package fleet
 import (
 	"bufio"
 	"errors"
+	"fmt"
 	"path/filepath"
 	"regexp"
 	"strings"
@@ -195,6 +196,7 @@ type HostScriptResultPayload struct {
 	Output      string `json:"output"`
 	Runtime     int    `json:"runtime"`
 	ExitCode    int    `json:"exit_code"`
+	Timeout     int    `json:"timeout"`
 }
 
 // HostScriptResult represents a script result that was requested to execute on
@@ -218,6 +220,9 @@ type HostScriptResult struct {
 	// host. It is -1 if it was received but the script did not terminate
 	// normally (same as how Go handles this: https://pkg.go.dev/os#ProcessState.ExitCode)
 	ExitCode *int64 `json:"exit_code" db:"exit_code"`
+	// Timeout is the maximum time in seconds that the script was allowed to run
+	// at the time of execution.
+	Timeout *int `json:"timeout" db:"timeout"`
 	// CreatedAt is the creation timestamp of the script execution request. It is
 	// not returned as part of the payloads, but is used to determine if the script
 	// is too old to still expect a response from the host.
@@ -266,7 +271,7 @@ func (hsr HostScriptResult) AuthzType() string {
 // for running a script synchronously (so that fleetctl can display it) and to
 // get the script results for an execution ID (e.g. when looking at the details
 // screen of a script execution activity in the website).
-func (hsr HostScriptResult) UserMessage(hostTimeout bool) string {
+func (hsr HostScriptResult) UserMessage(hostTimeout bool, hostTimeoutValue *int) string {
 	if hostTimeout {
 		return RunScriptHostTimeoutErrMsg
 	}
@@ -285,7 +290,7 @@ func (hsr HostScriptResult) UserMessage(hostTimeout bool) string {
 
 	switch *hsr.ExitCode {
 	case -1:
-		return RunScriptScriptTimeoutErrMsg
+		return HostScriptTimeoutMessage(hostTimeoutValue)
 	case -2:
 		return RunScriptDisabledErrMsg
 	default:
@@ -293,6 +298,16 @@ func (hsr HostScriptResult) UserMessage(hostTimeout bool) string {
 	}
 }
 
+func HostScriptTimeoutMessage(seconds *int) string {
+	var timeout int
+	if seconds == nil {
+		timeout = int(scripts.MaxHostExecutionTime.Seconds())
+	} else {
+		timeout = *seconds
+	}
+	return fmt.Sprintf("Timeout. Fleet stopped the script after %d seconds to protect host performance.", timeout)
+}
+
 func (hsr HostScriptResult) HostTimeout(waitForResultTime time.Duration) bool {
 	return hsr.SyncRequest && hsr.ExitCode == nil && time.Now().After(hsr.CreatedAt.Add(waitForResultTime))
 }

server/service/integration_enterprise_test.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/fleetdm/fleet/v4/ee/server/calendar"
 	"github.com/fleetdm/fleet/v4/pkg/optjson"
+	"github.com/fleetdm/fleet/v4/pkg/scripts"
 	"github.com/fleetdm/fleet/v4/server/config"
 	"github.com/fleetdm/fleet/v4/server/contexts/license"
 	"github.com/fleetdm/fleet/v4/server/cron"
@@ -5566,6 +5567,21 @@ func (s *integrationEnterpriseTestSuite) TestRunHostScript() {
 	require.Contains(t, extractServerErrorText(res.Body), fleet.RunScriptDisabledErrMsg)
 	res = s.Do("POST", "/api/latest/fleet/scripts/run/sync", fleet.HostScriptRequestPayload{HostID: plainOsqueryHost.ID, ScriptContents: "echo"}, http.StatusUnprocessableEntity)
 	require.Contains(t, extractServerErrorText(res.Body), fleet.RunScriptDisabledErrMsg)
+
+	// create a execution request that will return a timeout
+	s.DoJSON("POST", "/api/latest/fleet/scripts/run", fleet.HostScriptRequestPayload{HostID: host.ID, ScriptContents: "echo"}, http.StatusAccepted, &runResp)
+
+	// simulate a host response
+	s.DoJSON("POST", "/api/fleet/orbit/scripts/result",
+		json.RawMessage(fmt.Sprintf(`{"orbit_node_key": %q, "execution_id": %q, "exit_code": -1, "output": "script execution error: signal: killed", "timeout": 900}`, *host.OrbitNodeKey, runSyncResp.ExecutionID)),
+		http.StatusOK, &orbitPostScriptResp)
+
+	s.DoJSON("GET", "/api/latest/fleet/scripts/results/"+runSyncResp.ExecutionID, nil, http.StatusOK, &scriptResultResp)
+	require.Equal(t, host.ID, scriptResultResp.HostID)
+	require.Equal(t, "echo", scriptResultResp.ScriptContents)
+	require.Equal(t, int64(-1), *scriptResultResp.ExitCode)
+	require.Equal(t, "Timeout. Fleet stopped the script after 900 seconds to protect host performance.", scriptResultResp.Message)
+	require.Equal(t, "script execution error: signal: killed", scriptResultResp.Output)
 }
 
 func (s *integrationEnterpriseTestSuite) TestRunHostSavedScript() {
@@ -6733,7 +6749,7 @@ VALUES
 				name:       "script-timeout",
 				exitCode:   ptr.Int64(-1),
 				executedAt: now.Add(-1 * time.Hour),
-				expected:   fleet.RunScriptScriptTimeoutErrMsg,
+				expected:   fleet.HostScriptTimeoutMessage(ptr.Int(int(scripts.MaxHostExecutionTime.Seconds()))),
 			},
 			{
 				name:       "pending",

server/service/orbit.go

@@ -324,11 +324,12 @@ func (svc *Service) GetOrbitConfig(ctx context.Context) (fleet.OrbitConfig, erro
 		}
 
 		return fleet.OrbitConfig{
-			Flags:          opts.CommandLineStartUpFlags,
-			Extensions:     extensionsFiltered,
-			Notifications:  notifs,
-			NudgeConfig:    nudgeConfig,
-			UpdateChannels: updateChannels,
+			ScriptExeTimeout: opts.ScriptExecutionTimeout,
+			Flags:            opts.CommandLineStartUpFlags,
+			Extensions:       extensionsFiltered,
+			Notifications:    notifs,
+			NudgeConfig:      nudgeConfig,
+			UpdateChannels:   updateChannels,
 		}, nil
 	}
 
@@ -386,11 +387,12 @@ func (svc *Service) GetOrbitConfig(ctx context.Context) (fleet.OrbitConfig, erro
 	}
 
 	return fleet.OrbitConfig{
-		Flags:          opts.CommandLineStartUpFlags,
-		Extensions:     extensionsFiltered,
-		Notifications:  notifs,
-		NudgeConfig:    nudgeConfig,
-		UpdateChannels: updateChannels,
+		ScriptExeTimeout: opts.ScriptExecutionTimeout,
+		Flags:            opts.CommandLineStartUpFlags,
+		Extensions:       extensionsFiltered,
+		Notifications:    notifs,
+		NudgeConfig:      nudgeConfig,
+		UpdateChannels:   updateChannels,
 	}, nil
 }
 

server/service/scripts.go

@@ -114,7 +114,7 @@ func runScriptSyncEndpoint(ctx context.Context, request interface{}, svc fleet.S
 		// response struct.
 		hostTimeout = true
 	}
-	result.Message = result.UserMessage(hostTimeout)
+	result.Message = result.UserMessage(hostTimeout, result.Timeout)
 	return runScriptSyncResponse{
 		HostScriptResult: result,
 		HostTimeout:      hostTimeout,
@@ -371,7 +371,7 @@ func getScriptResultEndpoint(ctx context.Context, request interface{}, svc fleet
 	// TODO: move this logic out of the endpoint function and consolidate in either the service
 	// method or the fleet package
 	hostTimeout := scriptResult.HostTimeout(scripts.MaxServerWaitTime)
-	scriptResult.Message = scriptResult.UserMessage(hostTimeout)
+	scriptResult.Message = scriptResult.UserMessage(hostTimeout, scriptResult.Timeout)
 
 	return &getScriptResultResponse{
 		ScriptContents: scriptResult.ScriptContents,