Skip to content

Commit

Permalink
added edge case coverage and sped up test
Browse files Browse the repository at this point in the history
  • Loading branch information
@cesarfda
cesarfda committed Jan 14, 2025
1 parent f0a3c8a commit 25e9a93
Show file tree
Hide file tree
Showing 2 changed files with 123 additions and 37 deletions.
57 changes: 34 additions & 23 deletions ee/tuf/autoupdate.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,22 +86,24 @@ type querier interface {
}

type TufAutoupdater struct {
metadataClient *client.Client
libraryManager librarian
osquerier querier // used to query for current running osquery version
osquerierRetryInterval time.Duration
knapsack types.Knapsack
store types.KVStore // stores autoupdater errors for kolide_tuf_autoupdater_errors table
updateChannel string
pinnedVersions map[autoupdatableBinary]string // maps the binaries to their pinned versions
pinnedVersionGetters map[autoupdatableBinary]func() string // maps the binaries to the knapsack function to retrieve updated pinned versions
initialDelayEnd time.Time
updateLock *sync.Mutex
interrupt chan struct{}
interrupted atomic.Bool
signalRestart chan error
slogger *slog.Logger
restartFuncs map[autoupdatableBinary]func() error
metadataClient *client.Client
libraryManager librarian
osquerier querier // used to query for current running osquery version
osquerierRetryInterval time.Duration
knapsack types.Knapsack
store types.KVStore // stores autoupdater errors for kolide_tuf_autoupdater_errors table
updateChannel string
pinnedVersions map[autoupdatableBinary]string // maps the binaries to their pinned versions
pinnedVersionGetters map[autoupdatableBinary]func() string // maps the binaries to the knapsack function to retrieve updated pinned versions
initialDelayEnd time.Time
updateLock *sync.Mutex
interrupt chan struct{}
interrupted atomic.Bool
signalRestart chan error
slogger *slog.Logger
restartFuncs map[autoupdatableBinary]func() error
osquerierBackoffTimeout time.Duration
osquerierBackoffInterval time.Duration
}

type TufAutoupdaterOption func(*TufAutoupdater)
Expand All @@ -115,6 +117,13 @@ func WithOsqueryRestart(restart func() error) TufAutoupdaterOption {
}
}

func WithOsquerierBackoff(timeout, interval time.Duration) TufAutoupdaterOption {
return func(ta *TufAutoupdater) {
ta.osquerierBackoffTimeout = timeout
ta.osquerierBackoffInterval = interval
}
}

func NewTufAutoupdater(ctx context.Context, k types.Knapsack, metadataHttpClient *http.Client, mirrorHttpClient *http.Client,
osquerier querier, opts ...TufAutoupdaterOption) (*TufAutoupdater, error) {
ctx, span := traces.StartSpan(ctx)
Expand All @@ -134,12 +143,14 @@ func NewTufAutoupdater(ctx context.Context, k types.Knapsack, metadataHttpClient
binaryLauncher: func() string { return k.PinnedLauncherVersion() },
binaryOsqueryd: func() string { return k.PinnedOsquerydVersion() },
},
initialDelayEnd: time.Now().Add(k.AutoupdateInitialDelay()),
updateLock: &sync.Mutex{},
osquerier: osquerier,
osquerierRetryInterval: 30 * time.Second,
slogger: k.Slogger().With("component", "tuf_autoupdater"),
restartFuncs: make(map[autoupdatableBinary]func() error),
initialDelayEnd: time.Now().Add(k.AutoupdateInitialDelay()),
updateLock: &sync.Mutex{},
osquerier: osquerier,
osquerierRetryInterval: 30 * time.Second,
slogger: k.Slogger().With("component", "tuf_autoupdater"),
restartFuncs: make(map[autoupdatableBinary]func() error),
osquerierBackoffTimeout: 30 * time.Second, // Default production value
osquerierBackoffInterval: 5 * time.Second, // Default production value
}

