wip: add back-off mechanism when validating Fleet Desktop token

Author: roperzh

orbit/cmd/desktop/desktop.go

@@ -168,13 +168,15 @@ func main() {
 			return newToken
 		})
 
-		refetchToken := func() {
+		refetchTokenValueFromDisk := func() {
 			if _, err := tokenReader.Read(); err != nil {
 				log.Error().Err(err).Msg("refetch token")
 			}
 			log.Debug().Msg("successfully refetched the token from disk")
 		}
 
+		tokenChecker := token.NewChecker(identifierPath, client)
+
 		disableTray := func() {
 			log.Debug().Msg("disabling tray items")
 			myDeviceItem.SetTitle("Connecting...")
@@ -184,46 +186,24 @@ func main() {
 			migrateMDMItem.Hide()
 		}
 
-		// checkToken performs API test calls to enable the "My device" item as
-		// soon as the device auth token is registered by Fleet.
-		checkToken := func() <-chan interface{} {
-			done := make(chan interface{})
-
-			go func() {
-				ticker := time.NewTicker(5 * time.Second)
-				defer ticker.Stop()
-				defer close(done)
-
-				for {
-					refetchToken()
-					_, err := client.DesktopSummary(tokenReader.GetCached())
-
-					if err == nil || errors.Is(err, service.ErrMissingLicense) {
-						log.Debug().Msg("enabling tray items")
-						myDeviceItem.SetTitle("My device")
-						myDeviceItem.Enable()
-						transparencyItem.Enable()
-						return
-					}
-
-					log.Error().Err(err).Msg("get device URL")
-
-					<-ticker.C
-				}
-			}()
-
-			return done
+		// checkToken perform API calls until a token is valid. When we
+		// find a valid token it refreshes the cached alue and enables
+		// basic tray items.
+		checkToken := func() {
+			tokenChecker.AwaitValid()
+			refetchTokenValueFromDisk()
+			log.Debug().Msg("enabling tray items")
+			myDeviceItem.SetTitle("My device")
+			myDeviceItem.Enable()
+			transparencyItem.Enable()
 		}
 
-		// start a check as soon as the app starts
-		deviceEnabledChan := checkToken()
-
 		// this loop checks the `mtime` value of the token file and:
 		// 1. if the token file was modified, it disables the tray items until we
 		// verify the token is valid
 		// 2. calls (blocking) `checkToken` to verify the token is valid
 		go func() {
-			<-deviceEnabledChan
+			tokenChecker.AwaitValid()
 			tic := time.NewTicker(1 * time.Second)
 			defer tic.Stop()
 
@@ -234,9 +214,8 @@ func main() {
 				case err != nil:
 					log.Error().Err(err).Msg("check token file")
 				case expired:
-					log.Info().Msg("token file changed, rechecking")
 					disableTray()
-					<-checkToken()
+					checkToken()
 				}
 			}
 		}()
@@ -279,7 +258,6 @@ func main() {
 		// poll the server to check the policy status of the host and update the
 		// tray icon accordingly
 		go func() {
-			<-deviceEnabledChan
 			tic := time.NewTicker(5 * time.Minute)
 			defer tic.Stop()
 
@@ -290,11 +268,13 @@ func main() {
 				case err == nil:
 					// OK
 				case errors.Is(err, service.ErrMissingLicense):
+					// This case is for devices using the free plan.
 					myDeviceItem.SetTitle("My device")
+					transparencyItem.Enable()
 					continue
 				case errors.Is(err, service.ErrUnauthenticated):
 					disableTray()
-					<-checkToken()
+					checkToken()
 					continue
 				default:
 					log.Error().Err(err).Msg("get failing policies")
@@ -326,6 +306,7 @@ func main() {
 					}
 				}
 				myDeviceItem.Enable()
+				transparencyItem.Enable()
 
 				shouldRunMigrator := sum.Notifications.NeedsMDMMigration || sum.Notifications.RenewEnrollmentProfile
 

orbit/pkg/token/checker.go

@@ -0,0 +1,69 @@
+package token
+
+import (
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/cenkalti/backoff"
+	"github.com/fleetdm/fleet/v4/server/service"
+)
+
+type client interface {
+	CheckToken(string) error
+}
+
+type reader interface {
+	Read() (string, error)
+	GetCached() string
+}
+
+// TODO: docs
+type RemoteChecker struct {
+	reader reader
+	client client
+}
+
+// TODO: docs
+func NewChecker(path string, client client) *RemoteChecker {
+	return &RemoteChecker{
+		reader: &Reader{Path: path},
+		client: client,
+	}
+}
+
+func (c *RemoteChecker) isValid(err error) bool {
+	return err == nil || errors.Is(err, service.ErrMissingLicense)
+}
+
+// TODO: docs
+func (c *RemoteChecker) AwaitValid() {
+	// TODO: find appropriate values
+	// randomized interval = RetryInterval * (random value in range [1 - RandomizationFactor, 1 + RandomizationFactor])
+	retryStrategy := backoff.NewExponentialBackOff()
+	retryStrategy.InitialInterval = 2 * time.Second
+	retryStrategy.MaxElapsedTime = 1 * time.Minute
+
+	for {
+		if err := backoff.Retry(
+			func() error {
+				if _, err := c.reader.Read(); err != nil {
+					return err
+				}
+				err := c.client.CheckToken(c.reader.GetCached())
+				if !c.isValid(err) {
+					// TODO: better error
+					return errors.New("invalid")
+				}
+				return nil
+			},
+			retryStrategy,
+		); err != nil {
+			// TODO: what do we do here?
+			fmt.Println("backoff gave up")
+			continue
+		}
+
+		return
+	}
+}

orbit/pkg/token/checker_test.go

@@ -0,0 +1,146 @@
+package token
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/fleetdm/fleet/v4/server/service"
+	"github.com/stretchr/testify/require"
+)
+
+type mockClient struct {
+	TokenValidationFunc func(string) error
+}
+
+func (s *mockClient) CheckToken(token string) error {
+	return s.TokenValidationFunc(token)
+}
+
+type mockReader struct {
+	ReadFunc  func() (string, error)
+	CachedVal string
+}
+
+func (s *mockReader) Read() (string, error) {
+	return s.ReadFunc()
+}
+
+func (s *mockReader) GetCached() string {
+	return s.CachedVal
+}
+
+func TestNewChecker(t *testing.T) {
+	client := &mockClient{}
+	checker := NewChecker("path/to/token", client)
+
+	require.NotNil(t, checker)
+}
+
+func TestIsValid(t *testing.T) {
+	client := &mockClient{}
+	checker := NewChecker("path/to/token", client)
+
+	testCases := []struct {
+		err      error
+		expected bool
+	}{
+		{nil, true},
+		{errors.New("random error"), false},
+		{service.ErrMissingLicense, true},
+	}
+
+	for _, tc := range testCases {
+		result := checker.isValid(tc.err)
+		require.Equal(t, tc.expected, result)
+	}
+}
+
+func TestAwaitValid_Success(t *testing.T) {
+	client := &mockClient{
+		TokenValidationFunc: func(token string) error {
+			return nil
+		},
+	}
+	reader := &mockReader{
+		ReadFunc: func() (string, error) {
+			return "valid token", nil
+		},
+	}
+
+	checker := &RemoteChecker{
+		reader: reader,
+		client: client,
+	}
+
+	done := make(chan bool)
+	go func() {
+		checker.AwaitValid()
+		done <- true
+	}()
+
+	select {
+	case <-done:
+		// test passed
+	case <-time.After(2 * time.Second):
+		t.Fatal("Test timed out - AwaitValid did not complete in expected time")
+	}
+}
+
+func TestAwaitValid_Failure(t *testing.T) {
+	client := &mockClient{
+		TokenValidationFunc: func(token string) error {
+			return errors.New("invalid token")
+		},
+	}
+	reader := &mockReader{
+		ReadFunc: func() (string, error) {
+			return "invalid token", nil
+		},
+	}
+
+	checker := &RemoteChecker{
+		reader: reader,
+		client: client,
+	}
+
+	done := make(chan bool)
+	go func() {
+		checker.AwaitValid()
+		done <- true
+	}()
+
+	select {
+	case <-done:
+		t.Fatal("Test failed - AwaitValid should not have completed")
+	case <-time.After(2 * time.Second):
+		// test passed
+	}
+}
+
+func TestAwaitValid_ReaderError(t *testing.T) {
+	client := &mockClient{}
+	reader := &mockReader{
+		ReadFunc: func() (string, error) {
+			return "", errors.New("reader error")
+		},
+	}
+
+	checker := &RemoteChecker{
+		reader: reader,
+		client: client,
+	}
+
+	done := make(chan bool)
+	go func() {
+		checker.AwaitValid()
+		done <- true
+	}()
+
+	select {
+	case <-done:
+		t.Fatal("Test failed - AwaitValid should not have completed")
+	case <-time.After(2 * time.Second):
+		// test passed
+	}
+}