Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add basic registration tracking to knapsack #1980

Merged
merged 2 commits into from
Dec 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions ee/agent/knapsack/knapsack.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions ee/agent/types/knapsack.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
20 changes: 20 additions & 0 deletions ee/agent/types/mocks/knapsack.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 12 additions & 0 deletions ee/agent/types/registration.go
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 1 addition & 1 deletion pkg/osquery/runtime/osqueryinstance.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}

Expand Down
17 changes: 9 additions & 8 deletions pkg/osquery/runtime/osqueryinstance_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")
}
Expand All @@ -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")
Expand All @@ -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))
Expand Down
3 changes: 2 additions & 1 deletion pkg/osquery/runtime/osqueryinstance_windows_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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",
Expand Down
21 changes: 8 additions & 13 deletions pkg/osquery/runtime/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@ import (
)

const (
defaultRegistrationId = "default"

launchRetryDelay = 30 * time.Second
)

Expand All @@ -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,
Expand Down Expand Up @@ -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")
}
Expand Down
15 changes: 9 additions & 6 deletions pkg/osquery/runtime/runtime_posix_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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")
Expand Down
Loading
Loading