Improve Fleet Desktop "My Device" menu item UX at install time (#5915)

* Improve Fleet Desktop My Device link availability

* Use svc.clock and add test

* Revert change and add check for LastEnrolledAt on tests

changes/issue-4733-improve-fleet-desktop-my-device-url

@@ -0,0 +1 @@
+* Fixed Fleet Desktop's "Initializing..." menu item taking up to 1h (on install time) to change to "My device".

server/datastore/mysql/hosts_test.go

@@ -843,6 +843,7 @@ func testHostsEnroll(t *testing.T, ds *Datastore) {
 	for _, tt := range enrollTests {
 		h, err := ds.EnrollHost(context.Background(), tt.uuid, tt.nodeKey, &team.ID, 0)
 		require.NoError(t, err)
+		assert.NotZero(t, h.LastEnrolledAt)
 
 		assert.Equal(t, tt.uuid, h.OsqueryHostID)
 		assert.Equal(t, tt.nodeKey, h.NodeKey)
@@ -850,10 +851,12 @@ func testHostsEnroll(t *testing.T, ds *Datastore) {
 		// This host should be allowed to re-enroll immediately if cooldown is disabled
 		_, err = ds.EnrollHost(context.Background(), tt.uuid, tt.nodeKey+"new", nil, 0)
 		require.NoError(t, err)
+		assert.NotZero(t, h.LastEnrolledAt)
 
 		// This host should not be allowed to re-enroll immediately if cooldown is enabled
 		_, err = ds.EnrollHost(context.Background(), tt.uuid, tt.nodeKey+"new", nil, 10*time.Second)
 		require.Error(t, err)
+		assert.NotZero(t, h.LastEnrolledAt)
 	}
 
 	hosts, err = ds.ListHosts(context.Background(), filter, fleet.HostListOptions{})

server/service/osquery.go

@@ -490,6 +490,9 @@ func getDistributedQueriesEndpoint(ctx context.Context, request interface{}, svc
 	}, nil
 }
 
+// orbitInfoRefetchAfterEnrollDur value assumes the default distributed_interval value set by Fleet of 10s.
+const orbitInfoRefetchAfterEnrollDur = 1 * time.Minute
+
 func (svc *Service) GetDistributedQueries(ctx context.Context) (queries map[string]string, discovery map[string]string, accelerate uint, err error) {
 	// skipauth: Authorization is currently for user endpoints only.
 	svc.authz.SkipAuthorization(ctx)
@@ -513,6 +516,20 @@ func (svc *Service) GetDistributedQueries(ctx context.Context) (queries map[stri
 		discovery[name] = query
 	}
 
+	// The following is added to improve Fleet Desktop's UX at install time.
+	//
+	// At install (enroll) time, the "orbit_info" extension takes longer to load than the first
+	// query check-in (distributed/read request).
+	// To avoid having to wait for the next check-in to ingest the data (after
+	// svc.config.Osquery.DetailUpdateInterval, 1h by default),
+	// we make the best effort to retrieve such "device auth token" from the device, but with a
+	// limit of orbitInfoRefetchAfterEnrollDur to not generate too much write database overhead
+	// (writes to `host_device_auth` table).
+	if svc.clock.Now().Sub(host.LastEnrolledAt) < orbitInfoRefetchAfterEnrollDur {
+		queries[hostDetailQueryPrefix+osquery_utils.OrbitInfoQueryName] = osquery_utils.OrbitInfoDetailQuery.Query
+		discovery[hostDetailQueryPrefix+osquery_utils.OrbitInfoQueryName] = osquery_utils.OrbitInfoDetailQuery.Discovery
+	}
+
 	labelQueries, err := svc.labelQueriesForHost(ctx, host)
 	if err != nil {
 		return nil, nil, 0, osqueryError{message: err.Error()}

server/service/osquery_test.go

@@ -2569,3 +2569,49 @@ func TestLiveQueriesFailing(t *testing.T) {
 	require.Contains(t, string(logs), "level=error")
 	require.Contains(t, string(logs), "failed to get queries for host")
 }
+
+// TestFleetDesktopOrbitInfo tests that the orbit_info table extension is
+// refetched for "orbitInfoRefetchAfterEnrollDur" after enroll.
+func TestFleetDesktopOrbitInfo(t *testing.T) {
+	ds := new(mock.Store)
+	lq := new(live_query.MockLiveQuery)
+	mockClock := clock.NewMockClock()
+	fleetConfig := config.TestConfig()
+	fleetConfig.Osquery.LabelUpdateInterval = 5 * time.Minute
+	fleetConfig.Osquery.PolicyUpdateInterval = 5 * time.Minute
+	fleetConfig.Osquery.DetailUpdateInterval = 5 * time.Minute
+	svc := newTestServiceWithConfig(t, ds, fleetConfig, nil, lq, &TestServerOpts{Clock: mockClock})
+
+	lq.On("QueriesForHost", uint(0)).Return(map[string]string{}, nil)
+
+	now := time.Now().UTC()
+	mockClock.SetTime(now)
+
+	host := &fleet.Host{
+		Platform: "darwin",
+		// Host has enrolled 30 seconds ago.
+		LastEnrolledAt: now.Add(-30 * time.Second),
+		// Host is up-to-date with details, labels and policies
+		// because update interval for each is 5m.
+		DetailUpdatedAt: now.Add(-5 * time.Second),
+		LabelUpdatedAt:  now.Add(-5 * time.Second),
+		PolicyUpdatedAt: now.Add(-5 * time.Second),
+	}
+
+	ctx := hostctx.NewContext(context.Background(), host)
+
+	queries, discovery, _, err := svc.GetDistributedQueries(ctx)
+	require.NoError(t, err)
+	require.Len(t, queries, 1)
+	verifyDiscovery(t, queries, discovery)
+	require.Contains(t, queries, "fleet_detail_query_orbit_info")
+
+	// Advance mock clock
+	mockClock.AddTime(orbitInfoRefetchAfterEnrollDur)
+	ctx = hostctx.NewContext(context.Background(), host)
+
+	queries, discovery, _, err = svc.GetDistributedQueries(ctx)
+	require.NoError(t, err)
+	require.Len(t, queries, 0)
+	require.Len(t, discovery, 0)
+}

server/service/osquery_utils/queries.go

@@ -320,11 +320,17 @@ FROM logical_drives WHERE file_system = 'NTFS' LIMIT 1;`,
 		DirectIngestFunc: directIngestChromeProfiles,
 		Discovery:        discoveryTable("google_chrome_profiles"),
 	},
-	"orbit_info": {
-		Query:            `SELECT * FROM orbit_info`,
-		DirectIngestFunc: directIngestOrbitInfo,
-		Discovery:        discoveryTable("orbit_info"),
-	},
+	OrbitInfoQueryName: OrbitInfoDetailQuery,
+}
+
+// OrbitInfoQueryName is the name of the query to ingest orbit_info table extension data.
+const OrbitInfoQueryName = "orbit_info"
+
+// OrbitInfoDetailQuery holds the query and ingestion function for the orbit_info table extension.
+var OrbitInfoDetailQuery = DetailQuery{
+	Query:            `SELECT * FROM orbit_info`,
+	DirectIngestFunc: directIngestOrbitInfo,
+	Discovery:        discoveryTable("orbit_info"),
 }
 
 // discoveryTable returns a query to determine whether a table exists or not.