From edf4a4d02f3d71c0ba86d1262ddadb53883fef92 Mon Sep 17 00:00:00 2001 From: Martin Angers Date: Wed, 23 Aug 2023 18:31:47 -0400 Subject: [PATCH] Add script execution simulation to osquery-perf in preparation for load testing (part 3 of ticket) (#13456) --- cmd/osquery-perf/agent.go | 78 +++++++++++++++++++++++++- server/datastore/redis/redis_test.go | 2 +- server/service/integration_mdm_test.go | 17 ++++-- server/service/orbit_client.go | 25 +++++++++ 4 files changed, 113 insertions(+), 9 deletions(-) diff --git a/cmd/osquery-perf/agent.go b/cmd/osquery-perf/agent.go index 32504c8680..3e0dc60094 100644 --- a/cmd/osquery-perf/agent.go +++ b/cmd/osquery-perf/agent.go @@ -3,12 +3,15 @@ package main import ( "bytes" "compress/bzip2" + cryptorand "crypto/rand" "crypto/tls" "embed" + "encoding/base64" "encoding/json" "errors" "flag" "fmt" + "io" "log" "math/rand" "net/http" @@ -16,6 +19,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "text/template" "time" @@ -283,6 +287,11 @@ type agent struct { // isEnrolledToMDMMu protects isEnrolledToMDM. isEnrolledToMDMMu sync.Mutex + disableScriptExec bool + // atomic boolean is set to true when executing scripts, so that only a + // single goroutine at a time can execute scripts. + scriptExecRunning atomic.Bool + // // The following are exported to be used by the templates. // @@ -327,6 +336,7 @@ func newAgent( mdmSCEPChallenge string, liveQueryFailProb float64, liveQueryNoResultsProb float64, + disableScriptExec bool, ) *agent { var deviceAuthToken *string if rand.Float64() <= orbitProb { @@ -377,7 +387,8 @@ func newAgent( UUID: uuid, SerialNumber: serialNumber, - mdmClient: mdmClient, + mdmClient: mdmClient, + disableScriptExec: disableScriptExec, } } @@ -519,7 +530,7 @@ func (a *agent) runOrbitLoop() { // fleet desktop performs a burst of check token requests when it's initialized checkToken() - // orbit makes a call to check the config and update the CLI flags every 5 + // orbit makes a call to check the config and update the CLI flags every 30 // seconds orbitConfigTicker := time.Tick(30 * time.Second) // orbit makes a call every 5 minutes to check the validity of the device @@ -536,10 +547,16 @@ func (a *agent) runOrbitLoop() { for { select { case <-orbitConfigTicker: - if _, err := orbitClient.GetConfig(); err != nil { + cfg, err := orbitClient.GetConfig() + if err != nil { a.stats.IncrementOrbitErrors() log.Println("orbitClient.GetConfig: ", err) } + if len(cfg.Notifications.PendingScriptExecutionIDs) > 0 { + // there are pending scripts to execute on this host, start a goroutine + // that will simulate executing them. + go a.execScripts(cfg.Notifications.PendingScriptExecutionIDs, orbitClient) + } case <-orbitTokenRemoteCheckTicker: if tokenRotationEnabled { if err := deviceClient.CheckToken(*a.deviceAuthToken); err != nil { @@ -595,6 +612,58 @@ func (a *agent) runMDMLoop() { } } +func (a *agent) execScripts(execIDs []string, orbitClient *service.OrbitClient) { + if a.scriptExecRunning.Swap(true) { + // if Swap returns true, the goroutine was already running, exit + return + } + defer a.scriptExecRunning.Store(false) + + log.Printf("running scripts: %v\n", execIDs) + for _, execID := range execIDs { + if a.disableScriptExec { + // send a no-op result without executing if script exec is disabled + if err := orbitClient.SaveHostScriptResult(&fleet.HostScriptResultPayload{ + ExecutionID: execID, + Output: "script execution is disabled", + Runtime: 0, + ExitCode: -1, + }); err != nil { + log.Println("save disabled host script result:", err) + return + } + log.Printf("did save disabled host script result: id=%s\n", execID) + continue + } + + script, err := orbitClient.GetHostScript(execID) + if err != nil { + log.Println("get host script:", err) + return + } + + // simulate script execution + outputLen := rand.Intn(11000) // base64 encoding will make the actual output a bit bigger + buf := make([]byte, outputLen) + n, _ := io.ReadFull(cryptorand.Reader, buf) + exitCode := rand.Intn(2) + runtime := rand.Intn(5) + time.Sleep(time.Duration(runtime) * time.Second) + + if err := orbitClient.SaveHostScriptResult(&fleet.HostScriptResultPayload{ + HostID: script.HostID, + ExecutionID: script.ExecutionID, + Output: base64.StdEncoding.EncodeToString(buf[:n]), + Runtime: runtime, + ExitCode: exitCode, + }); err != nil { + log.Println("save host script result:", err) + return + } + log.Printf("did exec and save host script result: id=%s, output size=%d, runtime=%d, exit code=%d\n", execID, base64.StdEncoding.EncodedLen(n), runtime, exitCode) + } +} + func (a *agent) waitingDo(req *fasthttp.Request, res *fasthttp.Response) { err := a.fastClient.Do(req, res) for err != nil || res.StatusCode() != http.StatusOK { @@ -1345,6 +1414,8 @@ func main() { liveQueryFailProb = flag.Float64("live_query_fail_prob", 0.0, "Probability of a live query failing execution in the host") liveQueryNoResultsProb = flag.Float64("live_query_no_results_prob", 0.2, "Probability of a live query returning no results") + + disableScriptExec = flag.Bool("disable_script_exec", false, "Disable script execution support") ) flag.Parse() @@ -1426,6 +1497,7 @@ func main() { *mdmSCEPChallenge, *liveQueryFailProb, *liveQueryNoResultsProb, + *disableScriptExec, ) a.stats = stats a.nodeKeyManager = nodeKeyManager diff --git a/server/datastore/redis/redis_test.go b/server/datastore/redis/redis_test.go index b6ee30e3ed..c7e2da05d7 100644 --- a/server/datastore/redis/redis_test.go +++ b/server/datastore/redis/redis_test.go @@ -78,7 +78,7 @@ func TestConnectRetry(t *testing.T) { &netError{error: io.EOF, allowedCalls: 10}, 2, 1, 0, 100 * time.Millisecond, }, // net error, but non-retryable { - &netError{error: io.EOF, timeout: true, allowedCalls: 1}, 10, 2, 250 * time.Millisecond, 750 * time.Millisecond, + &netError{error: io.EOF, timeout: true, allowedCalls: 1}, 10, 2, 250 * time.Millisecond, 800 * time.Millisecond, }, // retryable, but succeeded after one retry } for _, c := range cases { diff --git a/server/service/integration_mdm_test.go b/server/service/integration_mdm_test.go index 88189df10e..0ea8fa4ff3 100644 --- a/server/service/integration_mdm_test.go +++ b/server/service/integration_mdm_test.go @@ -1498,15 +1498,22 @@ func (s *integrationMDMTestSuite) TestDEPProfileAssignment() { // TODO: seems like we're doing this request on each loop? require.Len(t, profileAssignmentReqs[0].Devices, 1) require.Equal(t, devices[0].SerialNumber, profileAssignmentReqs[0].Devices[0]) + + // profileAssignmentReqs[1] and [2] can be in any order + ix2Devices, ix1Device := 1, 2 + if len(profileAssignmentReqs[1].Devices) == 1 { + ix2Devices, ix1Device = ix1Device, ix2Devices + } + // - existing device with "added" // - new device with "added" - require.Len(t, profileAssignmentReqs[1].Devices, 2) - require.Equal(t, devices[0].SerialNumber, profileAssignmentReqs[1].Devices[0]) - require.Equal(t, addedSerial, profileAssignmentReqs[1].Devices[1]) + require.Len(t, profileAssignmentReqs[ix2Devices].Devices, 2, "%#+v", profileAssignmentReqs) + require.Equal(t, devices[0].SerialNumber, profileAssignmentReqs[ix2Devices].Devices[0]) + require.Equal(t, addedSerial, profileAssignmentReqs[ix2Devices].Devices[1]) // - existing device with "modified" and a different team (thus different profile request) - require.Len(t, profileAssignmentReqs[2].Devices, 1) - require.Equal(t, devices[1].SerialNumber, profileAssignmentReqs[2].Devices[0]) + require.Len(t, profileAssignmentReqs[ix1Device].Devices, 1) + require.Equal(t, devices[1].SerialNumber, profileAssignmentReqs[ix1Device].Devices[0]) // entries for all hosts except for the one with OpType = "deleted" assignment, err := s.ds.GetHostDEPAssignment(ctx, deletedHostID) diff --git a/server/service/orbit_client.go b/server/service/orbit_client.go index 057911c605..c5c094bba0 100644 --- a/server/service/orbit_client.go +++ b/server/service/orbit_client.go @@ -131,6 +131,31 @@ func (oc *OrbitClient) SetOrUpdateDeviceToken(deviceAuthToken string) error { return nil } +// GetHostScript returns the script fetched from Fleet server to run on this +// host. +func (oc *OrbitClient) GetHostScript(execID string) (*fleet.HostScriptResult, error) { + verb, path := "POST", "/api/fleet/orbit/scripts/request" + var resp orbitGetScriptResponse + if err := oc.authenticatedRequest(verb, path, &orbitGetScriptRequest{ + ExecutionID: execID, + }, &resp); err != nil { + return nil, err + } + return resp.HostScriptResult, nil +} + +// SaveHostScriptResult saves the result of running the script on this host. +func (oc *OrbitClient) SaveHostScriptResult(result *fleet.HostScriptResultPayload) error { + verb, path := "POST", "/api/fleet/orbit/scripts/result" + var resp orbitPostScriptResultResponse + if err := oc.authenticatedRequest(verb, path, &orbitPostScriptResultRequest{ + HostScriptResultPayload: result, + }, &resp); err != nil { + return err + } + return nil +} + // Ping sends a ping request to the orbit/ping endpoint. func (oc *OrbitClient) Ping() error { verb, path := "HEAD", "/api/fleet/orbit/ping"