diff --git a/ee/agent/knapsack/knapsack.go b/ee/agent/knapsack/knapsack.go index 84c3fb000..f52c63ce1 100644 --- a/ee/agent/knapsack/knapsack.go +++ b/ee/agent/knapsack/knapsack.go @@ -92,6 +92,11 @@ func (k *knapsack) SetInstanceQuerier(q types.InstanceQuerier) { k.querier = q } +// RegistrationTracker interface methods +func (k *knapsack) RegistrationIDs() []string { + return []string{types.DefaultRegistrationID} +} + // InstanceStatuses returns the current status of each osquery instance. // It performs a healthcheck against each existing instance. func (k *knapsack) InstanceStatuses() map[string]types.InstanceStatus { diff --git a/ee/agent/types/knapsack.go b/ee/agent/types/knapsack.go index ddb1d6880..45a441cb1 100644 --- a/ee/agent/types/knapsack.go +++ b/ee/agent/types/knapsack.go @@ -9,6 +9,7 @@ type Knapsack interface { BboltDB Flags Slogger + RegistrationTracker InstanceQuerier SetInstanceQuerier(q InstanceQuerier) // LatestOsquerydPath finds the path to the latest osqueryd binary, after accounting for updates. diff --git a/ee/agent/types/mocks/knapsack.go b/ee/agent/types/mocks/knapsack.go index 5b3f464a5..58bea406d 100644 --- a/ee/agent/types/mocks/knapsack.go +++ b/ee/agent/types/mocks/knapsack.go @@ -1102,6 +1102,26 @@ func (_m *Knapsack) RegisterChangeObserver(observer types.FlagsChangeObserver, f _m.Called(_ca...) } +// RegistrationIDs provides a mock function with given fields: +func (_m *Knapsack) RegistrationIDs() []string { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for RegistrationIDs") + } + + var r0 []string + if rf, ok := ret.Get(0).(func() []string); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]string) + } + } + + return r0 +} + // ResultLogsStore provides a mock function with given fields: func (_m *Knapsack) ResultLogsStore() types.GetterSetterDeleterIteratorUpdaterCounterAppender { ret := _m.Called() diff --git a/ee/agent/types/registration.go b/ee/agent/types/registration.go new file mode 100644 index 000000000..979889e6c --- /dev/null +++ b/ee/agent/types/registration.go @@ -0,0 +1,12 @@ +package types + +const ( + DefaultRegistrationID = "default" +) + +// RegistrationTracker manages the current set of registrations for this launcher installation. +// Right now, the list is hardcoded to only the default registration ID. In the future, this +// data may be provided by e.g. a control server subsystem. +type RegistrationTracker interface { + RegistrationIDs() []string +} diff --git a/pkg/osquery/runtime/osqueryinstance.go b/pkg/osquery/runtime/osqueryinstance.go index 190a71839..8e3a3961a 100644 --- a/pkg/osquery/runtime/osqueryinstance.go +++ b/pkg/osquery/runtime/osqueryinstance.go @@ -708,7 +708,7 @@ func calculateOsqueryPaths(rootDirectory string, registrationId string, runId st } // Keep default database path for default instance - if registrationId == defaultRegistrationId { + if registrationId == types.DefaultRegistrationID { osqueryFilePaths.databasePath = filepath.Join(rootDirectory, "osquery.db") } diff --git a/pkg/osquery/runtime/osqueryinstance_test.go b/pkg/osquery/runtime/osqueryinstance_test.go index 8e6c1b479..cf59febab 100644 --- a/pkg/osquery/runtime/osqueryinstance_test.go +++ b/pkg/osquery/runtime/osqueryinstance_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/kolide/kit/ulid" + "github.com/kolide/launcher/ee/agent/types" typesMocks "github.com/kolide/launcher/ee/agent/types/mocks" "github.com/kolide/launcher/pkg/log/multislogger" "github.com/stretchr/testify/assert" @@ -23,7 +24,7 @@ func TestCalculateOsqueryPaths(t *testing.T) { rootDir := t.TempDir() runId := ulid.New() - paths, err := calculateOsqueryPaths(rootDir, defaultRegistrationId, runId, osqueryOptions{}) + paths, err := calculateOsqueryPaths(rootDir, types.DefaultRegistrationID, runId, osqueryOptions{}) require.NoError(t, err) @@ -61,7 +62,7 @@ func TestCreateOsqueryCommand(t *testing.T) { k.On("OsqueryFlags").Return([]string{}) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient(), WithStdout(os.Stdout), WithStderr(os.Stderr)) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient(), WithStdout(os.Stdout), WithStderr(os.Stderr)) cmd, err := i.createOsquerydCommand(osquerydPath, paths) require.NoError(t, err) @@ -83,7 +84,7 @@ func TestCreateOsqueryCommandWithFlags(t *testing.T) { k.On("OsqueryVerbose").Return(true) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) cmd, err := i.createOsquerydCommand( testOsqueryBinaryDirectory, @@ -116,7 +117,7 @@ func TestCreateOsqueryCommand_SetsEnabledWatchdogSettingsAppropriately(t *testin k.On("OsqueryVerbose").Return(true) k.On("OsqueryFlags").Return([]string{}) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) cmd, err := i.createOsquerydCommand( testOsqueryBinaryDirectory, @@ -165,7 +166,7 @@ func TestCreateOsqueryCommand_SetsDisabledWatchdogSettingsAppropriately(t *testi k.On("OsqueryVerbose").Return(true) k.On("OsqueryFlags").Return([]string{}) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) cmd, err := i.createOsquerydCommand( testOsqueryBinaryDirectory, @@ -206,7 +207,7 @@ func TestHealthy_DoesNotPassForUnlaunchedInstance(t *testing.T) { k := typesMocks.NewKnapsack(t) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) require.Error(t, i.Healthy(), "unlaunched instance should not return healthy status") } @@ -217,7 +218,7 @@ func TestQuery_ReturnsErrorForUnlaunchedInstance(t *testing.T) { k := typesMocks.NewKnapsack(t) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) _, err := i.Query("select * from osquery_info;") require.Error(t, err, "should not be able to query unlaunched instance") @@ -228,7 +229,7 @@ func Test_healthcheckWithRetries(t *testing.T) { k := typesMocks.NewKnapsack(t) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) // No client available, so healthcheck should fail despite retries require.Error(t, i.healthcheckWithRetries(context.TODO(), 5, 100*time.Millisecond)) diff --git a/pkg/osquery/runtime/osqueryinstance_windows_test.go b/pkg/osquery/runtime/osqueryinstance_windows_test.go index 67c26d3a8..c5b31528b 100644 --- a/pkg/osquery/runtime/osqueryinstance_windows_test.go +++ b/pkg/osquery/runtime/osqueryinstance_windows_test.go @@ -7,6 +7,7 @@ import ( "strings" "testing" + "github.com/kolide/launcher/ee/agent/types" typesMocks "github.com/kolide/launcher/ee/agent/types/mocks" "github.com/kolide/launcher/pkg/log/multislogger" "github.com/stretchr/testify/require" @@ -26,7 +27,7 @@ func TestCreateOsqueryCommandEnvVars(t *testing.T) { k.On("OsqueryFlags").Return([]string{}) k.On("Slogger").Return(multislogger.NewNopLogger()) - i := newInstance(defaultRegistrationId, k, mockServiceClient()) + i := newInstance(types.DefaultRegistrationID, k, mockServiceClient()) cmd, err := i.createOsquerydCommand(osquerydPath, &osqueryFilePaths{ pidfilePath: "/foo/bar/osquery-abcd.pid", diff --git a/pkg/osquery/runtime/runner.go b/pkg/osquery/runtime/runner.go index b1bb4ea15..3f926c46d 100644 --- a/pkg/osquery/runtime/runner.go +++ b/pkg/osquery/runtime/runner.go @@ -15,8 +15,6 @@ import ( ) const ( - defaultRegistrationId = "default" - launchRetryDelay = 30 * time.Second ) @@ -34,16 +32,13 @@ type Runner struct { func New(k types.Knapsack, serviceClient service.KolideService, opts ...OsqueryInstanceOption) *Runner { runner := &Runner{ - registrationIds: []string{ - // For now, we only have one (default) instance and we use it for all queries - defaultRegistrationId, - }, - instances: make(map[string]*OsqueryInstance), - slogger: k.Slogger().With("component", "osquery_runner"), - knapsack: k, - serviceClient: serviceClient, - shutdown: make(chan struct{}), - opts: opts, + registrationIds: k.RegistrationIDs(), + instances: make(map[string]*OsqueryInstance), + slogger: k.Slogger().With("component", "osquery_runner"), + knapsack: k, + serviceClient: serviceClient, + shutdown: make(chan struct{}), + opts: opts, } k.RegisterChangeObserver(runner, @@ -182,7 +177,7 @@ func (r *Runner) Query(query string) ([]map[string]string, error) { defer r.instanceLock.Unlock() // For now, grab the default (i.e. only) instance - instance, ok := r.instances[defaultRegistrationId] + instance, ok := r.instances[types.DefaultRegistrationID] if !ok { return nil, errors.New("no default instance exists, cannot query") } diff --git a/pkg/osquery/runtime/runtime_posix_test.go b/pkg/osquery/runtime/runtime_posix_test.go index f0fa70bae..03554e176 100644 --- a/pkg/osquery/runtime/runtime_posix_test.go +++ b/pkg/osquery/runtime/runtime_posix_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/kolide/launcher/ee/agent/types" typesMocks "github.com/kolide/launcher/ee/agent/types/mocks" "github.com/osquery/osquery-go" "github.com/stretchr/testify/assert" @@ -41,6 +42,7 @@ func TestOsquerySlowStart(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RootDirectory").Return(rootDirectory).Maybe() @@ -89,6 +91,7 @@ func TestExtensionSocketPath(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -134,24 +137,24 @@ func TestRestart(t *testing.T) { runner, logBytes, teardown := setupOsqueryInstanceForTests(t) defer teardown() - previousStats := runner.instances[defaultRegistrationId].stats + previousStats := runner.instances[types.DefaultRegistrationID].stats require.NoError(t, runner.Restart()) waitHealthy(t, runner, logBytes) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.StartTime, "start time should be set on latest instance stats after restart") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ConnectTime, "connect time should be set on latest instance stats after restart") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.StartTime, "start time should be set on latest instance stats after restart") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ConnectTime, "connect time should be set on latest instance stats after restart") require.NotEmpty(t, previousStats.ExitTime, "exit time should be set on last instance stats when restarted") require.NotEmpty(t, previousStats.Error, "stats instance should have an error on restart") - previousStats = runner.instances[defaultRegistrationId].stats + previousStats = runner.instances[types.DefaultRegistrationID].stats require.NoError(t, runner.Restart()) waitHealthy(t, runner, logBytes) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.StartTime, "start time should be added to latest instance stats after restart") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ConnectTime, "connect time should be added to latest instance stats after restart") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.StartTime, "start time should be added to latest instance stats after restart") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ConnectTime, "connect time should be added to latest instance stats after restart") require.NotEmpty(t, previousStats.ExitTime, "exit time should be set on instance stats when restarted") require.NotEmpty(t, previousStats.Error, "stats instance should have an error on restart") diff --git a/pkg/osquery/runtime/runtime_test.go b/pkg/osquery/runtime/runtime_test.go index d57f833ba..75d981ed6 100644 --- a/pkg/osquery/runtime/runtime_test.go +++ b/pkg/osquery/runtime/runtime_test.go @@ -122,6 +122,7 @@ func TestBadBinaryPath(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -158,6 +159,7 @@ func TestWithOsqueryFlags(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -186,6 +188,7 @@ func TestFlagsChanged(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() // First, it should return false, then on the next call, it should return true k.On("WatchdogEnabled").Return(false).Once() @@ -215,7 +218,7 @@ func TestFlagsChanged(t *testing.T) { // Confirm watchdog is disabled watchdogDisabled := false - for _, a := range runner.instances[defaultRegistrationId].cmd.Args { + for _, a := range runner.instances[types.DefaultRegistrationID].cmd.Args { if strings.Contains(a, "disable_watchdog") { watchdogDisabled = true break @@ -223,7 +226,7 @@ func TestFlagsChanged(t *testing.T) { } require.True(t, watchdogDisabled, "instance not set up with watchdog disabled") - startingInstance := runner.instances[defaultRegistrationId] + startingInstance := runner.instances[types.DefaultRegistrationID] runner.FlagsChanged(keys.WatchdogEnabled) @@ -232,13 +235,13 @@ func TestFlagsChanged(t *testing.T) { waitHealthy(t, runner, logBytes) // Now confirm that the instance is new - require.NotEqual(t, startingInstance, runner.instances[defaultRegistrationId], "instance not replaced") + require.NotEqual(t, startingInstance, runner.instances[types.DefaultRegistrationID], "instance not replaced") // Confirm osquery watchdog is now enabled watchdogMemoryLimitMBFound := false watchdogUtilizationLimitPercentFound := false watchdogDelaySecFound := false - for _, a := range runner.instances[defaultRegistrationId].cmd.Args { + for _, a := range runner.instances[types.DefaultRegistrationID].cmd.Args { if strings.Contains(a, "disable_watchdog") { t.Error("disable_watchdog flag set") t.FailNow() @@ -299,10 +302,10 @@ func waitHealthy(t *testing.T, runner *Runner, logBytes *threadsafebuffer.Thread } // Confirm osquery instance setup is complete - if runner.instances[defaultRegistrationId] == nil { + if runner.instances[types.DefaultRegistrationID] == nil { return errors.New("default instance does not exist yet") } - if runner.instances[defaultRegistrationId].stats == nil || runner.instances[defaultRegistrationId].stats.ConnectTime == "" { + if runner.instances[types.DefaultRegistrationID].stats == nil || runner.instances[types.DefaultRegistrationID].stats.ConnectTime == "" { return errors.New("no connect time set yet") } @@ -321,6 +324,7 @@ func TestSimplePath(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -340,8 +344,8 @@ func TestSimplePath(t *testing.T) { waitHealthy(t, runner, logBytes) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.StartTime, "start time should be added to instance stats on start up") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ConnectTime, "connect time should be added to instance stats on start up") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.StartTime, "start time should be added to instance stats on start up") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ConnectTime, "connect time should be added to instance stats on start up") waitShutdown(t, runner, logBytes) } @@ -352,7 +356,11 @@ func TestMultipleInstances(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) + // Add in an extra instance + extraRegistrationId := ulid.New() + k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID, extraRegistrationId}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -370,19 +378,15 @@ func TestMultipleInstances(t *testing.T) { runner := New(k, serviceClient, opts...) - // Add in an extra instance - extraRegistrationId := ulid.New() - runner.registrationIds = append(runner.registrationIds, extraRegistrationId) - // Start the instance go runner.Run() waitHealthy(t, runner, logBytes) // Confirm the default instance was started - require.Contains(t, runner.instances, defaultRegistrationId) - require.NotNil(t, runner.instances[defaultRegistrationId].stats) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.StartTime, "start time should be added to default instance stats on start up") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ConnectTime, "connect time should be added to default instance stats on start up") + require.Contains(t, runner.instances, types.DefaultRegistrationID) + require.NotNil(t, runner.instances[types.DefaultRegistrationID].stats) + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.StartTime, "start time should be added to default instance stats on start up") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ConnectTime, "connect time should be added to default instance stats on start up") // Confirm the additional instance was started require.Contains(t, runner.instances, extraRegistrationId) @@ -392,17 +396,17 @@ func TestMultipleInstances(t *testing.T) { // Confirm instance statuses are reported correctly instanceStatuses := runner.InstanceStatuses() - require.Contains(t, instanceStatuses, defaultRegistrationId) - require.Equal(t, instanceStatuses[defaultRegistrationId], types.InstanceStatusHealthy) + require.Contains(t, instanceStatuses, types.DefaultRegistrationID) + require.Equal(t, instanceStatuses[types.DefaultRegistrationID], types.InstanceStatusHealthy) require.Contains(t, instanceStatuses, extraRegistrationId) require.Equal(t, instanceStatuses[extraRegistrationId], types.InstanceStatusHealthy) waitShutdown(t, runner, logBytes) // Confirm both instances exited - require.Contains(t, runner.instances, defaultRegistrationId) - require.NotNil(t, runner.instances[defaultRegistrationId].stats) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ExitTime, "exit time should be added to default instance stats on shutdown") + require.Contains(t, runner.instances, types.DefaultRegistrationID) + require.NotNil(t, runner.instances[types.DefaultRegistrationID].stats) + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ExitTime, "exit time should be added to default instance stats on shutdown") require.Contains(t, runner.instances, extraRegistrationId) require.NotNil(t, runner.instances[extraRegistrationId].stats) require.NotEmpty(t, runner.instances[extraRegistrationId].stats.ExitTime, "exit time should be added to secondary instance stats on shutdown") @@ -415,6 +419,7 @@ func TestRunnerHandlesImmediateShutdownWithMultipleInstances(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -444,11 +449,11 @@ func TestRunnerHandlesImmediateShutdownWithMultipleInstances(t *testing.T) { waitShutdown(t, runner, logBytes) // Confirm the default instance was started, and then exited - require.Contains(t, runner.instances, defaultRegistrationId) - require.NotNil(t, runner.instances[defaultRegistrationId].stats) - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.StartTime, "start time should be added to default instance stats on start up") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ConnectTime, "connect time should be added to default instance stats on start up") - require.NotEmpty(t, runner.instances[defaultRegistrationId].stats.ExitTime, "exit time should be added to default instance stats on shutdown") + require.Contains(t, runner.instances, types.DefaultRegistrationID) + require.NotNil(t, runner.instances[types.DefaultRegistrationID].stats) + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.StartTime, "start time should be added to default instance stats on start up") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ConnectTime, "connect time should be added to default instance stats on start up") + require.NotEmpty(t, runner.instances[types.DefaultRegistrationID].stats.ExitTime, "exit time should be added to default instance stats on shutdown") // Confirm the additional instance was started, and then exited require.Contains(t, runner.instances, extraRegistrationId) @@ -465,6 +470,7 @@ func TestMultipleShutdowns(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -496,6 +502,7 @@ func TestOsqueryDies(t *testing.T) { logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(false) k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) @@ -515,14 +522,14 @@ func TestOsqueryDies(t *testing.T) { waitHealthy(t, runner, logBytes) - require.Contains(t, runner.instances, defaultRegistrationId) - require.NotNil(t, runner.instances[defaultRegistrationId].stats) - previousStats := runner.instances[defaultRegistrationId].stats + require.Contains(t, runner.instances, types.DefaultRegistrationID) + require.NotNil(t, runner.instances[types.DefaultRegistrationID].stats) + previousStats := runner.instances[types.DefaultRegistrationID].stats // Simulate the osquery process unexpectedly dying runner.instanceLock.Lock() - require.NoError(t, killProcessGroup(runner.instances[defaultRegistrationId].cmd)) - runner.instances[defaultRegistrationId].errgroup.Wait() + require.NoError(t, killProcessGroup(runner.instances[types.DefaultRegistrationID].cmd)) + runner.instances[types.DefaultRegistrationID].errgroup.Wait() runner.instanceLock.Unlock() waitHealthy(t, runner, logBytes) @@ -537,6 +544,7 @@ func TestNotStarted(t *testing.T) { rootDirectory := t.TempDir() k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("RootDirectory").Return(rootDirectory).Maybe() k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return() @@ -566,10 +574,10 @@ func TestExtensionIsCleanedUp(t *testing.T) { runner, logBytes, teardown := setupOsqueryInstanceForTests(t) defer teardown() - requirePgidMatch(t, runner.instances[defaultRegistrationId].cmd.Process.Pid) + requirePgidMatch(t, runner.instances[types.DefaultRegistrationID].cmd.Process.Pid) // kill the current osquery process but not the extension - require.NoError(t, runner.instances[defaultRegistrationId].cmd.Process.Kill()) + require.NoError(t, runner.instances[types.DefaultRegistrationID].cmd.Process.Kill()) // We need to (a) let the runner restart osquery, and (b) wait for // the extension to die. Both of these may take up to 30s. We'll @@ -593,6 +601,7 @@ func setupOsqueryInstanceForTests(t *testing.T) (runner *Runner, logBytes *threa logBytes, slogger, opts := setUpTestSlogger(rootDirectory) k := typesMocks.NewKnapsack(t) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) k.On("OsqueryHealthcheckStartupDelay").Return(0 * time.Second).Maybe() k.On("WatchdogEnabled").Return(true) k.On("WatchdogMemoryLimitMB").Return(150) @@ -614,7 +623,7 @@ func setupOsqueryInstanceForTests(t *testing.T) (runner *Runner, logBytes *threa go runner.Run() waitHealthy(t, runner, logBytes) - requirePgidMatch(t, runner.instances[defaultRegistrationId].cmd.Process.Pid) + requirePgidMatch(t, runner.instances[types.DefaultRegistrationID].cmd.Process.Pid) teardown = func() { waitShutdown(t, runner, logBytes)