From a78a084efc12b495105d29fb8f319bf0f95c82b7 Mon Sep 17 00:00:00 2001 From: RebeccaMahany Date: Mon, 9 Dec 2024 13:21:37 -0500 Subject: [PATCH 1/3] Store osquery data per-registration in config store --- ee/agent/storage/keys.go | 32 +++++++++++++++++++ ee/agent/storage/keys_test.go | 48 ++++++++++++++++++++++++++++ pkg/osquery/extension.go | 29 +++++++++-------- pkg/osquery/extension_test.go | 8 ++--- pkg/osquery/table/launcher_config.go | 2 +- pkg/osquery/table/launcher_info.go | 2 +- 6 files changed, 101 insertions(+), 20 deletions(-) create mode 100644 ee/agent/storage/keys_test.go diff --git a/ee/agent/storage/keys.go b/ee/agent/storage/keys.go index 2f5ba7c05..8e029ca33 100644 --- a/ee/agent/storage/keys.go +++ b/ee/agent/storage/keys.go @@ -1,5 +1,37 @@ package storage +import "bytes" + var ( + // Well-known keys ObservabilityIngestAuthTokenKey = []byte("observability_ingest_auth_token") + + // Identifier types in complex keys + IdentifierTypeRegistration = []byte("registration") + + defaultIdentifier = []byte("default") +) + +const ( + keyDelimiter byte = 58 // : ) + +func KeyByIdentifier(key []byte, identifierType []byte, identifier []byte) []byte { + // The default value is stored under `key`, without any identifier + if len(identifier) == 0 || bytes.Equal(identifier, defaultIdentifier) { + return key + } + + // Key will take the form `::` -- allocate + // a new key with the appropriate capacity. + totalSize := len(key) + 1 + len(identifierType) + 1 + len(identifier) + newKey := make([]byte, 0, totalSize) + + newKey = append(newKey, key...) + newKey = append(newKey, keyDelimiter) + newKey = append(newKey, identifierType...) + newKey = append(newKey, keyDelimiter) + newKey = append(newKey, identifier...) + + return newKey +} diff --git a/ee/agent/storage/keys_test.go b/ee/agent/storage/keys_test.go new file mode 100644 index 000000000..0d5d68b29 --- /dev/null +++ b/ee/agent/storage/keys_test.go @@ -0,0 +1,48 @@ +package storage + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestKeyByIdentifier(t *testing.T) { + t.Parallel() + + for _, tt := range []struct { + testCaseName string + key []byte + identifierType []byte + identifier []byte + expectedKey []byte + }{ + { + testCaseName: "default identifier", + key: []byte("nodeKey"), + identifierType: IdentifierTypeRegistration, + identifier: []byte("default"), + expectedKey: []byte("nodeKey"), + }, + { + testCaseName: "empty identifier", + key: []byte("config"), + identifierType: IdentifierTypeRegistration, + identifier: nil, + expectedKey: []byte("config"), + }, + { + testCaseName: "registration identifier", + key: []byte("uuid"), + identifierType: IdentifierTypeRegistration, + identifier: []byte("some-test-registration-id"), + expectedKey: []byte("uuid:registration:some-test-registration-id"), + }, + } { + tt := tt + t.Run(tt.testCaseName, func(t *testing.T) { + t.Parallel() + + require.Equal(t, tt.expectedKey, KeyByIdentifier(tt.key, tt.identifierType, tt.identifier)) + }) + } +} diff --git a/pkg/osquery/extension.go b/pkg/osquery/extension.go index ecf1c7feb..2898550a8 100644 --- a/pkg/osquery/extension.go +++ b/pkg/osquery/extension.go @@ -15,6 +15,7 @@ import ( "github.com/google/uuid" "github.com/kolide/launcher/ee/agent/startupsettings" + "github.com/kolide/launcher/ee/agent/storage" "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/ee/uninstall" "github.com/kolide/launcher/pkg/backoff" @@ -118,7 +119,7 @@ func NewExtension(ctx context.Context, client service.KolideService, k types.Kna configStore := k.ConfigStore() - nodekey, err := NodeKey(configStore) + nodekey, err := NodeKey(configStore, registrationId) if err != nil { slogger.Log(ctx, slog.LevelDebug, "NewExtension got error reading nodekey. Ignoring", @@ -183,7 +184,7 @@ func (e *Extension) Shutdown(_ error) { // there is an existing identifier, that should be returned. If not, the // identifier should be randomly generated and persisted. func (e *Extension) getHostIdentifier() (string, error) { - return IdentifierFromDB(e.knapsack.ConfigStore()) + return IdentifierFromDB(e.knapsack.ConfigStore(), e.registrationId) } // SetupLauncherKeys configures the various keys used for communication. @@ -276,9 +277,9 @@ func PublicRSAKeyFromDB(configStore types.Getter) (string, string, error) { // IdentifierFromDB returns the built-in launcher identifier from the config bucket. // The function is exported to allow for building the kolide_launcher_info table. -func IdentifierFromDB(configStore types.GetterSetter) (string, error) { +func IdentifierFromDB(configStore types.GetterSetter, registrationId string) (string, error) { var identifier string - uuidBytes, _ := configStore.Get([]byte(uuidKey)) + uuidBytes, _ := configStore.Get(storage.KeyByIdentifier([]byte(uuidKey), storage.IdentifierTypeRegistration, []byte(registrationId))) gotID, err := uuid.ParseBytes(uuidBytes) // Use existing UUID @@ -295,7 +296,7 @@ func IdentifierFromDB(configStore types.GetterSetter) (string, error) { identifier = gotID.String() // Save new UUID - err = configStore.Set([]byte(uuidKey), []byte(identifier)) + err = configStore.Set(storage.KeyByIdentifier([]byte(uuidKey), storage.IdentifierTypeRegistration, []byte(registrationId)), []byte(identifier)) if err != nil { return "", fmt.Errorf("saving new UUID: %w", err) } @@ -304,8 +305,8 @@ func IdentifierFromDB(configStore types.GetterSetter) (string, error) { } // NodeKey returns the device node key from the storage layer -func NodeKey(getter types.Getter) (string, error) { - key, err := getter.Get([]byte(nodeKeyKey)) +func NodeKey(getter types.Getter, registrationId string) (string, error) { + key, err := getter.Get(storage.KeyByIdentifier([]byte(nodeKeyKey), storage.IdentifierTypeRegistration, []byte(registrationId))) if err != nil { return "", fmt.Errorf("error getting node key: %w", err) } @@ -317,8 +318,8 @@ func NodeKey(getter types.Getter) (string, error) { } // Config returns the device config from the storage layer -func Config(getter types.Getter) (string, error) { - key, err := getter.Get([]byte(configKey)) +func Config(getter types.Getter, registrationId string) (string, error) { + key, err := getter.Get(storage.KeyByIdentifier([]byte(configKey), storage.IdentifierTypeRegistration, []byte(registrationId))) if err != nil { return "", fmt.Errorf("error getting config key: %w", err) } @@ -366,7 +367,7 @@ func (e *Extension) Enroll(ctx context.Context) (string, bool, error) { } // Look up a node key cached in the local store - key, err := NodeKey(e.knapsack.ConfigStore()) + key, err := NodeKey(e.knapsack.ConfigStore(), e.registrationId) if err != nil { traces.SetError(span, fmt.Errorf("error reading node key from db: %w", err)) return "", false, fmt.Errorf("error reading node key from db: %w", err) @@ -459,7 +460,7 @@ func (e *Extension) Enroll(ctx context.Context) (string, bool, error) { } // Save newly acquired node key if successful - err = e.knapsack.ConfigStore().Set([]byte(nodeKeyKey), []byte(keyString)) + err = e.knapsack.ConfigStore().Set(storage.KeyByIdentifier([]byte(nodeKeyKey), storage.IdentifierTypeRegistration, []byte(e.registrationId)), []byte(keyString)) if err != nil { return "", true, fmt.Errorf("saving node key: %w", err) } @@ -485,7 +486,7 @@ func (e *Extension) RequireReenroll(ctx context.Context) { defer e.enrollMutex.Unlock() // Clear the node key such that reenrollment is required. e.NodeKey = "" - e.knapsack.ConfigStore().Delete([]byte(nodeKeyKey)) + e.knapsack.ConfigStore().Delete(storage.KeyByIdentifier([]byte(nodeKeyKey), storage.IdentifierTypeRegistration, []byte(e.registrationId))) } // GenerateConfigs will request the osquery configuration from the server. If @@ -501,7 +502,7 @@ func (e *Extension) GenerateConfigs(ctx context.Context) (map[string]string, err ) // Try to use cached config var confBytes []byte - confBytes, _ = e.knapsack.ConfigStore().Get([]byte(configKey)) + confBytes, _ = e.knapsack.ConfigStore().Get(storage.KeyByIdentifier([]byte(configKey), storage.IdentifierTypeRegistration, []byte(e.registrationId))) if len(confBytes) == 0 { if !e.enrolled() { @@ -513,7 +514,7 @@ func (e *Extension) GenerateConfigs(ctx context.Context) (map[string]string, err config = string(confBytes) } else { // Store good config - e.knapsack.ConfigStore().Set([]byte(configKey), []byte(config)) + e.knapsack.ConfigStore().Set(storage.KeyByIdentifier([]byte(configKey), storage.IdentifierTypeRegistration, []byte(e.registrationId)), []byte(config)) // open the start up settings writer just to trigger a write of the config, // then we can immediately close it diff --git a/pkg/osquery/extension_test.go b/pkg/osquery/extension_test.go index 51e5bb749..8e8cbc5d5 100644 --- a/pkg/osquery/extension_test.go +++ b/pkg/osquery/extension_test.go @@ -171,7 +171,7 @@ func TestExtensionEnrollTransportError(t *testing.T) { defer cleanup() k := makeKnapsack(t, db) - e, err := NewExtension(context.TODO(), m, k, ulid.New(), ExtensionOpts{}) + e, err := NewExtension(context.TODO(), m, k, types.DefaultRegistrationID, ExtensionOpts{}) require.Nil(t, err) key, invalid, err := e.Enroll(context.Background()) @@ -220,7 +220,7 @@ func TestExtensionEnroll(t *testing.T) { expectedEnrollSecret := "foo_secret" k.On("ReadEnrollSecret").Maybe().Return(expectedEnrollSecret, nil) - e, err := NewExtension(context.TODO(), m, k, ulid.New(), ExtensionOpts{}) + e, err := NewExtension(context.TODO(), m, k, types.DefaultRegistrationID, ExtensionOpts{}) require.Nil(t, err) key, invalid, err := e.Enroll(context.Background()) @@ -239,7 +239,7 @@ func TestExtensionEnroll(t *testing.T) { assert.Equal(t, expectedNodeKey, key) assert.Equal(t, expectedEnrollSecret, gotEnrollSecret) - e, err = NewExtension(context.TODO(), m, k, ulid.New(), ExtensionOpts{}) + e, err = NewExtension(context.TODO(), m, k, types.DefaultRegistrationID, ExtensionOpts{}) require.Nil(t, err) // Still should not re-enroll (because node key stored in DB) key, invalid, err = e.Enroll(context.Background()) @@ -273,7 +273,7 @@ func TestExtensionGenerateConfigsTransportError(t *testing.T) { defer cleanup() k := makeKnapsack(t, db) k.ConfigStore().Set([]byte(nodeKeyKey), []byte("some_node_key")) - e, err := NewExtension(context.TODO(), m, k, ulid.New(), ExtensionOpts{}) + e, err := NewExtension(context.TODO(), m, k, types.DefaultRegistrationID, ExtensionOpts{}) require.Nil(t, err) configs, err := e.GenerateConfigs(context.Background()) diff --git a/pkg/osquery/table/launcher_config.go b/pkg/osquery/table/launcher_config.go index b585dac1a..44b140d03 100644 --- a/pkg/osquery/table/launcher_config.go +++ b/pkg/osquery/table/launcher_config.go @@ -20,7 +20,7 @@ func generateLauncherConfig(store types.Getter, registrationTracker types.Regist return func(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) { results := make([]map[string]string, 0) for _, registrationId := range registrationTracker.RegistrationIDs() { - config, err := osquery.Config(store) + config, err := osquery.Config(store, registrationId) if err != nil { return nil, err } diff --git a/pkg/osquery/table/launcher_info.go b/pkg/osquery/table/launcher_info.go index e148d599b..74b150958 100644 --- a/pkg/osquery/table/launcher_info.go +++ b/pkg/osquery/table/launcher_info.go @@ -53,7 +53,7 @@ func LauncherInfoTable(configStore types.GetterSetter, LauncherHistoryStore type func generateLauncherInfoTable(configStore types.GetterSetter, LauncherHistoryStore types.GetterSetter) table.GenerateFunc { return func(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) { - identifier, err := osquery.IdentifierFromDB(configStore) + identifier, err := osquery.IdentifierFromDB(configStore, types.DefaultRegistrationID) if err != nil { return nil, err } From 9fd8070f302ad73e89684789c34b0b71fde6cc64 Mon Sep 17 00:00:00 2001 From: RebeccaMahany Date: Tue, 10 Dec 2024 11:02:18 -0500 Subject: [PATCH 2/3] Write data per-registration into startup settings store --- ee/agent/reset.go | 20 ++++++++--- ee/agent/startupsettings/writer.go | 48 ++++++++++++++----------- ee/agent/startupsettings/writer_test.go | 4 +++ ee/agent/storage/keys.go | 15 ++++++++ ee/agent/storage/keys_test.go | 39 ++++++++++++++++++++ ee/uninstall/uninstall_test.go | 1 + pkg/osquery/interactive/interactive.go | 7 ++-- pkg/osquery/runtime/osqueryinstance.go | 2 +- pkg/osquery/table/table.go | 24 ++++++++----- 9 files changed, 124 insertions(+), 36 deletions(-) diff --git a/ee/agent/reset.go b/ee/agent/reset.go index 4f698a683..f9c3ac343 100644 --- a/ee/agent/reset.go +++ b/ee/agent/reset.go @@ -9,9 +9,11 @@ import ( "fmt" "log/slog" "os" + "strings" "time" "github.com/golang-jwt/jwt/v5" + "github.com/kolide/launcher/ee/agent/storage" "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/pkg/osquery/runsimple" "github.com/kolide/launcher/pkg/traces" @@ -287,10 +289,20 @@ func currentMunemo(k types.Knapsack) (string, error) { // as a record of the current state of this database before reset. It appends this record // to previous records if they exist, and returns the collection ready for storage. func prepareDatabaseResetRecords(ctx context.Context, k types.Knapsack, resetReason string) ([]byte, error) { // nolint:unused - nodeKey, err := k.ConfigStore().Get([]byte("nodeKey")) - if err != nil { - k.Slogger().Log(ctx, slog.LevelWarn, "could not get node key from store", "err", err) + nodeKeys := make([]string, 0) + for _, registrationId := range k.RegistrationIDs() { + nodeKey, err := k.ConfigStore().Get(storage.KeyByIdentifier([]byte("nodeKey"), storage.IdentifierTypeRegistration, []byte(registrationId))) + if err != nil { + k.Slogger().Log(ctx, slog.LevelWarn, + "could not get node key from store", + "registration_id", registrationId, + "err", err, + ) + continue + } + nodeKeys = append(nodeKeys, string(nodeKey)) } + nodeKey := strings.Join(nodeKeys, ",") localPubKey, err := getLocalPubKey(k) if err != nil { @@ -328,7 +340,7 @@ func prepareDatabaseResetRecords(ctx context.Context, k types.Knapsack, resetRea } dataToStore := dbResetRecord{ - NodeKey: string(nodeKey), + NodeKey: nodeKey, PubKeys: [][]byte{localPubKey}, Serial: string(serial), HardwareUUID: string(hardwareUuid), diff --git a/ee/agent/startupsettings/writer.go b/ee/agent/startupsettings/writer.go index e70b9e453..2e9cc7f71 100644 --- a/ee/agent/startupsettings/writer.go +++ b/ee/agent/startupsettings/writer.go @@ -10,6 +10,7 @@ import ( "log/slog" "github.com/kolide/launcher/ee/agent/flags/keys" + "github.com/kolide/launcher/ee/agent/storage" agentsqlite "github.com/kolide/launcher/ee/agent/storage/sqlite" "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/pkg/traces" @@ -70,23 +71,27 @@ func (s *startupSettingsWriter) WriteSettings() error { } updatedFlags["use_tuf_autoupdater"] = "enabled" // Hardcode for backwards compatibility circa v1.5.3 - atcConfig, err := s.extractAutoTableConstructionConfig() - if err != nil { - s.knapsack.Slogger().Log(context.TODO(), slog.LevelDebug, - "extracting auto_table_construction config", - "err", err, - ) - } else { - updatedFlags["auto_table_construction"] = atcConfig - } + for _, registrationId := range s.knapsack.RegistrationIDs() { + atcConfig, err := s.extractAutoTableConstructionConfig(registrationId) + if err != nil { + s.knapsack.Slogger().Log(context.TODO(), slog.LevelDebug, + "extracting auto_table_construction config", + "err", err, + ) + } else { + atcConfigKey := storage.KeyByIdentifier([]byte("auto_table_construction"), storage.IdentifierTypeRegistration, []byte(registrationId)) + updatedFlags[string(atcConfigKey)] = atcConfig + } - if katcConfig, err := s.extractKATCConstructionConfig(); err != nil { - s.knapsack.Slogger().Log(context.TODO(), slog.LevelDebug, - "extracting katc_config", - "err", err, - ) - } else { - updatedFlags["katc_config"] = katcConfig + if katcConfig, err := s.extractKATCConstructionConfig(registrationId); err != nil { + s.knapsack.Slogger().Log(context.TODO(), slog.LevelDebug, + "extracting katc_config", + "err", err, + ) + } else { + katcConfigKey := storage.KeyByIdentifier([]byte("katc_config"), storage.IdentifierTypeRegistration, []byte(registrationId)) + updatedFlags[string(katcConfigKey)] = katcConfig + } } if _, err := s.kvStore.Update(updatedFlags); err != nil { @@ -112,8 +117,8 @@ func (s *startupSettingsWriter) Close() error { return s.kvStore.Close() } -func (s *startupSettingsWriter) extractAutoTableConstructionConfig() (string, error) { - osqConfig, err := s.knapsack.ConfigStore().Get([]byte("config")) +func (s *startupSettingsWriter) extractAutoTableConstructionConfig(registrationId string) (string, error) { + osqConfig, err := s.knapsack.ConfigStore().Get(storage.KeyByIdentifier([]byte("config"), storage.IdentifierTypeRegistration, []byte(registrationId))) if err != nil { return "", fmt.Errorf("could not get osquery config from store: %w", err) } @@ -140,10 +145,13 @@ func (s *startupSettingsWriter) extractAutoTableConstructionConfig() (string, er return string(atcJson), nil } -func (s *startupSettingsWriter) extractKATCConstructionConfig() (string, error) { +func (s *startupSettingsWriter) extractKATCConstructionConfig(registrationId string) (string, error) { kolideCfg := make(map[string]string) if err := s.knapsack.KatcConfigStore().ForEach(func(k []byte, v []byte) error { - kolideCfg[string(k)] = string(v) + key, identifier := storage.SplitKey(k) + if string(identifier) == registrationId { + kolideCfg[string(key)] = string(v) + } return nil }); err != nil { return "", fmt.Errorf("could not get Kolide ATC config from store: %w", err) diff --git a/ee/agent/startupsettings/writer_test.go b/ee/agent/startupsettings/writer_test.go index a46e5ae2a..1554d4271 100644 --- a/ee/agent/startupsettings/writer_test.go +++ b/ee/agent/startupsettings/writer_test.go @@ -11,6 +11,7 @@ import ( "github.com/kolide/launcher/ee/agent/flags/keys" "github.com/kolide/launcher/ee/agent/storage/inmemory" agentsqlite "github.com/kolide/launcher/ee/agent/storage/sqlite" + "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/mock" @@ -35,6 +36,7 @@ func TestOpenWriter_NewDatabase(t *testing.T) { k.On("ConfigStore").Return(inmemory.NewStore()) k.On("Slogger").Return(multislogger.NewNopLogger()) k.On("KatcConfigStore").Return(inmemory.NewStore()) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) // Set up storage db, which should create the database and set all flags s, err := OpenWriter(context.TODO(), k) @@ -87,6 +89,7 @@ func TestOpenWriter_DatabaseAlreadyExists(t *testing.T) { k.On("RegisterChangeObserver", mock.Anything, keys.UpdateChannel) k.On("RegisterChangeObserver", mock.Anything, keys.PinnedLauncherVersion) k.On("RegisterChangeObserver", mock.Anything, keys.PinnedOsquerydVersion) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) // Set up flag updateChannelVal := "alpha" @@ -132,6 +135,7 @@ func TestFlagsChanged(t *testing.T) { k.On("RegisterChangeObserver", mock.Anything, keys.UpdateChannel) k.On("RegisterChangeObserver", mock.Anything, keys.PinnedLauncherVersion) k.On("RegisterChangeObserver", mock.Anything, keys.PinnedOsquerydVersion) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) updateChannelVal := "beta" k.On("UpdateChannel").Return(updateChannelVal).Once() pinnedLauncherVersion := "1.2.3" diff --git a/ee/agent/storage/keys.go b/ee/agent/storage/keys.go index 8e029ca33..a780a28db 100644 --- a/ee/agent/storage/keys.go +++ b/ee/agent/storage/keys.go @@ -35,3 +35,18 @@ func KeyByIdentifier(key []byte, identifierType []byte, identifier []byte) []byt return newKey } + +func SplitKey(key []byte) ([]byte, []byte) { + if !bytes.Contains(key, []byte{keyDelimiter}) { + return key, defaultIdentifier + } + + // Key takes the form `::` -- split + // on the keyDelimiter. + parts := bytes.SplitN(key, []byte{keyDelimiter}, 3) + if len(parts) != 3 { + return key, defaultIdentifier + } + + return parts[0], parts[2] +} diff --git a/ee/agent/storage/keys_test.go b/ee/agent/storage/keys_test.go index 0d5d68b29..0209bb37a 100644 --- a/ee/agent/storage/keys_test.go +++ b/ee/agent/storage/keys_test.go @@ -46,3 +46,42 @@ func TestKeyByIdentifier(t *testing.T) { }) } } + +func TestSplitKey(t *testing.T) { + t.Parallel() + + for _, tt := range []struct { + testCaseName string + key []byte + expectedKey []byte + expectedIdentifier []byte + }{ + { + testCaseName: "default node key", + key: []byte("nodeKey"), + expectedKey: []byte("nodeKey"), + expectedIdentifier: []byte("default"), + }, + { + testCaseName: "uuid by registration", + key: []byte("uuid:registration:some-test-registration-id"), + expectedKey: []byte("uuid"), + expectedIdentifier: []byte("some-test-registration-id"), + }, + { + testCaseName: "katc table by registration", + key: []byte("katc_some_test_table:registration:another-test-registration-id"), + expectedKey: []byte("katc_some_test_table"), + expectedIdentifier: []byte("another-test-registration-id"), + }, + } { + tt := tt + t.Run(tt.testCaseName, func(t *testing.T) { + t.Parallel() + + splitKey, identifier := SplitKey(tt.key) + require.Equal(t, tt.expectedKey, splitKey) + require.Equal(t, tt.expectedIdentifier, identifier) + }) + } +} diff --git a/ee/uninstall/uninstall_test.go b/ee/uninstall/uninstall_test.go index 2216732de..556e8ebd9 100644 --- a/ee/uninstall/uninstall_test.go +++ b/ee/uninstall/uninstall_test.go @@ -60,6 +60,7 @@ func TestUninstall(t *testing.T) { k.On("EnrollSecretPath").Return(enrollSecretPath) k.On("Slogger").Return(multislogger.NewNopLogger()) k.On("RootDirectory").Return(tempRootDir) + k.On("RegistrationIDs").Return([]string{types.DefaultRegistrationID}) testConfigStore, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.ConfigStore.String()) require.NoError(t, err, "could not create test config store") k.On("ConfigStore").Return(testConfigStore) diff --git a/pkg/osquery/interactive/interactive.go b/pkg/osquery/interactive/interactive.go index 660eabad8..1309f3b38 100644 --- a/pkg/osquery/interactive/interactive.go +++ b/pkg/osquery/interactive/interactive.go @@ -13,6 +13,7 @@ import ( "github.com/kolide/kit/fsutil" "github.com/kolide/kit/ulid" "github.com/kolide/launcher/ee/agent/startupsettings" + "github.com/kolide/launcher/ee/agent/storage" "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/pkg/augeas" osqueryRuntime "github.com/kolide/launcher/pkg/osquery/runtime" @@ -60,7 +61,7 @@ func StartProcess(knapsack types.Knapsack, interactiveRootDir string) (*os.Proce } // start building list of osq plugins with the kolide tables - osqPlugins := table.PlatformTables(knapsack, knapsack.Slogger(), knapsack.OsquerydPath()) + osqPlugins := table.PlatformTables(knapsack, types.DefaultRegistrationID, knapsack.Slogger(), knapsack.OsquerydPath()) osqueryFlags := knapsack.OsqueryFlags() // if we were not provided a config path flag, try to add default config @@ -200,7 +201,9 @@ func generateConfigPlugin(launcherDaemonRootDir string) (*config.Plugin, error) } defer r.Close() - atcConfig, err := r.Get("auto_table_construction") + // Use the default registration's config + atcConfigKey := storage.KeyByIdentifier([]byte("auto_table_construction"), storage.IdentifierTypeRegistration, []byte(types.DefaultRegistrationID)) + atcConfig, err := r.Get(string(atcConfigKey)) if err != nil { return nil, fmt.Errorf("error getting auto_table_construction from startup settings: %w", err) } diff --git a/pkg/osquery/runtime/osqueryinstance.go b/pkg/osquery/runtime/osqueryinstance.go index 824af3703..994b92b02 100644 --- a/pkg/osquery/runtime/osqueryinstance.go +++ b/pkg/osquery/runtime/osqueryinstance.go @@ -430,7 +430,7 @@ func (i *OsqueryInstance) Launch() error { distributed.NewPlugin(KolideSaasExtensionName, i.saasExtension.GetQueries, i.saasExtension.WriteResults), osquerylogger.NewPlugin(KolideSaasExtensionName, i.saasExtension.LogString), } - kolideSaasPlugins = append(kolideSaasPlugins, table.PlatformTables(i.knapsack, i.knapsack.Slogger().With("component", "platform_tables"), currentOsquerydBinaryPath)...) + kolideSaasPlugins = append(kolideSaasPlugins, table.PlatformTables(i.knapsack, i.registrationId, i.knapsack.Slogger().With("component", "platform_tables"), currentOsquerydBinaryPath)...) kolideSaasPlugins = append(kolideSaasPlugins, table.LauncherTables(i.knapsack)...) if err := i.StartOsqueryExtensionManagerServer(KolideSaasExtensionName, paths.extensionSocketPath, i.extensionManagerClient, kolideSaasPlugins); err != nil { diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go index d357e7f3b..2bd17b609 100644 --- a/pkg/osquery/table/table.go +++ b/pkg/osquery/table/table.go @@ -8,6 +8,7 @@ import ( "log/slog" "github.com/kolide/launcher/ee/agent/startupsettings" + "github.com/kolide/launcher/ee/agent/storage" "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/ee/allowedcmd" "github.com/kolide/launcher/ee/katc" @@ -43,7 +44,7 @@ func LauncherTables(k types.Knapsack) []osquery.OsqueryPlugin { } // PlatformTables returns all tables for the launcher build platform. -func PlatformTables(k types.Knapsack, slogger *slog.Logger, currentOsquerydBinaryPath string) []osquery.OsqueryPlugin { +func PlatformTables(k types.Knapsack, registrationId string, slogger *slog.Logger, currentOsquerydBinaryPath string) []osquery.OsqueryPlugin { // Common tables to all platforms tables := []osquery.OsqueryPlugin{ ChromeLoginDataEmails(slogger), @@ -72,23 +73,23 @@ func PlatformTables(k types.Knapsack, slogger *slog.Logger, currentOsquerydBinar tables = append(tables, platformSpecificTables(slogger, currentOsquerydBinaryPath)...) // Add in the Kolide custom ATC tables - tables = append(tables, kolideCustomAtcTables(k, slogger)...) + tables = append(tables, kolideCustomAtcTables(k, registrationId, slogger)...) return tables } // kolideCustomAtcTables retrieves Kolide ATC config from the appropriate data store(s), // then constructs the tables. -func kolideCustomAtcTables(k types.Knapsack, slogger *slog.Logger) []osquery.OsqueryPlugin { +func kolideCustomAtcTables(k types.Knapsack, registrationId string, slogger *slog.Logger) []osquery.OsqueryPlugin { // Fetch tables from KVStore or from startup settings - config, err := katcFromDb(k) + config, err := katcFromDb(k, registrationId) if err != nil { slogger.Log(context.TODO(), slog.LevelDebug, "could not retrieve KATC config from store, may not have access -- falling back to startup settings", "err", err, ) - config, err = katcFromStartupSettings(k) + config, err = katcFromStartupSettings(k, registrationId) if err != nil { slogger.Log(context.TODO(), slog.LevelWarn, "could not retrieve KATC config from startup settings", @@ -101,13 +102,17 @@ func kolideCustomAtcTables(k types.Knapsack, slogger *slog.Logger) []osquery.Osq return katc.ConstructKATCTables(config, slogger) } -func katcFromDb(k types.Knapsack) (map[string]string, error) { +func katcFromDb(k types.Knapsack, registrationId string) (map[string]string, error) { if k == nil || k.KatcConfigStore() == nil { return nil, errors.New("stores in knapsack not available") } katcCfg := make(map[string]string) if err := k.KatcConfigStore().ForEach(func(k []byte, v []byte) error { - katcCfg[string(k)] = string(v) + key, identifier := storage.SplitKey(k) + if string(identifier) == registrationId { + katcCfg[string(key)] = string(v) + } + return nil }); err != nil { return nil, fmt.Errorf("retrieving contents of Kolide ATC config store: %w", err) @@ -116,14 +121,15 @@ func katcFromDb(k types.Knapsack) (map[string]string, error) { return katcCfg, nil } -func katcFromStartupSettings(k types.Knapsack) (map[string]string, error) { +func katcFromStartupSettings(k types.Knapsack, registrationId string) (map[string]string, error) { r, err := startupsettings.OpenReader(context.TODO(), k.RootDirectory()) if err != nil { return nil, fmt.Errorf("error opening startup settings reader: %w", err) } defer r.Close() - katcConfig, err := r.Get("katc_config") + katcConfigKey := storage.KeyByIdentifier([]byte("katc_config"), storage.IdentifierTypeRegistration, []byte(registrationId)) + katcConfig, err := r.Get(string(katcConfigKey)) if err != nil { return nil, fmt.Errorf("error getting katc_config from startup settings: %w", err) } From ffb48cf2c03a6729a2782230a12e15184639df8e Mon Sep 17 00:00:00 2001 From: RebeccaMahany Date: Wed, 11 Dec 2024 15:58:56 -0500 Subject: [PATCH 3/3] Return identifier type from SplitKey too --- ee/agent/startupsettings/writer.go | 2 +- ee/agent/storage/keys.go | 8 +++--- ee/agent/storage/keys_test.go | 39 +++++++++++++++++------------- pkg/osquery/table/table.go | 2 +- 4 files changed, 28 insertions(+), 23 deletions(-) diff --git a/ee/agent/startupsettings/writer.go b/ee/agent/startupsettings/writer.go index 2e9cc7f71..d489228ad 100644 --- a/ee/agent/startupsettings/writer.go +++ b/ee/agent/startupsettings/writer.go @@ -148,7 +148,7 @@ func (s *startupSettingsWriter) extractAutoTableConstructionConfig(registrationI func (s *startupSettingsWriter) extractKATCConstructionConfig(registrationId string) (string, error) { kolideCfg := make(map[string]string) if err := s.knapsack.KatcConfigStore().ForEach(func(k []byte, v []byte) error { - key, identifier := storage.SplitKey(k) + key, _, identifier := storage.SplitKey(k) if string(identifier) == registrationId { kolideCfg[string(key)] = string(v) } diff --git a/ee/agent/storage/keys.go b/ee/agent/storage/keys.go index a780a28db..84ea26b42 100644 --- a/ee/agent/storage/keys.go +++ b/ee/agent/storage/keys.go @@ -36,17 +36,17 @@ func KeyByIdentifier(key []byte, identifierType []byte, identifier []byte) []byt return newKey } -func SplitKey(key []byte) ([]byte, []byte) { +func SplitKey(key []byte) ([]byte, []byte, []byte) { if !bytes.Contains(key, []byte{keyDelimiter}) { - return key, defaultIdentifier + return key, nil, defaultIdentifier } // Key takes the form `::` -- split // on the keyDelimiter. parts := bytes.SplitN(key, []byte{keyDelimiter}, 3) if len(parts) != 3 { - return key, defaultIdentifier + return key, nil, defaultIdentifier } - return parts[0], parts[2] + return parts[0], parts[1], parts[2] } diff --git a/ee/agent/storage/keys_test.go b/ee/agent/storage/keys_test.go index 0209bb37a..4e4b84fb7 100644 --- a/ee/agent/storage/keys_test.go +++ b/ee/agent/storage/keys_test.go @@ -51,36 +51,41 @@ func TestSplitKey(t *testing.T) { t.Parallel() for _, tt := range []struct { - testCaseName string - key []byte - expectedKey []byte - expectedIdentifier []byte + testCaseName string + key []byte + expectedKey []byte + expectedIdentifierType []byte + expectedIdentifier []byte }{ { - testCaseName: "default node key", - key: []byte("nodeKey"), - expectedKey: []byte("nodeKey"), - expectedIdentifier: []byte("default"), + testCaseName: "default node key", + key: []byte("nodeKey"), + expectedKey: []byte("nodeKey"), + expectedIdentifierType: nil, + expectedIdentifier: []byte("default"), }, { - testCaseName: "uuid by registration", - key: []byte("uuid:registration:some-test-registration-id"), - expectedKey: []byte("uuid"), - expectedIdentifier: []byte("some-test-registration-id"), + testCaseName: "uuid by registration", + key: []byte("uuid:registration:some-test-registration-id"), + expectedKey: []byte("uuid"), + expectedIdentifierType: IdentifierTypeRegistration, + expectedIdentifier: []byte("some-test-registration-id"), }, { - testCaseName: "katc table by registration", - key: []byte("katc_some_test_table:registration:another-test-registration-id"), - expectedKey: []byte("katc_some_test_table"), - expectedIdentifier: []byte("another-test-registration-id"), + testCaseName: "katc table by registration", + key: []byte("katc_some_test_table:registration:another-test-registration-id"), + expectedKey: []byte("katc_some_test_table"), + expectedIdentifierType: IdentifierTypeRegistration, + expectedIdentifier: []byte("another-test-registration-id"), }, } { tt := tt t.Run(tt.testCaseName, func(t *testing.T) { t.Parallel() - splitKey, identifier := SplitKey(tt.key) + splitKey, identifierType, identifier := SplitKey(tt.key) require.Equal(t, tt.expectedKey, splitKey) + require.Equal(t, tt.expectedIdentifierType, identifierType) require.Equal(t, tt.expectedIdentifier, identifier) }) } diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go index 2bd17b609..e02ede881 100644 --- a/pkg/osquery/table/table.go +++ b/pkg/osquery/table/table.go @@ -108,7 +108,7 @@ func katcFromDb(k types.Knapsack, registrationId string) (map[string]string, err } katcCfg := make(map[string]string) if err := k.KatcConfigStore().ForEach(func(k []byte, v []byte) error { - key, identifier := storage.SplitKey(k) + key, _, identifier := storage.SplitKey(k) if string(identifier) == registrationId { katcCfg[string(key)] = string(v) }