diff --git a/plugin/table/table.go b/plugin/table/table.go index 711cb20..686208e 100644 --- a/plugin/table/table.go +++ b/plugin/table/table.go @@ -182,6 +182,30 @@ type QueryContext struct { // Constraints is a map from column name to the details of the // constraints on that column. Constraints map[string]ConstraintList + + // ColumnsUsed contains the columns osquery reports as needed by the query. + // If this is nil, the extension should assume all columns may be needed. + ColumnsUsed []string + + // ColumnsUsedBitset is the raw SQLite column-use bitset passed through by + // osquery. Consumers that need name-based checks should prefer ColumnsUsed. + ColumnsUsedBitset *uint64 +} + +func (q QueryContext) HasColumnUsage() bool { + return q.ColumnsUsed != nil || q.ColumnsUsedBitset != nil +} + +func (q QueryContext) IsColumnUsed(name string) bool { + if q.ColumnsUsed == nil { + return true + } + for _, column := range q.ColumnsUsed { + if column == name { + return true + } + } + return false } // ConstraintList contains the details of the constraints for the given column. @@ -226,7 +250,9 @@ const ( // The following types and functions exist for parsing of the queryContext // JSON and are not made public. type queryContextJSON struct { - Constraints []constraintListJSON `json:"constraints"` + Constraints []constraintListJSON `json:"constraints"` + ColumnsUsed []string `json:"colsUsed"` + ColumnsUsedBitset *uint64 `json:"colsUsedBitset"` } type constraintListJSON struct { @@ -243,7 +269,11 @@ func parseQueryContext(ctxJSON string) (*QueryContext, error) { return nil, errors.Wrap(err, "unmarshaling context JSON") } - ctx := QueryContext{map[string]ConstraintList{}} + ctx := QueryContext{ + Constraints: map[string]ConstraintList{}, + ColumnsUsed: parsed.ColumnsUsed, + ColumnsUsedBitset: parsed.ColumnsUsedBitset, + } for _, cList := range parsed.Constraints { constraints, err := parseConstraintList(cList.List) if err != nil { diff --git a/plugin/table/table_test.go b/plugin/table/table_test.go index 4530e3e..452b2c6 100644 --- a/plugin/table/table_test.go +++ b/plugin/table/table_test.go @@ -57,7 +57,7 @@ func TestTablePlugin(t *testing.T) { // Call with good action and context resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "generate", "context": "{}"}) - assert.Equal(t, QueryContext{map[string]ConstraintList{}}, calledQueryCtx) + assert.Equal(t, QueryContext{Constraints: map[string]ConstraintList{}}, calledQueryCtx) assert.Equal(t, &StatusOK, resp.Status) assert.Equal(t, osquery.ExtensionPluginResponse{ { @@ -189,7 +189,7 @@ func TestParseQueryContext(t *testing.T) { } ] }`, - context: QueryContext{map[string]ConstraintList{ + context: QueryContext{Constraints: map[string]ConstraintList{ "big_int": ConstraintList{ColumnTypeBigInt, []Constraint{}}, "double": ConstraintList{ColumnTypeDouble, []Constraint{}}, "integer": ConstraintList{ColumnTypeInteger, []Constraint{}}, @@ -233,7 +233,7 @@ func TestParseQueryContext(t *testing.T) { ] } `, - context: QueryContext{map[string]ConstraintList{ + context: QueryContext{Constraints: map[string]ConstraintList{ "big_int": ConstraintList{ColumnTypeBigInt, []Constraint{}}, "double": ConstraintList{ColumnTypeDouble, []Constraint{{OperatorGreaterThanOrEquals, "3.1"}}}, "integer": ConstraintList{ColumnTypeInteger, []Constraint{}}, @@ -318,3 +318,41 @@ func TestParseVaryingQueryContexts(t *testing.T) { }) } } + +func TestParseQueryContextColumnsUsed(t *testing.T) { + context, err := parseQueryContext(`{ + "constraints": [], + "colsUsed": ["source", "session_id"], + "colsUsedBitset": 3 +}`) + require.NoError(t, err) + + require.NotNil(t, context) + require.True(t, context.HasColumnUsage()) + assert.Equal(t, []string{"source", "session_id"}, context.ColumnsUsed) + require.NotNil(t, context.ColumnsUsedBitset) + assert.Equal(t, uint64(3), *context.ColumnsUsedBitset) + assert.True(t, context.IsColumnUsed("source")) + assert.False(t, context.IsColumnUsed("text")) +} + +func TestParseQueryContextEmptyColumnsUsed(t *testing.T) { + context, err := parseQueryContext(`{ + "constraints": [], + "colsUsed": [], + "colsUsedBitset": 0 +}`) + require.NoError(t, err) + + require.NotNil(t, context) + require.True(t, context.HasColumnUsage()) + assert.NotNil(t, context.ColumnsUsed) + assert.False(t, context.IsColumnUsed("text")) +} + +func TestQueryContextIsColumnUsedFallback(t *testing.T) { + context := QueryContext{Constraints: map[string]ConstraintList{}} + + assert.False(t, context.HasColumnUsage()) + assert.True(t, context.IsColumnUsed("text")) +}