Discharge7443/fleet
1
0
Fork 0
Code Issues Pull Requests Actions 19 Packages Projects Releases Wiki Activity

Add script execution simulation to osquery-perf in preparation for load testing (part 3 of ticket) (#13456)

Browse Source
This commit is contained in:
Martin Angers
2023-08-23 18:31:47 -04:00
committed by GitHub
parent d5c7e7eb51
commit edf4a4d02f
4 changed files with 113 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/osquery-perf/agent.go
+75 -3
View File
@@ -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
server/datastore/redis/redis_test.go
+1 -1
View File
@@ -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 {
server/service/integration_mdm_test.go
+12 -5
View File
@@ -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)
server/service/orbit_client.go
+25
View File
@@ -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"