for _, opt := range opts {
Expand Down Expand Up @@ -784,7 +795,7 @@ func (ta *TufAutoupdater) collectAndSetEnrollmentDetails(ctx context.Context) er
)
}
return err
}, 30*time.Second, 5*time.Second); err != nil {
}, ta.osquerierBackoffTimeout, ta.osquerierBackoffInterval); err != nil {
if os.Getenv("LAUNCHER_DEBUG_ENROLL_DETAILS_REQUIRED") == "true" {
return fmt.Errorf("query osq enrollment details: %w", err)
}
Expand Down
103 changes: 89 additions & 14 deletions ee/tuf/autoupdate_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ func TestNewTufAutoupdater(t *testing.T) {
mockKnapsack.On("PinnedOsquerydVersion").Return("")
mockKnapsack.On("RegisterChangeObserver", mock.Anything, keys.UpdateChannel, keys.PinnedLauncherVersion, keys.PinnedOsquerydVersion).Return()

_, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, newMockQuerier(t))
_, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, newMockQuerier(t), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Confirm we pulled all config items as expected
Expand Down Expand Up @@ -122,7 +122,7 @@ func TestExecute_launcherUpdate(t *testing.T) {
mockKnapsack.On("Slogger").Return(slogger.Logger)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier)
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsquerierBackoff(1*time.Second, 1*time.Second))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -222,7 +222,7 @@ func TestExecute_osquerydUpdate(t *testing.T) {
mockKnapsack.On("Slogger").Return(slogger.Logger)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -305,7 +305,7 @@ func TestExecute_downgrade(t *testing.T) {
mockKnapsack.On("Slogger").Return(slogger.Logger)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -395,7 +395,7 @@ func TestExecute_withInitialDelay(t *testing.T) {

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient,
mockQuerier, WithOsqueryRestart(func() error { return nil }))
mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Expect that we interrupt
Expand Down Expand Up @@ -458,7 +458,7 @@ func TestExecute_inModernStandby(t *testing.T) {

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient,
mockQuerier, WithOsqueryRestart(func() error { return nil }))
mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Set up library manager: we should expect to tidy the library on startup, but NOT add anything to it
Expand Down Expand Up @@ -519,7 +519,7 @@ func TestInterrupt_Multiple(t *testing.T) {

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient,
mockQuerier, WithOsqueryRestart(func() error { return nil }))
mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Set up normal library and querier interactions
Expand Down Expand Up @@ -655,7 +655,7 @@ func TestDo(t *testing.T) {
mockKnapsack.On("SetEnrollmentDetails", mock.Anything).Return(nil).Maybe()

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -731,7 +731,7 @@ func TestDo_HandlesSimultaneousUpdates(t *testing.T) {
mockKnapsack.On("SetEnrollmentDetails", mock.Anything).Return(nil)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -816,7 +816,7 @@ func TestDo_WillNotExecuteDuringInitialDelay(t *testing.T) {
mockKnapsack.On("LatestOsquerydPath", mock.Anything).Return(fakeOsqBinaryPath)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

// Update the metadata client with our test root JSON
Expand Down Expand Up @@ -900,7 +900,7 @@ func TestFlagsChanged_UpdateChannelChanged(t *testing.T) {
mockKnapsack.On("UpdateChannel").Return("nightly")

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")
require.Equal(t, "beta", autoupdater.updateChannel)

Expand Down Expand Up @@ -967,7 +967,7 @@ func TestFlagsChanged_PinnedVersionChanged(t *testing.T) {
mockKnapsack.On("PinnedOsquerydVersion").Return(pinnedOsquerydVersion)

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")
require.Equal(t, "", autoupdater.pinnedVersions[binaryOsqueryd])

Expand Down Expand Up @@ -1026,7 +1026,7 @@ func TestFlagsChanged_DuringInitialDelay(t *testing.T) {
mockKnapsack.On("PinnedLauncherVersion").Return("")

// Set up autoupdater
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }))
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsqueryRestart(func() error { return nil }), WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")
require.Equal(t, pinnedLauncherVersion, autoupdater.pinnedVersions[binaryLauncher])

Expand Down Expand Up @@ -1157,7 +1157,7 @@ func Test_storeError(t *testing.T) {
mockKnapsack.On("LatestOsquerydPath", mock.Anything).Return(fakeOsqBinaryPath)
mockQuerier := newMockQuerier(t)

autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier)
autoupdater, err := NewTufAutoupdater(context.TODO(), mockKnapsack, http.DefaultClient, http.DefaultClient, mockQuerier, WithOsquerierBackoff(100*time.Millisecond, 20*time.Millisecond))
require.NoError(t, err, "could not initialize new TUF autoupdater")

mockLibraryManager := NewMocklibrarian(t)
Expand Down Expand Up @@ -1239,6 +1239,81 @@ func Test_cleanUpOldErrors(t *testing.T) {
require.Equal(t, 1, keyCount, "cleanup routine did not clean up correct number of old errors")
}

func TestCollectAndSetEnrollmentDetails(t *testing.T) {
t.Parallel()

testCases := []struct {
name string
osquerydPath string
queryResponse []map[string]string
queryError error
expectError bool
}{
{
name: "successful collection",
osquerydPath: "/test/osqueryd",
queryResponse: []map[string]string{{
"osquery_version": "1.1.1",
"os_build": "test",
"os_name": "test",
"os_version": "test",
"hardware_model": "test",
"hardware_serial": "test",
"hardware_vendor": "test",
"hostname": "test",
"hardware_uuid": "test",
}},
expectError: false,
},
{
name: "no osqueryd path",
osquerydPath: "",
expectError: false,
},
{
name: "query fails",
osquerydPath: "/test/osqueryd",
queryError: errors.New("test error"),
expectError: false, // We don't return error unless env var is set
},
}

for _, tt := range testCases {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

mockKnapsack := typesmocks.NewKnapsack(t)
mockKnapsack.On("LatestOsquerydPath", mock.Anything).Return(tt.osquerydPath)
mockKnapsack.On("SetEnrollmentDetails", mock.AnythingOfType("types.EnrollmentDetails")).Return(nil).Maybe()
mockKnapsack.On("Slogger").Return(multislogger.NewNopLogger())

mockQuerier := newMockQuerier(t)
if tt.osquerydPath != "" {
mockQuerier.On("Query", mock.Anything).Return(tt.queryResponse, tt.queryError).Maybe()
}

autoupdater := &TufAutoupdater{
knapsack: mockKnapsack,
osquerier: mockQuerier,
slogger: mockKnapsack.Slogger(),
osquerierBackoffTimeout: 100 * time.Millisecond,
osquerierRetryInterval: 20 * time.Millisecond,
}

err := autoupdater.collectAndSetEnrollmentDetails(context.Background())
if tt.expectError {
require.Error(t, err)
} else {
require.NoError(t, err)
}

mockKnapsack.AssertExpectations(t)
mockQuerier.AssertExpectations(t)
})
}
}

func setupStorage(t *testing.T) types.KVStore {
s, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.AutoupdateErrorsStore.String())
require.NoError(t, err)
Expand Down

0 comments on commit 25e9a93

Please sign in to comment.