Repository: kolide/launcher

Issue:
Support plain `leveldb` source type for KATC
Some extension data is stored in a plain leveldb (_not_ indexeddb on top of leveldb) database. We would like to be able to query this via the KATC tables.

Add new `leveldb` source type to https://github.com/kolide/launcher/blob/main/ee/katc/config.go, and support reading a key-value leveldb. We can use some of the basic code from https://github.com/kolide/launcher/tree/main/ee/indexeddb -- there's no encoding and no nesting present, so we can simply iterate over the database keys instead of having to find the correct database and object store.

Discussion:

Solution patch:
diff --git a/ee/indexeddb/indexeddb.go b/ee/indexeddb/indexeddb.go
index 9ca76653..49d1f378 100644
--- a/ee/indexeddb/indexeddb.go
+++ b/ee/indexeddb/indexeddb.go
@@ -37,7 +37,7 @@ func QueryIndexeddbObjectStore(ctx context.Context, dbLocation string, dbName st
 	defer span.End()
 
 	// If Chrome is open, we won't be able to open the db. So, copy it to a temporary location first.
-	tempDbCopyLocation, err := copyIndexeddb(ctx, dbLocation)
+	tempDbCopyLocation, err := CopyLeveldb(ctx, dbLocation)
 	if err != nil {
 		if tempDbCopyLocation != "" {
 			_ = os.RemoveAll(tempDbCopyLocation)
@@ -49,20 +49,9 @@ func QueryIndexeddbObjectStore(ctx context.Context, dbLocation string, dbName st
 
 	objs := make([]map[string][]byte, 0)
 
-	opts := &opt.Options{
-		Comparer:               indexeddbComparer,
-		DisableSeeksCompaction: true,               // no need to perform compaction
-		Strict:                 opt.StrictRecovery, // we prefer to drop corrupted data rather than fail to open the db altogether
-	}
-	db, dbOpenErr := leveldb.OpenFile(tempDbCopyLocation, opts)
-	if dbOpenErr != nil {
-		// Perform recover in case we missed something while copying
-		var dbRecoverErr error
-		db, dbRecoverErr = leveldb.RecoverFile(tempDbCopyLocation, opts)
-		if dbRecoverErr != nil {
-			return nil, fmt.Errorf("opening db: `%+v`; recovering db: %w", dbOpenErr, dbRecoverErr)
-		}
-		span.AddEvent("db_recovered")
+	db, err := OpenLeveldb(tempDbCopyLocation)
+	if err != nil {
+		return nil, fmt.Errorf("opening leveldb: %w", err)
 	}
 	defer db.Close()
 	span.AddEvent("db_opened")
@@ -133,7 +122,7 @@ func QueryIndexeddbObjectStore(ctx context.Context, dbLocation string, dbName st
 	return objs, nil
 }
 
-func copyIndexeddb(ctx context.Context, sourceDb string) (string, error) {
+func CopyLeveldb(ctx context.Context, sourceDb string) (string, error) {
 	ctx, span := traces.StartSpan(ctx)
 	defer span.End()
 
@@ -192,3 +181,22 @@ func copyFile(ctx context.Context, src string, dest string) error {
 
 	return nil
 }
+
+func OpenLeveldb(dbLocation string) (*leveldb.DB, error) {
+	opts := &opt.Options{
+		Comparer:               indexeddbComparer,
+		DisableSeeksCompaction: true,               // no need to perform compaction
+		Strict:                 opt.StrictRecovery, // we prefer to drop corrupted data rather than fail to open the db altogether
+	}
+	db, dbOpenErr := leveldb.OpenFile(dbLocation, opts)
+	if dbOpenErr != nil {
+		// Perform recover in case we missed something while copying
+		var dbRecoverErr error
+		db, dbRecoverErr = leveldb.RecoverFile(dbLocation, opts)
+		if dbRecoverErr != nil {
+			return nil, fmt.Errorf("opening db: `%+v`; recovering db: %w", dbOpenErr, dbRecoverErr)
+		}
+	}
+
+	return db, nil
+}
diff --git a/ee/katc/config.go b/ee/katc/config.go
index 215ea877..1670672d 100644
--- a/ee/katc/config.go
+++ b/ee/katc/config.go
@@ -33,6 +33,7 @@ type sourceData struct {
 const (
 	sqliteSourceType           = "sqlite"
 	indexeddbLeveldbSourceType = "indexeddb_leveldb"
+	leveldbSourceType          = "leveldb"
 )
 
 func (kst *katcSourceType) UnmarshalJSON(data []byte) error {
@@ -51,6 +52,10 @@ func (kst *katcSourceType) UnmarshalJSON(data []byte) error {
 		kst.name = indexeddbLeveldbSourceType
 		kst.dataFunc = indexeddbLeveldbData
 		return nil
+	case leveldbSourceType:
+		kst.name = leveldbSourceType
+		kst.dataFunc = leveldbData
+		return nil
 	default:
 		return fmt.Errorf("unknown table type %s", s)
 	}
@@ -146,6 +151,20 @@ func ConstructKATCTables(config map[string]string, flags types.Flags, slogger *s
 		}
 
 		t, columns := newKatcTable(tableName, cfg, slogger)
+
+		// Validate that the columns are valid for this table type -- only checked
+		// for LevelDB tables currently
+		if t.sourceType.name == leveldbSourceType {
+			if err := validateLeveldbTableColumns(columns); err != nil {
+				slogger.Log(context.TODO(), slog.LevelWarn,
+					"invalid columns for leveldb table",
+					"table_name", tableName,
+					"err", err,
+				)
+				continue
+			}
+		}
+
 		plugins = append(plugins, tablewrapper.New(flags, slogger, tableName, columns, t.generate))
 	}
 
diff --git a/ee/katc/indexeddb_leveldb.go b/ee/katc/indexeddb_leveldb.go
index 63f2f104..9bc03ae1 100644
--- a/ee/katc/indexeddb_leveldb.go
+++ b/ee/katc/indexeddb_leveldb.go
@@ -32,7 +32,7 @@ func indexeddbLeveldbData(ctx context.Context, slogger *slog.Logger, sourcePaths
 		}
 
 		// Extract database and table from query
-		dbName, objectStoreName, err := extractQueryTargets(query)
+		dbName, objectStoreName, err := extractIndexeddbQueryTargets(query)
 		if err != nil {
 			return nil, fmt.Errorf("getting db and object store names: %w", err)
 		}
@@ -63,10 +63,10 @@ func indexeddbLeveldbData(ctx context.Context, slogger *slog.Logger, sourcePaths
 	return results, nil
 }
 
-// extractQueryTargets retrieves the targets of the query (the database name and the object store name)
+// extractIndexeddbQueryTargets retrieves the targets of the query (the database name and the object store name)
 // from the query. IndexedDB is a NoSQL database, so we expect to retrieve all rows from the given
 // object store within the given database name.
-func extractQueryTargets(query string) (string, string, error) {
+func extractIndexeddbQueryTargets(query string) (string, string, error) {
 	parts := strings.Split(query, ".")
 	if len(parts) != 2 {
 		return "", "", fmt.Errorf("unable to extract query targets from query: expected `<db name>.<obj store name>`, got `%s`", query)
diff --git a/ee/katc/leveldb.go b/ee/katc/leveldb.go
new file mode 100644
index 00000000..2de53ce8
--- /dev/null
+++ b/ee/katc/leveldb.go
@@ -0,0 +1,144 @@
+package katc
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/kolide/launcher/ee/indexeddb"
+	"github.com/kolide/launcher/pkg/traces"
+	"github.com/osquery/osquery-go/plugin/table"
+)
+
+const (
+	keyColumnName   = "key"
+	valueColumnName = "value"
+)
+
+var leveldbExpectedColumns = map[string]struct{}{
+	keyColumnName:   {},
+	valueColumnName: {},
+	pathColumnName:  {},
+}
+
+func validateLeveldbTableColumns(columns []table.ColumnDefinition) error {
+	for _, c := range columns {
+		if _, ok := leveldbExpectedColumns[c.Name]; !ok {
+			return fmt.Errorf("unsupported column %s for leveldb table", c.Name)
+		}
+	}
+
+	return nil
+}
+
+// leveldbData is the dataFunc for plain LevelDB databases without additional encoding
+// or nesting. IndexedDB databases leveraging LevelDB should use `indexeddbLeveldbData`
+// instead. If set, the query is a comma-separated allowlist of keys to return; if empty,
+// all keys are returned.
+func leveldbData(ctx context.Context, slogger *slog.Logger, sourcePaths []string, query string, queryContext table.QueryContext) ([]sourceData, error) {
+	ctx, span := traces.StartSpan(ctx)
+	defer span.End()
+
+	// Pull out path constraints from the query against the KATC table, to avoid querying more leveldb files than we need to.
+	pathConstraintsFromQuery := getPathConstraint(queryContext)
+
+	allowedKeyMap := extractLeveldbQueryTargets(query)
+
+	results := make([]sourceData, 0)
+	for _, sourcePath := range sourcePaths {
+		pathPattern := sourcePatternToGlobbablePattern(sourcePath)
+		leveldbs, err := filepath.Glob(pathPattern)
+		if err != nil {
+			return nil, fmt.Errorf("globbing for leveldb files: %w", err)
+		}
+
+		// Query databases
+		for _, dbPath := range leveldbs {
+			// Check to make sure `db` adheres to pathConstraintsFromQuery. This is an
+			// optimization to avoid work, if osquery sqlite filtering is going to exclude it.
+			valid, err := checkPathConstraints(dbPath, pathConstraintsFromQuery)
+			if err != nil {
+				return nil, fmt.Errorf("checking source path constraints: %w", err)
+			}
+			if !valid {
+				continue
+			}
+
+			rowsFromDb, err := queryLeveldb(ctx, dbPath, allowedKeyMap)
+			if err != nil {
+				return nil, fmt.Errorf("querying leveldb at %s: %w", dbPath, err)
+			}
+			results = append(results, sourceData{
+				path: dbPath,
+				rows: rowsFromDb,
+			})
+		}
+	}
+
+	return results, nil
+}
+
+func queryLeveldb(ctx context.Context, path string, allowedKeyMap map[string]struct{}) ([]map[string][]byte, error) {
+	ctx, span := traces.StartSpan(ctx)
+	defer span.End()
+
+	// If Chrome is open, we won't be able to open the db. So, copy it to a temporary location first.
+	tempDbCopyLocation, err := indexeddb.CopyLeveldb(ctx, path)
+	if err != nil {
+		if tempDbCopyLocation != "" {
+			_ = os.RemoveAll(tempDbCopyLocation)
+		}
+		return nil, fmt.Errorf("unable to copy db: %w", err)
+	}
+	// The copy was successful -- make sure we clean it up after we're done
+	defer os.RemoveAll(tempDbCopyLocation)
+
+	db, err := indexeddb.OpenLeveldb(tempDbCopyLocation)
+	if err != nil {
+		return nil, fmt.Errorf("opening leveldb: %w", err)
+	}
+	defer db.Close()
+
+	rowsFromDb := make([]map[string][]byte, 0)
+	iter := db.NewIterator(nil, nil)
+	for iter.Next() {
+		// We have to copy over the key/value to prevent their data from
+		// being overwritten on subsequent iterations.
+		k := make([]byte, len(iter.Key()))
+		copy(k, iter.Key())
+		if len(allowedKeyMap) > 0 {
+			if _, ok := allowedKeyMap[string(k)]; !ok {
+				// Key is not allowlisted -- skip it
+				continue
+			}
+		}
+		v := make([]byte, len(iter.Value()))
+		copy(v, iter.Value())
+		rowsFromDb = append(rowsFromDb, map[string][]byte{
+			"key":   k,
+			"value": v,
+		})
+	}
+	iter.Release()
+	if err := iter.Error(); err != nil {
+		return rowsFromDb, fmt.Errorf("iterator error: %w", err)
+	}
+
+	return rowsFromDb, nil
+}
+
+func extractLeveldbQueryTargets(query string) map[string]struct{} {
+	allowedKeyMap := make(map[string]struct{})
+
+	for _, allowedKey := range strings.Split(query, ",") {
+		if len(allowedKey) == 0 {
+			continue
+		}
+		allowedKeyMap[allowedKey] = struct{}{}
+	}
+
+	return allowedKeyMap
+}



Test patch:
diff --git a/ee/katc/config_test.go b/ee/katc/config_test.go
index 372ff8b0..90ee9a34 100644
--- a/ee/katc/config_test.go
+++ b/ee/katc/config_test.go
@@ -50,6 +50,20 @@ func TestConstructKATCTables(t *testing.T) {
 			},
 			expectedPluginCount: 1,
 		},
+		{
+			testCaseName: "leveldb",
+			katcConfig: map[string]string{
+				"kolide_leveldb_test": `{
+					"source_type": "leveldb",
+					"columns": ["key", "value"],
+					"source_paths": ["/some/path/to/db.leveldb"],
+					"source_query": "",
+					"row_transform_steps": [],
+					"overlays": []
+				}`,
+			},
+			expectedPluginCount: 1,
+		},
 		{
 			testCaseName: "overlay",
 			katcConfig: map[string]string{
@@ -148,6 +162,20 @@ func TestConstructKATCTables(t *testing.T) {
 			},
 			expectedPluginCount: 0,
 		},
+		{
+			testCaseName: "invalid leveldb column",
+			katcConfig: map[string]string{
+				"kolide_leveldb_test": `{
+					"source_type": "leveldb",
+					"columns": ["key", "config"],
+					"source_paths": ["/some/path/to/db.leveldb"],
+					"source_query": "",
+					"row_transform_steps": [],
+					"overlays": []
+				}`,
+			},
+			expectedPluginCount: 0,
+		},
 	} {
 		tt := tt
 		t.Run(tt.testCaseName, func(t *testing.T) {
diff --git a/ee/katc/indexeddb_leveldb_test.go b/ee/katc/indexeddb_leveldb_test.go
index e305ff97..31234652 100644
--- a/ee/katc/indexeddb_leveldb_test.go
+++ b/ee/katc/indexeddb_leveldb_test.go
@@ -48,7 +48,7 @@ func Test_extractQueryTargets(t *testing.T) {
 		t.Run(tt.testCaseName, func(t *testing.T) {
 			t.Parallel()
 
-			dbName, objStoreName, err := extractQueryTargets(tt.query)
+			dbName, objStoreName, err := extractIndexeddbQueryTargets(tt.query)
 
 			if tt.expectErr {
 				require.Error(t, err)
diff --git a/ee/katc/table_test.go b/ee/katc/table_test.go
index 46e61358..b5cd41ec 100644
--- a/ee/katc/table_test.go
+++ b/ee/katc/table_test.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/kolide/goleveldb/leveldb/opt"
 	"github.com/kolide/launcher/ee/indexeddb"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/osquery/osquery-go/plugin/table"
@@ -277,6 +278,73 @@ func TestQueryChromeIndexedDB(t *testing.T) {
 	}
 }
 
+func TestQueryLevelDB(t *testing.T) {
+	t.Parallel()
+
+	// Create a level db to query
+	tempDir := t.TempDir()
+	db, err := indexeddb.OpenLeveldb(tempDir)
+	require.NoError(t, err)
+
+	// Add some data
+	expectedKey := "some-key"
+	expectedValue := "some-value"
+	require.NoError(t, db.Put([]byte(expectedKey), []byte(expectedValue), &opt.WriteOptions{}))
+	require.NoError(t, db.Close())
+
+	// Construct table
+	sourceQuery := strings.Join([]string{expectedKey}, ",")
+	cfg := katcTableConfig{
+		Columns: []string{"key", "value"},
+		katcTableDefinition: katcTableDefinition{
+			SourceType: &katcSourceType{
+				name:     leveldbSourceType,
+				dataFunc: leveldbData,
+			},
+			SourcePaths:       &[]string{filepath.Join("some", "incorrect", "path")},
+			SourceQuery:       &sourceQuery,
+			RowTransformSteps: &[]rowTransformStep{},
+		},
+		Overlays: []katcTableConfigOverlay{
+			{
+				Filters: map[string]string{
+					"goos": runtime.GOOS,
+				},
+				katcTableDefinition: katcTableDefinition{
+					SourcePaths: &[]string{tempDir},
+				},
+			},
+		},
+	}
+	testTable, _ := newKatcTable("test_katc_table", cfg, multislogger.NewNopLogger())
+
+	// Make a query context restricting the source to our exact source indexeddb database
+	queryContext := table.QueryContext{
+		Constraints: map[string]table.ConstraintList{
+			pathColumnName: {
+				Constraints: []table.Constraint{
+					{
+						Operator:   table.OperatorEquals,
+						Expression: tempDir,
+					},
+				},
+			},
+		},
+	}
+
+	// At long last: run a query
+	results, err := testTable.generate(context.TODO(), queryContext)
+	require.NoError(t, err)
+
+	require.Equal(t, 1, len(results))
+	require.Contains(t, results[0], "key")
+	require.Equal(t, expectedKey, results[0]["key"])
+	require.Contains(t, results[0], "value")
+	require.Equal(t, expectedValue, results[0]["value"])
+	require.Contains(t, results[0], "path")
+	require.Equal(t, tempDir, results[0]["path"])
+}
+
 func Test_checkSourcePathConstraints(t *testing.T) {
 	t.Parallel()