Store configuration for Kolide custom ATC tables (#1761)

kolide · Jun 26, 2024 · 05ea065 · 05ea065
1 parent 25f1513
commit 05ea065
Show file tree

Hide file tree

Showing 15 changed files with 203 additions and 12 deletions.
diff --git a/cmd/launcher/launcher.go b/cmd/launcher/launcher.go
@@ -71,6 +71,7 @@ const (
 	serverDataSubsystemName  = "kolide_server_data"
 	desktopMenuSubsystemName = "kolide_desktop_menu"
 	authTokensSubsystemName  = "auth_tokens"
+	katcSubsystemName        = "katc_config" // Kolide ATC
 )
 
 // runLauncher is the entry point into running launcher. It creates a
@@ -242,13 +243,13 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl
 		startupSpan.AddEvent("log_shipper_init_completed")
 	}
 
-	s, err := startupsettings.OpenWriter(ctx, k)
+	startupSettingsWriter, err := startupsettings.OpenWriter(ctx, k)
 	if err != nil {
 		return fmt.Errorf("creating startup db: %w", err)
 	}
-	defer s.Close()
+	defer startupSettingsWriter.Close()
 
-	if err := s.WriteSettings(); err != nil {
+	if err := startupSettingsWriter.WriteSettings(); err != nil {
 		slogger.Log(ctx, slog.LevelError,
 			"writing startup settings",
 			"err", err,
@@ -400,6 +401,10 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl
 		controlService.RegisterConsumer(serverDataSubsystemName, keyvalueconsumer.New(k.ServerProvidedDataStore()))
 		// agentFlagConsumer handles agent flags pushed from the control server
 		controlService.RegisterConsumer(agentFlagsSubsystemName, keyvalueconsumer.New(flagController))
+		// katcConfigConsumer handles updates to Kolide's custom ATC tables
+		controlService.RegisterConsumer(katcSubsystemName, keyvalueconsumer.New(k.KatcConfigStore()))
+		controlService.RegisterSubscriber(katcSubsystemName, osqueryRunner)
+		controlService.RegisterSubscriber(katcSubsystemName, startupSettingsWriter)
 
 		runner, err = desktopRunner.New(
 			k,

diff --git a/ee/agent/knapsack/knapsack.go b/ee/agent/knapsack/knapsack.go
@@ -87,6 +87,10 @@ func (k *knapsack) AgentFlagsStore() types.KVStore {
 	return k.getKVStore(storage.AgentFlagsStore)
 }
 
+func (k *knapsack) KatcConfigStore() types.KVStore {
+	return k.getKVStore(storage.KatcConfigStore)
+}
+
 func (k *knapsack) AutoupdateErrorsStore() types.KVStore {
 	return k.getKVStore(storage.AutoupdateErrorsStore)
 }

diff --git a/ee/agent/startupsettings/writer.go b/ee/agent/startupsettings/writer.go
@@ -51,6 +51,17 @@ func OpenWriter(ctx context.Context, knapsack types.Knapsack) (*startupSettingsW
 	return s, nil
 }
 
+// Ping satisfies the control.subscriber interface -- the runner subscribes to changes to
+// the katc_config subsystem.
+func (s *startupSettingsWriter) Ping() {
+	if err := s.WriteSettings(); err != nil {
+		s.knapsack.Slogger().Log(context.TODO(), slog.LevelWarn,
+			"could not write updated settings",
+			"err", err,
+		)
+	}
+}
+
 // WriteSettings updates the flags with their values from the agent flag data store.
 func (s *startupSettingsWriter) WriteSettings() error {
 	updatedFlags := make(map[string]string)
@@ -69,6 +80,15 @@ func (s *startupSettingsWriter) WriteSettings() error {
 		updatedFlags["auto_table_construction"] = 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 _, err := s.kvStore.Update(updatedFlags); err != nil {
 		return fmt.Errorf("writing settings: %w", err)
 	}
@@ -119,3 +139,20 @@ func (s *startupSettingsWriter) extractAutoTableConstructionConfig() (string, er
 
 	return string(atcJson), nil
 }
+
+func (s *startupSettingsWriter) extractKATCConstructionConfig() (string, error) {
+	kolideCfg := make(map[string]string)
+	if err := s.knapsack.KatcConfigStore().ForEach(func(k []byte, v []byte) error {
+		kolideCfg[string(k)] = string(v)
+		return nil
+	}); err != nil {
+		return "", fmt.Errorf("could not get Kolide ATC config from store: %w", err)
+	}
+
+	atcJson, err := json.Marshal(kolideCfg)
+	if err != nil {
+		return "", fmt.Errorf("could not marshal katc_config: %w", err)
+	}
+
+	return string(atcJson), nil
+}
diff --git a/ee/agent/startupsettings/writer_test.go b/ee/agent/startupsettings/writer_test.go
@@ -34,6 +34,7 @@ func TestOpenWriter_NewDatabase(t *testing.T) {
 	k.On("PinnedOsquerydVersion").Return("5.11.0")
 	k.On("ConfigStore").Return(inmemory.NewStore())
 	k.On("Slogger").Return(multislogger.NewNopLogger())
+	k.On("KatcConfigStore").Return(inmemory.NewStore())
 
 	// Set up storage db, which should create the database and set all flags
 	s, err := OpenWriter(context.TODO(), k)
@@ -82,6 +83,7 @@ func TestOpenWriter_DatabaseAlreadyExists(t *testing.T) {
 	// Set up dependencies
 	k := typesmocks.NewKnapsack(t)
 	k.On("RootDirectory").Return(testRootDir)
+	k.On("KatcConfigStore").Return(inmemory.NewStore())
 	k.On("RegisterChangeObserver", mock.Anything, keys.UpdateChannel)
 	k.On("RegisterChangeObserver", mock.Anything, keys.PinnedLauncherVersion)
 	k.On("RegisterChangeObserver", mock.Anything, keys.PinnedOsquerydVersion)
@@ -126,6 +128,7 @@ func TestFlagsChanged(t *testing.T) {
 	testRootDir := t.TempDir()
 	k := typesmocks.NewKnapsack(t)
 	k.On("RootDirectory").Return(testRootDir)
+	k.On("KatcConfigStore").Return(inmemory.NewStore())
 	k.On("RegisterChangeObserver", mock.Anything, keys.UpdateChannel)
 	k.On("RegisterChangeObserver", mock.Anything, keys.PinnedLauncherVersion)
 	k.On("RegisterChangeObserver", mock.Anything, keys.PinnedOsquerydVersion)

diff --git a/ee/agent/storage/bbolt/stores_bbolt.go b/ee/agent/storage/bbolt/stores_bbolt.go
@@ -20,6 +20,7 @@ func MakeStores(ctx context.Context, slogger *slog.Logger, db *bbolt.DB) (map[st
 
 	var storeNames = []storage.Store{
 		storage.AgentFlagsStore,
+		storage.KatcConfigStore,
 		storage.AutoupdateErrorsStore,
 		storage.ConfigStore,
 		storage.ControlStore,

diff --git a/ee/agent/storage/ci/stores_ci.go b/ee/agent/storage/ci/stores_ci.go
@@ -18,6 +18,7 @@ import (
 func MakeStores(t *testing.T, slogger *slog.Logger, db *bbolt.DB) (map[storage.Store]types.KVStore, error) {
 	var storeNames = []storage.Store{
 		storage.AgentFlagsStore,
+		storage.KatcConfigStore,
 		storage.AutoupdateErrorsStore,
 		storage.ConfigStore,
 		storage.ControlStore,

diff --git a/ee/agent/storage/stores.go b/ee/agent/storage/stores.go
@@ -5,6 +5,7 @@ type Store string
 
 const (
 	AgentFlagsStore             Store = "agent_flags"              // The store used for agent control flags.
+	KatcConfigStore             Store = "katc_config"              // The store used for Kolide custom ATC configuration
 	AutoupdateErrorsStore       Store = "tuf_autoupdate_errors"    // The store used for tracking new autoupdater errors.
 	ConfigStore                 Store = "config"                   // The store used for launcher configuration.
 	ControlStore                Store = "control_service_data"     // The store used for control service caching data.

diff --git a/ee/agent/types/mocks/knapsack.go b/ee/agent/types/mocks/knapsack.go
diff --git a/ee/agent/types/stores.go b/ee/agent/types/stores.go
@@ -5,6 +5,7 @@ import "github.com/kolide/launcher/ee/agent/storage"
 type Stores interface {
 	Stores() map[storage.Store]KVStore
 	AgentFlagsStore() KVStore
+	KatcConfigStore() KVStore
 	AutoupdateErrorsStore() KVStore
 	ConfigStore() KVStore
 	ControlStore() KVStore

diff --git a/pkg/osquery/interactive/interactive.go b/pkg/osquery/interactive/interactive.go
@@ -56,7 +56,7 @@ func StartProcess(knapsack types.Knapsack, interactiveRootDir string) (*os.Proce
 	}
 
 	// start building list of osq plugins with the kolide tables
-	osqPlugins := table.PlatformTables(knapsack.Slogger(), knapsack.OsquerydPath())
+	osqPlugins := table.PlatformTables(knapsack, knapsack.Slogger(), knapsack.OsquerydPath())
 
 	osqueryFlags := knapsack.OsqueryFlags()
 	// if we were not provided a config path flag, try to add default config

diff --git a/pkg/osquery/interactive/interactive_test.go b/pkg/osquery/interactive/interactive_test.go
@@ -16,6 +16,8 @@ import (
 
 	"github.com/kolide/kit/fsutil"
 	"github.com/kolide/kit/ulid"
+	"github.com/kolide/launcher/ee/agent/storage"
+	storageci "github.com/kolide/launcher/ee/agent/storage/ci"
 	"github.com/kolide/launcher/ee/agent/types/mocks"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/packaging"
@@ -97,6 +99,9 @@ func TestProc(t *testing.T) {
 			mockSack.On("OsqueryFlags").Return(tt.osqueryFlags)
 			mockSack.On("Slogger").Return(multislogger.NewNopLogger())
 			mockSack.On("RootDirectory").Maybe().Return("whatever_the_root_launcher_dir_is")
+			store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+			require.NoError(t, err)
+			mockSack.On("KatcConfigStore").Return(store)
 
 			proc, _, err := StartProcess(mockSack, rootDir)
 

diff --git a/pkg/osquery/runtime/runner.go b/pkg/osquery/runtime/runner.go
@@ -183,6 +183,21 @@ func (r *Runner) FlagsChanged(flagKeys ...keys.FlagKey) {
 	}
 }
 
+// Ping satisfies the control.subscriber interface -- the runner subscribes to changes to
+// the katc_config subsystem.
+func (r *Runner) Ping() {
+	r.slogger.Log(context.TODO(), slog.LevelDebug,
+		"Kolide ATC configuration changed, restarting instance to apply",
+	)
+
+	if err := r.Restart(); err != nil {
+		r.slogger.Log(context.TODO(), slog.LevelError,
+			"could not restart osquery instance after Kolide ATC configuration changed",
+			"err", err,
+		)
+	}
+}
+
 // Restart allows you to cleanly shutdown the current instance and launch a new
 // instance with the same configurations.
 func (r *Runner) Restart() error {
@@ -469,7 +484,7 @@ func (r *Runner) launchOsqueryInstance() error {
 		)
 	}
 
-	// Now spawn an extension manage to for the tables. We need to
+	// Now spawn an extension manager for the tables. We need to
 	// start this one in the background, because the runner.Start
 	// function needs to return promptly enough for osquery to use
 	// it to enroll. Very racy
@@ -482,7 +497,7 @@ func (r *Runner) launchOsqueryInstance() error {
 			"errgroup", "kolide extension manager server launch",
 		)
 
-		plugins := table.PlatformTables(r.knapsack.Slogger().With("component", "platform_tables"), currentOsquerydBinaryPath)
+		plugins := table.PlatformTables(r.knapsack, r.knapsack.Slogger().With("component", "platform_tables"), currentOsquerydBinaryPath)
 
 		if len(plugins) == 0 {
 			return nil

diff --git a/pkg/osquery/runtime/runtime_posix_test.go b/pkg/osquery/runtime/runtime_posix_test.go
@@ -12,6 +12,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/kolide/launcher/ee/agent/storage"
+	storageci "github.com/kolide/launcher/ee/agent/storage/ci"
 	typesMocks "github.com/kolide/launcher/ee/agent/types/mocks"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/threadsafebuffer"
@@ -51,6 +53,9 @@ func TestOsquerySlowStart(t *testing.T) {
 	slogger := multislogger.New(slog.NewJSONHandler(&logBytes, &slog.HandlerOptions{Level: slog.LevelDebug}))
 	k.On("Slogger").Return(slogger.Logger)
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner := New(
 		k,
@@ -95,6 +100,9 @@ func TestExtensionSocketPath(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	extensionSocketPath := filepath.Join(rootDirectory, "sock")
 	runner := New(

diff --git a/pkg/osquery/runtime/runtime_test.go b/pkg/osquery/runtime/runtime_test.go
@@ -336,6 +336,9 @@ func TestWithOsqueryFlags(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner := New(
 		k,
@@ -368,6 +371,9 @@ func TestFlagsChanged(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	// Start the runner
 	runner := New(
@@ -461,6 +467,9 @@ func TestSimplePath(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner := New(
 		k,
@@ -489,6 +498,9 @@ func TestMultipleShutdowns(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner := New(
 		k,
@@ -516,6 +528,9 @@ func TestOsqueryDies(t *testing.T) {
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner := New(
 		k,
@@ -610,6 +625,9 @@ func setupOsqueryInstanceForTests(t *testing.T) (runner *Runner, teardown func()
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Maybe()
 	k.On("Slogger").Return(multislogger.NewNopLogger())
 	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	store, err := storageci.NewStore(t, multislogger.NewNopLogger(), storage.KatcConfigStore.String())
+	require.NoError(t, err)
+	k.On("KatcConfigStore").Return(store)
 
 	runner = New(
 		k,