enhance support for orbit and fleet desktop in osquery-perf (#8217)

This improves osquery-perf with support for a more realistic orbit + fleet desktop simulation as described in #8212

This was based on the work done by @sharvilshah in his branch.

cmd/osquery-perf/agent.go

@@ -10,7 +10,6 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"math/rand"
 	"net/http"
@@ -63,18 +62,48 @@ func init() {
 type Stats struct {
 	errors            int
 	enrollments       int
+	orbitenrollments  int
 	distributedwrites int
+	orbitErrors       int
+	desktopErrors     int
 
 	l sync.Mutex
 }
 
-func (s *Stats) RecordStats(errors int, enrollments int, distributedwrites int) {
+func (s *Stats) IncrementErrors(errors int) {
 	s.l.Lock()
 	defer s.l.Unlock()
-
 	s.errors += errors
-	s.enrollments += enrollments
-	s.distributedwrites += distributedwrites
+}
+
+func (s *Stats) IncrementEnrollments() {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.enrollments++
+}
+
+func (s *Stats) IncrementOrbitEnrollments() {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.orbitenrollments++
+}
+
+func (s *Stats) IncrementDistributedWrites() {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.distributedwrites++
+}
+
+func (s *Stats) IncrementOrbitErrors() {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.orbitErrors++
+}
+
+func (s *Stats) IncrementDesktopErrors() {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.desktopErrors++
 }
 
 func (s *Stats) Log() {
@@ -82,11 +111,14 @@ func (s *Stats) Log() {
 	defer s.l.Unlock()
 
 	fmt.Printf(
-		"%s :: error rate: %.2f \t enrollments: %d \t writes: %d\n",
+		"%s :: error rate: %.2f \t enrollments: %d \t orbit enrollments: %d \t writes: %d\n \t orbit errors: %d \t desktop errors: %d",
 		time.Now().String(),
 		float64(s.errors)/float64(s.enrollments),
 		s.enrollments,
+		s.orbitenrollments,
 		s.distributedwrites,
+		s.orbitErrors,
+		s.desktopErrors,
 	)
 }
 
@@ -172,6 +204,7 @@ type agent struct {
 	// Non-nil means the agent is identified as orbit osquery,
 	// nil means the agent is identified as vanilla osquery.
 	deviceAuthToken *string
+	orbitNodeKey    *string
 
 	scheduledQueries []string
 
@@ -242,7 +275,17 @@ type distributedReadResponse struct {
 	Queries map[string]string `json:"queries"`
 }
 
+func (a *agent) isOrbit() bool {
+	return a.deviceAuthToken != nil
+}
+
 func (a *agent) runLoop(i int, onlyAlreadyEnrolled bool) {
+	if a.isOrbit() {
+		if err := a.orbitEnroll(); err != nil {
+			return
+		}
+	}
+
 	if err := a.enroll(i, onlyAlreadyEnrolled); err != nil {
 		return
 	}
@@ -257,6 +300,10 @@ func (a *agent) runLoop(i int, onlyAlreadyEnrolled bool) {
 		}
 	}
 
+	if a.isOrbit() {
+		go a.runOrbitLoop()
+	}
+
 	configTicker := time.Tick(a.ConfigInterval)
 	liveQueryTicker := time.Tick(a.QueryInterval)
 	for {
@@ -274,20 +321,181 @@ func (a *agent) runLoop(i int, onlyAlreadyEnrolled bool) {
 	}
 }
 
+func (a *agent) runOrbitLoop() {
+	orbitClient, err := service.NewOrbitClient(
+		"",
+		a.serverAddress,
+		"",
+		true,
+		a.EnrollSecret,
+		a.UUID,
+	)
+	if err != nil {
+		log.Println("creating orbit client: ", err)
+	}
+
+	orbitClient.TestNodeKey = *a.orbitNodeKey
+
+	deviceClient, err := service.NewDeviceClient(a.serverAddress, true, "")
+	if err != nil {
+		log.Println("creating device client: ", err)
+	}
+
+	// orbit does a config check when it starts
+	if _, err := orbitClient.GetConfig(); err != nil {
+		a.stats.IncrementOrbitErrors()
+		log.Println("orbitClient.GetConfig: ", err)
+	}
+
+	tokenRotationEnabled := orbitClient.GetServerCapabilities().Has(fleet.CapabilityOrbitEndpoints) &&
+		orbitClient.GetServerCapabilities().Has(fleet.CapabilityTokenRotation)
+
+	// it also writes and checks the device token
+	if tokenRotationEnabled {
+		if err := orbitClient.SetOrUpdateDeviceToken(*a.deviceAuthToken); err != nil {
+			a.stats.IncrementOrbitErrors()
+			log.Println("orbitClient.SetOrUpdateDeviceToken: ", err)
+		}
+
+		if err := deviceClient.CheckToken(*a.deviceAuthToken); err != nil {
+			a.stats.IncrementOrbitErrors()
+			log.Println("deviceClient.CheckToken: ", err)
+		}
+	}
+
+	// checkToken is used to simulate Fleet Desktop polling until a token is
+	// valid, we make a random number of requests to properly emulate what
+	// happens in the real world as there are delays that are not accounted by
+	// the way this simulation is arranged.
+	checkToken := func() {
+		min := 1
+		max := 5
+		numberOfRequests := rand.Intn(max-min+1) + min
+		ticker := time.NewTicker(5 * time.Second)
+		defer ticker.Stop()
+
+		for {
+			<-ticker.C
+			numberOfRequests--
+			if err := deviceClient.CheckToken(*a.deviceAuthToken); err != nil {
+				log.Println("deviceClient.CheckToken: ", err)
+			}
+			if numberOfRequests == 0 {
+				break
+			}
+		}
+	}
+
+	// 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
+	// seconds
+	orbitConfigTicker := time.Tick(30 * time.Second)
+	// orbit makes a call every 5 minutes to check the validity of the device
+	// token on the server
+	orbitTokenRemoteCheckTicker := time.Tick(5 * time.Minute)
+	// orbit pings the server every 1 hour to rotate the device token
+	orbitTokenRotationTicker := time.Tick(1 * time.Hour)
+	// orbit polls the /orbit/ping endpoint every 5 minutes to check if the
+	// server capabilities have changed
+	capabilitiesCheckerTicker := time.Tick(5 * time.Minute)
+	// fleet desktop polls for policy compliance every 5 minutes
+	fleetDesktopPolicyTicker := time.Tick(5 * time.Minute)
+
+	for {
+		select {
+		case <-orbitConfigTicker:
+			if _, err := orbitClient.GetConfig(); err != nil {
+				a.stats.IncrementOrbitErrors()
+				log.Println("orbitClient.GetConfig: ", err)
+			}
+		case <-orbitTokenRemoteCheckTicker:
+			if tokenRotationEnabled {
+				if err := deviceClient.CheckToken(*a.deviceAuthToken); err != nil {
+					a.stats.IncrementOrbitErrors()
+					log.Println("deviceClient.CheckToken: ", err)
+				}
+			}
+		case <-orbitTokenRotationTicker:
+			if tokenRotationEnabled {
+				newToken := ptr.String(uuid.NewString())
+				if err := orbitClient.SetOrUpdateDeviceToken(*newToken); err != nil {
+					a.stats.IncrementOrbitErrors()
+					log.Println("orbitClient.SetOrUpdateDeviceToken: ", err)
+				}
+				a.deviceAuthToken = newToken
+				// fleet desktop performs a burst of check token requests after a token is rotated
+				checkToken()
+			}
+		case <-capabilitiesCheckerTicker:
+			if err := orbitClient.Ping(); err != nil {
+				a.stats.IncrementOrbitErrors()
+				log.Println("orbitClient.Ping: ", err)
+			}
+		case <-fleetDesktopPolicyTicker:
+			if _, err := deviceClient.NumberOfFailingPolicies(*a.deviceAuthToken); err != nil {
+				a.stats.IncrementDesktopErrors()
+				log.Println("deviceClient.NumberOfFailingPolicies: ", err)
+			}
+		}
+	}
+}
+
 func (a *agent) waitingDo(req *fasthttp.Request, res *fasthttp.Response) {
 	err := a.fastClient.Do(req, res)
 	for err != nil || res.StatusCode() != http.StatusOK {
 		fmt.Println(err, res.StatusCode())
-		a.stats.RecordStats(1, 0, 0)
+		a.stats.IncrementErrors(1)
 		<-time.Tick(time.Duration(rand.Intn(120)+1) * time.Second)
 		err = a.fastClient.Do(req, res)
 	}
 }
 
+// TODO: add support to `alreadyEnrolled` akin to the `enroll` function.  for
+// now, we assume that the agent is not already enrolled, if you kill the agent
+// process then those Orbit node keys are gone.
+func (a *agent) orbitEnroll() error {
+	params := service.EnrollOrbitRequest{EnrollSecret: a.EnrollSecret, HardwareUUID: a.UUID}
+
+	req := fasthttp.AcquireRequest()
+
+	jsonBytes, err := json.Marshal(params)
+	if err != nil {
+		log.Println("orbit json marshall:", err)
+		return err
+	}
+
+	req.SetBody(jsonBytes)
+	req.Header.SetMethod("POST")
+	req.Header.SetContentType("application/json")
+	req.Header.SetRequestURI(a.serverAddress + "/api/fleet/orbit/enroll")
+
+	resp := fasthttp.AcquireResponse()
+
+	a.waitingDo(req, resp)
+	defer fasthttp.ReleaseResponse(resp)
+
+	if resp.StatusCode() != http.StatusOK {
+		log.Println("orbit enroll status:", resp.StatusCode())
+		return fmt.Errorf("status code: %d", resp.StatusCode())
+	}
+
+	var parsedResp service.EnrollOrbitResponse
+	if err := json.Unmarshal(resp.Body(), &parsedResp); err != nil {
+		log.Println("orbit json parse:", err)
+		return err
+	}
+
+	a.orbitNodeKey = &parsedResp.OrbitNodeKey
+	a.stats.IncrementOrbitEnrollments()
+	return nil
+}
+
 func (a *agent) enroll(i int, onlyAlreadyEnrolled bool) error {
 	a.nodeKey = a.nodeKeyManager.Get(i)
 	if a.nodeKey != "" {
-		a.stats.RecordStats(0, 1, 0)
+		a.stats.IncrementEnrollments()
 		return nil
 	}
 
@@ -326,7 +534,7 @@ func (a *agent) enroll(i int, onlyAlreadyEnrolled bool) error {
 	}
 
 	a.nodeKey = parsedResp.NodeKey
-	a.stats.RecordStats(0, 1, 0)
+	a.stats.IncrementEnrollments()
 
 	a.nodeKeyManager.Add(a.nodeKey)
 
@@ -465,7 +673,7 @@ func loadSoftware(platform string, ver string) []map[string]string {
 		fmt.Sprintf("%s_%s-software.json.bz2", platform, ver),
 	)
 
-	tmpDir, err := ioutil.TempDir("", "osquery-perf")
+	tmpDir, err := os.MkdirTemp("", "osquery-perf")
 	if err != nil {
 		panic(err)
 	}
@@ -482,7 +690,7 @@ func loadSoftware(platform string, ver string) []map[string]string {
 	}
 
 	var software []softwareJSON
-	contents, err := ioutil.ReadFile(dstPath)
+	contents, err := os.ReadFile(dstPath)
 	if err != nil {
 		log.Printf("reading vuln software for %s %s: %s\n", platform, ver, err)
 		return nil
@@ -916,7 +1124,7 @@ func (a *agent) DistributedWrite(queries map[string]string) {
 	fasthttp.ReleaseRequest(req)
 	defer fasthttp.ReleaseResponse(res)
 
-	a.stats.RecordStats(0, 0, 1)
+	a.stats.IncrementDistributedWrites()
 	// No need to read the distributed write body
 }
 

cmd/osquery-perf/mac10.14.6.tmpl

@@ -53,7 +53,7 @@
             "uuid": "{{ .UUID }}"
         }
     },
-    "host_identifier": "{{ .CachedString "hostname" }}",
+    "host_identifier": "{{ .UUID }}",
     "platform_type": "16"
 }
 {{- end }}

infrastructure/loadtesting/terraform/loadtesting.tf

@@ -51,9 +51,8 @@ resource "aws_ecs_task_definition" "loadtest" {
         command = [
           "/go/osquery-perf",
           "-enroll_secret", data.aws_secretsmanager_secret_version.enroll_secret.secret_string,
-          "-host_count", "5000",
+          "-host_count", "500",
           "-server_url", "http://${aws_lb.internal.dns_name}",
-          "-node_key_file", "nodekeys",
           "--policy_pass_prob", "0.5",
           "--start_period", "5m",
         ]

server/service/base_client.go

@@ -15,6 +15,8 @@ import (
 	"github.com/fleetdm/fleet/v4/server/fleet"
 )
 
+var errInvalidScheme = errors.New("address must start with https:// for remote connections")
+
 // httpClient interface allows the HTTP methods to be mocked.
 type httpClient interface {
 	Do(req *http.Request) (*http.Response, error)
@@ -113,8 +115,9 @@ func newBaseClient(addr string, insecureSkipVerify bool, rootCA, urlPrefix strin
 		return nil, fmt.Errorf("parsing URL: %w", err)
 	}
 
-	if baseURL.Scheme != "https" && !strings.Contains(baseURL.Host, "localhost") && !strings.Contains(baseURL.Host, "127.0.0.1") {
-		return nil, errors.New("address must start with https:// for remote connections")
+	allowHTTP := insecureSkipVerify || strings.Contains(baseURL.Host, "localhost") || strings.Contains(baseURL.Host, "127.0.0.1")
+	if baseURL.Scheme != "https" && !allowHTTP {
+		return nil, errInvalidScheme
 	}
 
 	rootCAPool := x509.NewCertPool()

server/service/base_client_test.go

@@ -92,19 +92,30 @@ func TestParseResponseGeneralErrors(t *testing.T) {
 
 func TestNewBaseClient(t *testing.T) {
 	t.Run("invalid addresses are an error", func(t *testing.T) {
-		_, err := newBaseClient("invalid", true, "", "", fleet.CapabilityMap{})
+		_, err := newBaseClient("http://foo\x7f.com/", true, "", "", fleet.CapabilityMap{})
 		require.Error(t, err)
 	})
 
 	t.Run("http is only valid in development", func(t *testing.T) {
-		_, err := newBaseClient("http://test.com", true, "", "", fleet.CapabilityMap{})
-		require.Error(t, err)
+		cases := []struct {
+			name               string
+			address            string
+			insecureSkipVerify bool
+			expectedErr        error
+		}{
+			{"http non-local URL without insecureSkipVerify", "http://test.com", false, errInvalidScheme},
+			{"http non-local URL with insecureSkipVerify", "http://test.com", true, nil},
+			{"https", "https://test.com", false, nil},
+			{"http localhost with insecureSkipVerify", "http://localhost:8080", true, nil},
+			{"http localhost without insecureSkipVerify", "http://localhost:8080", false, nil},
+			{"http local ip with insecureSkipVerify", "http://127.0.0.1:8080", true, nil},
+			{"http local ip without insecureSkipVerify", "http://127.0.0.1:8080", false, nil},
+		}
 
-		_, err = newBaseClient("http://localhost:8080", true, "", "", fleet.CapabilityMap{})
-		require.NoError(t, err)
-
-		_, err = newBaseClient("http://127.0.0.1:8080", true, "", "", fleet.CapabilityMap{})
-		require.NoError(t, err)
+		for _, c := range cases {
+			_, err := newBaseClient(c.address, c.insecureSkipVerify, "", "", fleet.CapabilityMap{})
+			require.Equal(t, c.expectedErr, err, c.name)
+		}
 	})
 }
 

server/service/handler.go

@@ -485,7 +485,7 @@ func attachFleetAPIRoutes(r *mux.Router, svc fleet.Service, config config.FleetC
 		ne.HEAD(apple_mdm.InstallerPath, mdmAppleHeadInstallerEndpoint, mdmAppleHeadInstallerRequest{})
 	}
 
-	ne.POST("/api/fleet/orbit/enroll", enrollOrbitEndpoint, enrollOrbitRequest{})
+	ne.POST("/api/fleet/orbit/enroll", enrollOrbitEndpoint, EnrollOrbitRequest{})
 
 	// For some reason osquery does not provide a node key with the block data.
 	// Instead the carve session ID should be verified in the service method.

server/service/orbit.go

@@ -21,12 +21,12 @@ type orbitError struct {
 	message string
 }
 
-type enrollOrbitRequest struct {
+type EnrollOrbitRequest struct {
 	EnrollSecret string `json:"enroll_secret"`
 	HardwareUUID string `json:"hardware_uuid"`
 }
 
-type enrollOrbitResponse struct {
+type EnrollOrbitResponse struct {
 	OrbitNodeKey string `json:"orbit_node_key,omitempty"`
 	Err          error  `json:"error,omitempty"`
 }
@@ -52,12 +52,12 @@ func (e orbitError) Error() string {
 	return e.message
 }
 
-func (r enrollOrbitResponse) error() error { return r.Err }
+func (r EnrollOrbitResponse) error() error { return r.Err }
 
 // hijackRender so we can add a header with the server capabilities in the
 // response, allowing Orbit to know what features are available without the
 // need to enroll.
-func (r enrollOrbitResponse) hijackRender(ctx context.Context, w http.ResponseWriter) {
+func (r EnrollOrbitResponse) hijackRender(ctx context.Context, w http.ResponseWriter) {
 	writeCapabilitiesHeader(w, fleet.ServerOrbitCapabilities)
 	enc := json.NewEncoder(w)
 	enc.SetIndent("", "  ")
@@ -68,12 +68,12 @@ func (r enrollOrbitResponse) hijackRender(ctx context.Context, w http.ResponseWr
 }
 
 func enrollOrbitEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (interface{}, error) {
-	req := request.(*enrollOrbitRequest)
+	req := request.(*EnrollOrbitRequest)
 	nodeKey, err := svc.EnrollOrbit(ctx, req.HardwareUUID, req.EnrollSecret)
 	if err != nil {
-		return enrollOrbitResponse{Err: err}, nil
+		return EnrollOrbitResponse{Err: err}, nil
 	}
-	return enrollOrbitResponse{OrbitNodeKey: nodeKey}, nil
+	return EnrollOrbitResponse{OrbitNodeKey: nodeKey}, nil
 }
 
 func (svc *Service) AuthenticateOrbitHost(ctx context.Context, orbitNodeKey string) (*fleet.Host, bool, error) {

server/service/orbit_client.go

@@ -31,6 +31,9 @@ type OrbitClient struct {
 
 	lastRecordedErrMu sync.Mutex
 	lastRecordedErr   error
+
+	// TestNodeKey is used for testing only.
+	TestNodeKey string
 }
 
 func (oc *OrbitClient) request(verb string, path string, params interface{}, resp interface{}) error {
@@ -133,8 +136,8 @@ func (oc *OrbitClient) Ping() error {
 
 func (oc *OrbitClient) enroll() (string, error) {
 	verb, path := "POST", "/api/fleet/orbit/enroll"
-	params := enrollOrbitRequest{EnrollSecret: oc.enrollSecret, HardwareUUID: oc.uuid}
-	var resp enrollOrbitResponse
+	params := EnrollOrbitRequest{EnrollSecret: oc.enrollSecret, HardwareUUID: oc.uuid}
+	var resp EnrollOrbitResponse
 	err := oc.request(verb, path, params, &resp)
 	if err != nil {
 		return "", err
@@ -149,6 +152,10 @@ var enrollLock sync.Mutex
 // getNodeKeyOrEnroll attempts to read the orbit node key if the file exists on disk
 // otherwise it enrolls the host with Fleet and saves the node key to disk
 func (oc *OrbitClient) getNodeKeyOrEnroll() (string, error) {
+	if oc.TestNodeKey != "" {
+		return oc.TestNodeKey, nil
+	}
+
 	enrollLock.Lock()
 	defer enrollLock.Unlock()