From 8419858eacf5d29caf54d839903631a5f0161b7f Mon Sep 17 00:00:00 2001 From: zackattack01 Date: Thu, 21 Mar 2024 14:21:16 -0400 Subject: [PATCH 1/4] wire types.StorageStatTracker into knapsack, update direct bboltdb uses --- cmd/launcher/launcher.go | 7 +- ee/agent/knapsack/knapsack.go | 28 +++---- ee/agent/storage/bbolt/stats_bbolt.go | 112 ++++++++++++++++++++++++++ ee/agent/types/keyvalue_store.go | 10 +++ ee/agent/types/knapsack.go | 2 +- ee/debug/checkups/bboltdb.go | 17 ++-- ee/debug/checkups/host.go | 10 +-- pkg/osquery/table/launcher_db_info.go | 17 ++-- pkg/osquery/table/table.go | 2 +- 9 files changed, 162 insertions(+), 43 deletions(-) create mode 100644 ee/agent/storage/bbolt/stats_bbolt.go diff --git a/cmd/launcher/launcher.go b/cmd/launcher/launcher.go index 93d46309d..87372b8ae 100644 --- a/cmd/launcher/launcher.go +++ b/cmd/launcher/launcher.go @@ -185,9 +185,14 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl return fmt.Errorf("failed to create stores: %w", err) } + statTracker, err := agentbbolt.NewStatStore(slogger, db) + if err != nil { + return fmt.Errorf("failed to create stat tracker: %w", err) + } + fcOpts := []flags.Option{flags.WithCmdLineOpts(opts)} flagController := flags.NewFlagController(slogger, stores[storage.AgentFlagsStore], fcOpts...) - k := knapsack.New(stores, flagController, db, multiSlogger, systemMultiSlogger) + k := knapsack.New(stores, flagController, statTracker, multiSlogger, systemMultiSlogger) go runOsqueryVersionCheck(ctx, slogger, k.LatestOsquerydPath(ctx)) go timemachine.AddExclusions(ctx, k) diff --git a/ee/agent/knapsack/knapsack.go b/ee/agent/knapsack/knapsack.go index 9d5896df4..65532f2ce 100644 --- a/ee/agent/knapsack/knapsack.go +++ b/ee/agent/knapsack/knapsack.go @@ -14,7 +14,6 @@ import ( "github.com/kolide/launcher/ee/tuf" "github.com/kolide/launcher/pkg/autoupdate" "github.com/kolide/launcher/pkg/log/multislogger" - "go.etcd.io/bbolt" ) // type alias Flags, so that we can embed it inside knapsack, as `flags` and not `Flags` @@ -26,14 +25,7 @@ type knapsack struct { stores map[storage.Store]types.KVStore // Embed flags so we get all the flag interfaces flags - - // BboltDB is the underlying bbolt database. - // Ideally, we can eventually remove this. This is only here because some parts of the codebase - // like the osquery extension have a direct dependency on bbolt and need this reference. - // If we are able to abstract bbolt out completely in these areas, we should be able to - // remove this field and prevent "leaking" bbolt into places it doesn't need to. - db *bbolt.DB - + storageStatTracker types.StorageStatTracker slogger, systemSlogger *multislogger.MultiSlogger // This struct is a work in progress, and will be iteratively added to as needs arise. @@ -41,7 +33,7 @@ type knapsack struct { // Querier } -func New(stores map[storage.Store]types.KVStore, flags types.Flags, db *bbolt.DB, slogger, systemSlogger *multislogger.MultiSlogger) *knapsack { +func New(stores map[storage.Store]types.KVStore, flags types.Flags, sStatTracker types.StorageStatTracker, slogger, systemSlogger *multislogger.MultiSlogger) *knapsack { if slogger == nil { slogger = multislogger.New() } @@ -50,11 +42,11 @@ func New(stores map[storage.Store]types.KVStore, flags types.Flags, db *bbolt.DB } k := &knapsack{ - db: db, - flags: flags, - stores: stores, - slogger: slogger, - systemSlogger: systemSlogger, + storageStatTracker: sStatTracker, + flags: flags, + stores: stores, + slogger: slogger, + systemSlogger: systemSlogger, } return k @@ -74,9 +66,9 @@ func (k *knapsack) AddSlogHandler(handler ...slog.Handler) { k.systemSlogger.AddHandler(handler...) } -// BboltDB interface methods -func (k *knapsack) BboltDB() *bbolt.DB { - return k.db +// storage stat tracking interface methods +func (k *knapsack) StorageStatTracker() types.StorageStatTracker { + return k.storageStatTracker } // Stores interface methods diff --git a/ee/agent/storage/bbolt/stats_bbolt.go b/ee/agent/storage/bbolt/stats_bbolt.go new file mode 100644 index 000000000..a7cfdac28 --- /dev/null +++ b/ee/agent/storage/bbolt/stats_bbolt.go @@ -0,0 +1,112 @@ +package agentbbolt + +import ( + "encoding/json" + "fmt" + "log/slog" + + "go.etcd.io/bbolt" +) + +type bucketStatsHolder struct { + Stats bbolt.BucketStats + FillPercent float64 + NumberOfKeys int + Size int +} + +type dbStatsHolder struct { + Stats bbolt.TxStats + Size int64 +} + +type Stats struct { + DB dbStatsHolder + Buckets map[string]bucketStatsHolder +} + +type bboltStatStore struct { + slogger *slog.Logger + db *bbolt.DB +} + +// NewStatStore provides a wrapper around a bbolt.DB connection. This is done at the global (above bucket) level +// and should be used for any operations regarding the collection of storage statistics +func NewStatStore(slogger *slog.Logger, db *bbolt.DB) (*bboltStatStore, error) { + if db == nil { + return nil, NoDbError{} + } + + s := &bboltStatStore{ + slogger: slogger, + db: db, + } + + return s, nil +} + +func (s *bboltStatStore) SizeBytes() (int64, error) { + if s == nil || s.db == nil { + return 0, NoDbError{} + } + + var size int64 + + if err := s.db.View(func(tx *bbolt.Tx) error { + size = tx.Size() + return nil + }); err != nil { + return 0, fmt.Errorf("creating view tx to check size stat: %w", err) + } + + return size, nil +} + +// GetStats returns a json blob containing both global and bucket-level +// statistics. Note that the bucketName set does not impact the output, all buckets +// will be traversed for stats regardless +func (s *bboltStatStore) GetStats() ([]byte, error) { + if s == nil || s.db == nil { + return nil, NoDbError{} + } + + stats := &Stats{ + Buckets: make(map[string]bucketStatsHolder), + } + + if err := s.db.View(func(tx *bbolt.Tx) error { + stats.DB.Stats = tx.Stats() + stats.DB.Size = tx.Size() + + if err := tx.ForEach(bucketStatsFunc(stats)); err != nil { + return fmt.Errorf("dumping bucket: %w", err) + } + return nil + }); err != nil { + return nil, fmt.Errorf("creating view tx: %w", err) + } + + statsJson, err := json.Marshal(stats) + if err != nil { + return nil, err + } + + return statsJson, nil +} + +func bucketStatsFunc(stats *Stats) func([]byte, *bbolt.Bucket) error { + return func(name []byte, b *bbolt.Bucket) error { + bstats := b.Stats() + + // KeyN is the number of keys + // LeafAlloc is pretty close the number of bytes used + stats.Buckets[string(name)] = bucketStatsHolder{ + Stats: bstats, + FillPercent: b.FillPercent, + NumberOfKeys: bstats.KeyN, + Size: bstats.LeafAlloc, + } + + return nil + } +} diff --git a/ee/agent/types/keyvalue_store.go b/ee/agent/types/keyvalue_store.go index d8c9c10f8..8085dd6ac 100644 --- a/ee/agent/types/keyvalue_store.go +++ b/ee/agent/types/keyvalue_store.go @@ -55,6 +55,16 @@ type Appender interface { AppendValues(values ...[]byte) error } +// StorageStatTracker is an interface towards exposing the +// statistics of an underlying storage methodology +type StorageStatTracker interface { + // GetStats will return any relevant global and bucket-level statistics + // for the underlying storage methodology, expected as a json blob + GetStats() ([]byte, error) + // Size will return the total size in bytes for the stored key-value pairs + SizeBytes() (int64, error) +} + // GetterSetter is an interface that groups the Get and Set methods. type GetterSetter interface { Getter diff --git a/ee/agent/types/knapsack.go b/ee/agent/types/knapsack.go index 51ccd4a51..04f399cf7 100644 --- a/ee/agent/types/knapsack.go +++ b/ee/agent/types/knapsack.go @@ -6,7 +6,6 @@ import "context" // launcher code and are typically valid for the lifetime of the launcher application instance. type Knapsack interface { Stores - BboltDB Flags Slogger // LatestOsquerydPath finds the path to the latest osqueryd binary, after accounting for updates. @@ -15,4 +14,5 @@ type Knapsack interface { ReadEnrollSecret() (string, error) // CurrentEnrollmentStatus returns the current enrollment status of the launcher installation CurrentEnrollmentStatus() (EnrollmentStatus, error) + StorageStatTracker() StorageStatTracker } diff --git a/ee/debug/checkups/bboltdb.go b/ee/debug/checkups/bboltdb.go index d275692d1..4d71478b2 100644 --- a/ee/debug/checkups/bboltdb.go +++ b/ee/debug/checkups/bboltdb.go @@ -2,11 +2,11 @@ package checkups import ( "context" + "encoding/json" "errors" "fmt" "io" - "github.com/kolide/launcher/ee/agent" "github.com/kolide/launcher/ee/agent/types" ) @@ -20,21 +20,24 @@ func (c *bboltdbCheckup) Name() string { } func (c *bboltdbCheckup) Run(_ context.Context, _ io.Writer) error { - db := c.k.BboltDB() + db := c.k.StorageStatTracker() if db == nil { - return errors.New("no DB available") + return errors.New("no db connection available for storage stat tracking") } - stats, err := agent.GetStats(db) + stats, err := db.GetStats() if err != nil { return fmt.Errorf("getting db stats: %w", err) } - c.data = make(map[string]any) - for k, v := range stats.Buckets { - c.data[k] = v + data := make(map[string]any) + + if err := json.Unmarshal(stats, &data); err != nil { + return fmt.Errorf("unmarshalling storage stats json: %w", err) } + c.data = data + return nil } diff --git a/ee/debug/checkups/host.go b/ee/debug/checkups/host.go index 36525b005..5551242b2 100644 --- a/ee/debug/checkups/host.go +++ b/ee/debug/checkups/host.go @@ -60,17 +60,17 @@ func (hc *hostInfoCheckup) Run(ctx context.Context, extraFH io.Writer) error { } func (hc *hostInfoCheckup) bboltDbSize() string { - db := hc.k.BboltDB() + db := hc.k.StorageStatTracker() if db == nil { - return "error: bbolt db connection was not available via knapsack" + return "error: storage stat tracking db connection was not available via knapsack" } - boltStats, err := agent.GetStats(db) + sizeBytes, err := db.SizeBytes() if err != nil { - return fmt.Sprintf("encountered error accessing bbolt stats: %s", err.Error()) + return fmt.Sprintf("encountered error accessing storage stats: %s", err.Error()) } - return strconv.FormatInt(boltStats.DB.Size, 10) + return strconv.FormatInt(sizeBytes, 10) } func hostName() string { diff --git a/pkg/osquery/table/launcher_db_info.go b/pkg/osquery/table/launcher_db_info.go index 5cf4a86f9..11dc185b3 100644 --- a/pkg/osquery/table/launcher_db_info.go +++ b/pkg/osquery/table/launcher_db_info.go @@ -2,31 +2,28 @@ package table import ( "context" - "encoding/json" "fmt" "strings" - "github.com/kolide/launcher/ee/agent" + "github.com/kolide/launcher/ee/agent/types" "github.com/kolide/launcher/ee/dataflatten" "github.com/kolide/launcher/ee/tables/dataflattentable" "github.com/kolide/launcher/ee/tables/tablehelpers" "github.com/osquery/osquery-go/plugin/table" - "go.etcd.io/bbolt" ) -func LauncherDbInfo(db *bbolt.DB) *table.Plugin { +func LauncherDbInfo(sStatTracker types.StorageStatTracker) *table.Plugin { columns := dataflattentable.Columns() - return table.NewPlugin("kolide_launcher_db_info", columns, generateLauncherDbInfo(db)) + return table.NewPlugin("kolide_launcher_db_info", columns, generateLauncherDbInfo(sStatTracker)) } -func generateLauncherDbInfo(db *bbolt.DB) table.GenerateFunc { +func generateLauncherDbInfo(sStatTracker types.StorageStatTracker) table.GenerateFunc { return func(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) { - stats, err := agent.GetStats(db) - if err != nil { - return nil, err + if sStatTracker == nil { + return nil, fmt.Errorf("unable to gather db info without stat tracking connection") } - statsJson, err := json.Marshal(stats) + statsJson, err := sStatTracker.GetStats() if err != nil { return nil, err } diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go index aa5007ceb..5ccc9e9ed 100644 --- a/pkg/osquery/table/table.go +++ b/pkg/osquery/table/table.go @@ -23,7 +23,7 @@ import ( func LauncherTables(k types.Knapsack) []osquery.OsqueryPlugin { return []osquery.OsqueryPlugin{ LauncherConfigTable(k.ConfigStore()), - LauncherDbInfo(k.BboltDB()), + LauncherDbInfo(k.StorageStatTracker()), LauncherInfoTable(k.ConfigStore()), launcher_db.TablePlugin("kolide_server_data", k.ServerProvidedDataStore()), launcher_db.TablePlugin("kolide_control_flags", k.AgentFlagsStore()), From 67a7dc0d55ef006d836f37e8adf21204eef48bc9 Mon Sep 17 00:00:00 2001 From: zackattack01 Date: Thu, 21 Mar 2024 14:29:28 -0400 Subject: [PATCH 2/4] regenerate knapsack mocks --- ee/agent/types/mocks/knapsack.go | 35 +++++++++++++++----------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/ee/agent/types/mocks/knapsack.go b/ee/agent/types/mocks/knapsack.go index 8be4fe74f..e00963661 100644 --- a/ee/agent/types/mocks/knapsack.go +++ b/ee/agent/types/mocks/knapsack.go @@ -5,10 +5,7 @@ package mocks import ( context "context" - bbolt "go.etcd.io/bbolt" - keys "github.com/kolide/launcher/ee/agent/flags/keys" - mock "github.com/stretchr/testify/mock" slog "log/slog" @@ -110,22 +107,6 @@ func (_m *Knapsack) AutoupdateInterval() time.Duration { return r0 } -// BboltDB provides a mock function with given fields: -func (_m *Knapsack) BboltDB() *bbolt.DB { - ret := _m.Called() - - var r0 *bbolt.DB - if rf, ok := ret.Get(0).(func() *bbolt.DB); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(*bbolt.DB) - } - } - - return r0 -} - // CertPins provides a mock function with given fields: func (_m *Knapsack) CertPins() [][]byte { ret := _m.Called() @@ -1490,6 +1471,22 @@ func (_m *Knapsack) StatusLogsStore() types.GetterSetterDeleterIteratorUpdaterCo return r0 } +// StorageStatTracker provides a mock function with given fields: +func (_m *Knapsack) StorageStatTracker() types.StorageStatTracker { + ret := _m.Called() + + var r0 types.StorageStatTracker + if rf, ok := ret.Get(0).(func() types.StorageStatTracker); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(types.StorageStatTracker) + } + } + + return r0 +} + // Stores provides a mock function with given fields: func (_m *Knapsack) Stores() map[storage.Store]types.GetterSetterDeleterIteratorUpdaterCounterAppender { ret := _m.Called() From 8551be7dc955555b2c0bcac60277d5ac1f12d272 Mon Sep 17 00:00:00 2001 From: zackattack01 Date: Thu, 21 Mar 2024 15:20:32 -0400 Subject: [PATCH 3/4] clean up remaining bbolt refs, fix up tests --- ee/agent/types/bboltdb.go | 7 ------- ee/debug/checkups/bboltdb.go | 16 ++++++++-------- ee/debug/checkups/checkpoint_test.go | 6 +++++- ee/debug/checkups/checkups.go | 2 +- ee/debug/checkups/host.go | 4 ++-- 5 files changed, 16 insertions(+), 19 deletions(-) delete mode 100644 ee/agent/types/bboltdb.go diff --git a/ee/agent/types/bboltdb.go b/ee/agent/types/bboltdb.go deleted file mode 100644 index 7fd4c8dd2..000000000 --- a/ee/agent/types/bboltdb.go +++ /dev/null @@ -1,7 +0,0 @@ -package types - -import "go.etcd.io/bbolt" - -type BboltDB interface { - BboltDB() *bbolt.DB -} diff --git a/ee/debug/checkups/bboltdb.go b/ee/debug/checkups/bboltdb.go index 4d71478b2..f6f1cac01 100644 --- a/ee/debug/checkups/bboltdb.go +++ b/ee/debug/checkups/bboltdb.go @@ -10,16 +10,16 @@ import ( "github.com/kolide/launcher/ee/agent/types" ) -type bboltdbCheckup struct { +type kvStorageStatsCheckup struct { k types.Knapsack data map[string]any } -func (c *bboltdbCheckup) Name() string { - return "bboltdb" +func (c *kvStorageStatsCheckup) Name() string { + return "KV Storage Stats" } -func (c *bboltdbCheckup) Run(_ context.Context, _ io.Writer) error { +func (c *kvStorageStatsCheckup) Run(_ context.Context, _ io.Writer) error { db := c.k.StorageStatTracker() if db == nil { return errors.New("no db connection available for storage stat tracking") @@ -41,18 +41,18 @@ func (c *bboltdbCheckup) Run(_ context.Context, _ io.Writer) error { return nil } -func (c *bboltdbCheckup) ExtraFileName() string { +func (c *kvStorageStatsCheckup) ExtraFileName() string { return "" } -func (c *bboltdbCheckup) Status() Status { +func (c *kvStorageStatsCheckup) Status() Status { return Informational } -func (c *bboltdbCheckup) Summary() string { +func (c *kvStorageStatsCheckup) Summary() string { return "N/A" } -func (c *bboltdbCheckup) Data() any { +func (c *kvStorageStatsCheckup) Data() any { return c.data } diff --git a/ee/debug/checkups/checkpoint_test.go b/ee/debug/checkups/checkpoint_test.go index 1ff1b7b90..36b7f38e6 100644 --- a/ee/debug/checkups/checkpoint_test.go +++ b/ee/debug/checkups/checkpoint_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + agentbbolt "github.com/kolide/launcher/ee/agent/storage/bbolt" storageci "github.com/kolide/launcher/ee/agent/storage/ci" "github.com/kolide/launcher/ee/agent/types" typesmocks "github.com/kolide/launcher/ee/agent/types/mocks" @@ -15,10 +16,13 @@ import ( func TestInterrupt_Multiple(t *testing.T) { t.Parallel() + statStore, err := agentbbolt.NewStatStore(multislogger.NewNopLogger(), storageci.SetupDB(t)) + require.NoError(t, err) + mockKnapsack := typesmocks.NewKnapsack(t) mockKnapsack.On("UpdateChannel").Return("nightly").Maybe() mockKnapsack.On("TufServerURL").Return("localhost").Maybe() - mockKnapsack.On("BboltDB").Return(storageci.SetupDB(t)).Maybe() + mockKnapsack.On("StorageStatTracker").Return(statStore).Maybe() mockKnapsack.On("KolideHosted").Return(false).Maybe() mockKnapsack.On("KolideServerURL").Return("localhost").Maybe() mockKnapsack.On("ControlServerURL").Return("localhost").Maybe() diff --git a/ee/debug/checkups/checkups.go b/ee/debug/checkups/checkups.go index 6907460e3..0a30eadd2 100644 --- a/ee/debug/checkups/checkups.go +++ b/ee/debug/checkups/checkups.go @@ -103,7 +103,7 @@ func checkupsFor(k types.Knapsack, target targetBits) []checkupInt { {&launchdCheckup{}, doctorSupported | flareSupported}, {&runtimeCheckup{}, flareSupported}, {&enrollSecretCheckup{k: k}, doctorSupported | flareSupported}, - {&bboltdbCheckup{k: k}, flareSupported}, + {&kvStorageStatsCheckup{k: k}, flareSupported}, {&networkCheckup{}, doctorSupported | flareSupported}, {&installCheckup{}, flareSupported}, {&servicesCheckup{}, doctorSupported | flareSupported}, diff --git a/ee/debug/checkups/host.go b/ee/debug/checkups/host.go index 5551242b2..4b51f809c 100644 --- a/ee/debug/checkups/host.go +++ b/ee/debug/checkups/host.go @@ -36,7 +36,7 @@ func (hc *hostInfoCheckup) Run(ctx context.Context, extraFH io.Writer) error { hc.data = make(map[string]any) hc.data["hostname"] = hostName() hc.data["keyinfo"] = agentKeyInfo() - hc.data["bbolt_db_size"] = hc.bboltDbSize() + hc.data["bbolt_db_size"] = hc.kvStorageSizeBytes() desktopProcesses := runner.InstanceDesktopProcessRecords() hc.data["user_desktop_processes"] = desktopProcesses hc.data["enrollment_status"] = naIfError(hc.k.CurrentEnrollmentStatus()) @@ -59,7 +59,7 @@ func (hc *hostInfoCheckup) Run(ctx context.Context, extraFH io.Writer) error { return nil } -func (hc *hostInfoCheckup) bboltDbSize() string { +func (hc *hostInfoCheckup) kvStorageSizeBytes() string { db := hc.k.StorageStatTracker() if db == nil { return "error: storage stat tracking db connection was not available via knapsack" From 95589114c5dd526dc6e521e1c5b72c0bbfcce5d6 Mon Sep 17 00:00:00 2001 From: zackattack01 Date: Thu, 21 Mar 2024 15:49:53 -0400 Subject: [PATCH 4/4] remove unused dbdump from agent --- ee/agent/dbdump.go | 61 ---------------------------------------------- 1 file changed, 61 deletions(-) delete mode 100644 ee/agent/dbdump.go diff --git a/ee/agent/dbdump.go b/ee/agent/dbdump.go deleted file mode 100644 index 2f5745830..000000000 --- a/ee/agent/dbdump.go +++ /dev/null @@ -1,61 +0,0 @@ -package agent - -import ( - "fmt" - - "go.etcd.io/bbolt" -) - -type bucketStatsHolder struct { - Stats bbolt.BucketStats - FillPercent float64 - NumberOfKeys int - Size int -} - -type dbStatsHolder struct { - Stats bbolt.TxStats - Size int64 -} - -type Stats struct { - DB dbStatsHolder - Buckets map[string]bucketStatsHolder -} - -func GetStats(db *bbolt.DB) (*Stats, error) { - stats := &Stats{ - Buckets: make(map[string]bucketStatsHolder), - } - - if err := db.View(func(tx *bbolt.Tx) error { - stats.DB.Stats = tx.Stats() - stats.DB.Size = tx.Size() - - if err := tx.ForEach(bucketStatsFunc(stats)); err != nil { - return fmt.Errorf("dumping bucket: %w", err) - } - return nil - }); err != nil { - return nil, fmt.Errorf("creating view tx: %w", err) - } - - return stats, nil -} - -func bucketStatsFunc(stats *Stats) func([]byte, *bbolt.Bucket) error { - return func(name []byte, b *bbolt.Bucket) error { - bstats := b.Stats() - - // KeyN is the number of keys - // LeafAlloc is pretty close the number of bytes used - stats.Buckets[string(name)] = bucketStatsHolder{ - Stats: bstats, - FillPercent: b.FillPercent, - NumberOfKeys: bstats.KeyN, - Size: bstats.LeafAlloc, - } - - return nil - } -}