Add registration ID to osquery history (#1982)

kolide · Dec 9, 2024 · b36c719 · b36c719
1 parent 9818435
commit b36c719
Show file tree

Hide file tree

Showing 9 changed files with 138 additions and 35 deletions.
diff --git a/ee/debug/checkups/osquery_restarts.go b/ee/debug/checkups/osquery_restarts.go
@@ -45,12 +45,14 @@ func (orc *osqRestartCheckup) Run(ctx context.Context, extraFH io.Writer) error
 
 	for idx, instance := range restartHistory {
 		results[idx] = map[string]string{
-			"start_time":   instance.StartTime,
-			"connect_time": instance.ConnectTime,
-			"exit_time":    instance.ExitTime,
-			"instance_id":  instance.InstanceId,
-			"version":      instance.Version,
-			"errors":       instance.Error,
+			"registration_id": instance.RegistrationId,
+			"instance_run_id": instance.RunId,
+			"start_time":      instance.StartTime,
+			"connect_time":    instance.ConnectTime,
+			"exit_time":       instance.ExitTime,
+			"instance_id":     instance.InstanceId,
+			"version":         instance.Version,
+			"errors":          instance.Error,
 		}
 	}
 

diff --git a/ee/tables/osquery_instance_history/osquery_instance_history.go b/ee/tables/osquery_instance_history/osquery_instance_history.go
@@ -9,6 +9,7 @@ import (
 
 func TablePlugin() *table.Plugin {
 	columns := []table.ColumnDefinition{
+		table.TextColumn("registration_id"),
 		table.TextColumn("instance_run_id"),
 		table.TextColumn("start_time"),
 		table.TextColumn("connect_time"),
@@ -33,6 +34,7 @@ func generate() table.GenerateFunc {
 		for _, instance := range history {
 
 			results = append(results, map[string]string{
+				"registration_id": instance.RegistrationId,
 				"instance_run_id": instance.RunId,
 				"start_time":      instance.StartTime,
 				"connect_time":    instance.ConnectTime,

diff --git a/pkg/osquery/runtime/history/history.go b/pkg/osquery/runtime/history/history.go
@@ -61,13 +61,30 @@ func LatestInstance() (Instance, error) {
 	currentHistory.Lock()
 	defer currentHistory.Unlock()
 
-	if currentHistory.instances == nil || len(currentHistory.instances) == 0 {
+	if len(currentHistory.instances) == 0 {
 		return Instance{}, NoInstancesError{}
 	}
 
 	return *currentHistory.instances[len(currentHistory.instances)-1], nil
 }
 
+func LatestInstanceByRegistrationID(registrationId string) (Instance, error) {
+	currentHistory.Lock()
+	defer currentHistory.Unlock()
+
+	if len(currentHistory.instances) == 0 {
+		return Instance{}, NoInstancesError{}
+	}
+
+	for i := len(currentHistory.instances) - 1; i > -1; i -= 1 {
+		if currentHistory.instances[i].RegistrationId == registrationId {
+			return *currentHistory.instances[i], nil
+		}
+	}
+
+	return Instance{}, NoInstancesError{}
+}
+
 func LatestInstanceUptimeMinutes() (int64, error) {
 	lastInstance, err := LatestInstance()
 	if err != nil {
@@ -88,7 +105,7 @@ func LatestInstanceUptimeMinutes() (int64, error) {
 }
 
 // NewInstance adds a new instance to the osquery instance history and returns it
-func NewInstance(runId string) (*Instance, error) {
+func NewInstance(registrationId string, runId string) (*Instance, error) {
 	currentHistory.Lock()
 	defer currentHistory.Unlock()
 
@@ -98,9 +115,10 @@ func NewInstance(runId string) (*Instance, error) {
 	}
 
 	newInstance := &Instance{
-		RunId:     runId,
-		StartTime: timeNow(),
-		Hostname:  hostname,
+		RegistrationId: registrationId,
+		RunId:          runId,
+		StartTime:      timeNow(),
+		Hostname:       hostname,
 	}
 
 	currentHistory.addInstanceToHistory(newInstance)

diff --git a/pkg/osquery/runtime/history/history_test.go b/pkg/osquery/runtime/history/history_test.go
@@ -59,7 +59,7 @@ func TestNewInstance(t *testing.T) { // nolint:paralleltest
 
 			require.NoError(t, InitHistory(setupStorage(t, tt.initialInstances...)))
 
-			_, err := NewInstance(ulid.New())
+			_, err := NewInstance(ulid.New(), ulid.New())
 
 			assert.Equal(t, tt.wantNumInstances, len(currentHistory.instances), "expect history length to reflect new instance")
 
@@ -176,6 +176,82 @@ func TestLatestInstance(t *testing.T) { // nolint:paralleltest
 	}
 }
 
+func TestLatestInstanceForRegistrationID(t *testing.T) { // nolint:paralleltest
+	tests := []struct {
+		name             string
+		registrationID   string
+		initialInstances []*Instance
+		want             Instance
+		errString        string
+	}{
+		{
+			name:           "success",
+			registrationID: "test",
+			initialInstances: []*Instance{
+				{
+					StartTime:      "first_expected_start_time",
+					RegistrationId: "some other",
+				},
+				{
+					StartTime:      "second_expected_start_time",
+					RegistrationId: "test",
+				},
+				{
+					StartTime:      "third_expected_start_time",
+					RegistrationId: "test",
+				},
+				{
+					StartTime:      "fourth_expected_start_time",
+					RegistrationId: "another",
+				},
+			},
+			want: Instance{
+				StartTime:      "third_expected_start_time",
+				RegistrationId: "test",
+			},
+		},
+		{
+			name:           "no_instances_error",
+			registrationID: "test",
+			errString:      NoInstancesError{}.Error(),
+		},
+		{
+			name:           "no matching instances",
+			registrationID: "test3",
+			initialInstances: []*Instance{
+				{
+					StartTime:      "first_expected_start_time",
+					RegistrationId: "test",
+				},
+				{
+					StartTime:      "second_expected_start_time",
+					RegistrationId: "test1",
+				},
+				{
+					StartTime:      "third_expected_start_time",
+					RegistrationId: "test2",
+				},
+			},
+			errString: NoInstancesError{}.Error(),
+		},
+	}
+	for _, tt := range tests { // nolint:paralleltest
+		t.Run(tt.name, func(t *testing.T) {
+			t.Cleanup(func() { currentHistory = &History{} })
+
+			require.NoError(t, InitHistory(setupStorage(t, tt.initialInstances...)))
+
+			got, err := LatestInstanceByRegistrationID(tt.registrationID)
+
+			if tt.errString != "" {
+				assert.EqualError(t, err, tt.errString)
+			}
+
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
 func TestLatestInstanceUptimeMinutes(t *testing.T) { // nolint:paralleltest
 	tests := []struct {
 		name             string

diff --git a/pkg/osquery/runtime/history/instance.go b/pkg/osquery/runtime/history/instance.go
@@ -6,14 +6,15 @@ import (
 )
 
 type Instance struct {
-	RunId       string // ID for instance, assigned by launcher
-	StartTime   string
-	ConnectTime string
-	ExitTime    string
-	Hostname    string
-	InstanceId  string // ID from osquery
-	Version     string
-	Error       string
+	RegistrationId string // which registration this instance belongs to
+	RunId          string // ID for instance, assigned by launcher
+	StartTime      string
+	ConnectTime    string
+	ExitTime       string
+	Hostname       string
+	InstanceId     string // ID from osquery
+	Version        string
+	Error          string
 }
 
 type Querier interface {

diff --git a/pkg/osquery/runtime/osqueryinstance.go b/pkg/osquery/runtime/osqueryinstance.go
@@ -362,7 +362,7 @@ func (i *OsqueryInstance) Launch() error {
 		"osquery socket created",
 	)
 
-	stats, err := history.NewInstance(i.runId)
+	stats, err := history.NewInstance(i.registrationId, i.runId)
 	if err != nil {
 		i.slogger.Log(ctx, slog.LevelWarn,
 			"could not create new osquery instance history",

diff --git a/pkg/osquery/table/launcher_config.go b/pkg/osquery/table/launcher_config.go
@@ -8,25 +8,27 @@ import (
 	"github.com/osquery/osquery-go/plugin/table"
 )
 
-func LauncherConfigTable(store types.Getter) *table.Plugin {
+func LauncherConfigTable(store types.Getter, registrationTracker types.RegistrationTracker) *table.Plugin {
 	columns := []table.ColumnDefinition{
 		table.TextColumn("config"),
+		table.TextColumn("registration_id"),
 	}
-	return table.NewPlugin("kolide_launcher_config", columns, generateLauncherConfig(store))
+	return table.NewPlugin("kolide_launcher_config", columns, generateLauncherConfig(store, registrationTracker))
 }
 
-func generateLauncherConfig(store types.Getter) table.GenerateFunc {
+func generateLauncherConfig(store types.Getter, registrationTracker types.RegistrationTracker) table.GenerateFunc {
 	return func(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
-		config, err := osquery.Config(store)
-		if err != nil {
-			return nil, err
+		results := make([]map[string]string, 0)
+		for _, registrationId := range registrationTracker.RegistrationIDs() {
+			config, err := osquery.Config(store)
+			if err != nil {
+				return nil, err
+			}
+			results = append(results, map[string]string{
+				"config":          config,
+				"registration_id": registrationId,
+			})
 		}
-		results := []map[string]string{
-			{
-				"config": config,
-			},
-		}
-
 		return results, nil
 	}
 }
diff --git a/pkg/osquery/table/launcher_info.go b/pkg/osquery/table/launcher_info.go
@@ -30,6 +30,7 @@ func LauncherInfoTable(configStore types.GetterSetter, LauncherHistoryStore type
 		table.TextColumn("revision"),
 		table.TextColumn("version"),
 		table.TextColumn("version_chain"),
+		table.TextColumn("registration_id"),
 		table.TextColumn("identifier"),
 		table.TextColumn("osquery_instance_id"),
 		table.TextColumn("uptime"),
@@ -57,7 +58,7 @@ func generateLauncherInfoTable(configStore types.GetterSetter, LauncherHistorySt
 			return nil, err
 		}
 
-		osqueryInstance, err := history.LatestInstance()
+		osqueryInstance, err := history.LatestInstanceByRegistrationID(types.DefaultRegistrationID)
 		if err != nil {
 			return nil, err
 		}
@@ -91,6 +92,7 @@ func generateLauncherInfoTable(configStore types.GetterSetter, LauncherHistorySt
 				"revision":            version.Version().Revision,
 				"version":             version.Version().Version,
 				"version_chain":       os.Getenv("KOLIDE_LAUNCHER_VERSION_CHAIN"),
+				"registration_id":     types.DefaultRegistrationID,
 				"identifier":          identifier,
 				"osquery_instance_id": osqueryInstance.InstanceId,
 				"fingerprint":         fingerprint,

diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go
@@ -29,7 +29,7 @@ import (
 // around _launcher_ things thus do not make sense in tables.ext
 func LauncherTables(k types.Knapsack) []osquery.OsqueryPlugin {
 	return []osquery.OsqueryPlugin{
-		LauncherConfigTable(k.ConfigStore()),
+		LauncherConfigTable(k.ConfigStore(), k),
 		LauncherDbInfo(k.BboltDB()),
 		LauncherInfoTable(k.ConfigStore(), k.LauncherHistoryStore()),
 		launcher_db.TablePlugin("kolide_server_data", k.ServerProvidedDataStore()),