Add basic registration tracking to knapsack (kolide#1980)

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 {

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.

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
+}

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")
 	}
 

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))

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",

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")
 	}

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")