Get timeout for table generate function from control server (#2085)

kolide · Feb 3, 2025 · ddeead5 · ddeead5
1 parent 7fa0f94
commit ddeead5
Show file tree

Hide file tree

Showing 82 changed files with 583 additions and 278 deletions.
diff --git a/ee/agent/flags/flag_controller.go b/ee/agent/flags/flag_controller.go
@@ -703,3 +703,14 @@ func (fc *FlagController) Identifier() string {
 
 	return identifier
 }
+
+func (fc *FlagController) SetTableGenerateTimeout(interval time.Duration) error {
+	return fc.setControlServerValue(keys.TableGenerateTimeout, durationToBytes(interval))
+}
+func (fc *FlagController) TableGenerateTimeout() time.Duration {
+	return NewDurationFlagValue(fc.slogger, keys.TableGenerateTimeout,
+		WithDefault(4*time.Minute),
+		WithMin(30*time.Second),
+		WithMax(10*time.Minute),
+	).get(fc.getControlServerValue(keys.TableGenerateTimeout))
+}
diff --git a/ee/agent/flags/keys/keys.go b/ee/agent/flags/keys/keys.go
@@ -58,6 +58,7 @@ const (
 	LauncherWatchdogEnabled         FlagKey = "launcher_watchdog_enabled" // note that this will only impact windows deployments for now
 	SystrayRestartEnabled           FlagKey = "systray_restart_enabled"
 	CurrentRunningOsqueryVersion    FlagKey = "osquery_version"
+	TableGenerateTimeout            FlagKey = "table_generate_timeout"
 )
 
 func (key FlagKey) String() string {

diff --git a/ee/agent/types/flags.go b/ee/agent/types/flags.go
@@ -230,4 +230,8 @@ type Flags interface {
 
 	// Identifier is the package build identifier used to namespace our paths and service names
 	Identifier() string
+
+	// TableGenerateTimeout is the maximum time a Kolide extension table is permitted to take
+	SetTableGenerateTimeout(interval time.Duration) error
+	TableGenerateTimeout() time.Duration
 }
diff --git a/ee/agent/types/mocks/flags.go b/ee/agent/types/mocks/flags.go
diff --git a/ee/agent/types/mocks/knapsack.go b/ee/agent/types/mocks/knapsack.go
diff --git a/ee/katc/config.go b/ee/katc/config.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log/slog"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/indexeddb"
 	"github.com/kolide/launcher/ee/tables/tablewrapper"
 	"github.com/osquery/osquery-go"
@@ -130,7 +131,7 @@ type (
 
 // ConstructKATCTables takes stored configuration of KATC tables, parses the configuration,
 // and returns the constructed tables.
-func ConstructKATCTables(config map[string]string, slogger *slog.Logger) []osquery.OsqueryPlugin {
+func ConstructKATCTables(config map[string]string, flags types.Flags, slogger *slog.Logger) []osquery.OsqueryPlugin {
 	plugins := make([]osquery.OsqueryPlugin, 0)
 
 	for tableName, tableConfigStr := range config {
@@ -145,7 +146,7 @@ func ConstructKATCTables(config map[string]string, slogger *slog.Logger) []osque
 		}
 
 		t, columns := newKatcTable(tableName, cfg, slogger)
-		plugins = append(plugins, tablewrapper.New(slogger, tableName, columns, t.generate))
+		plugins = append(plugins, tablewrapper.New(flags, slogger, tableName, columns, t.generate))
 	}
 
 	return plugins

diff --git a/ee/katc/config_test.go b/ee/katc/config_test.go
@@ -5,8 +5,12 @@ import (
 	"fmt"
 	"runtime"
 	"testing"
+	"time"
 
+	"github.com/kolide/launcher/ee/agent/flags/keys"
+	typesmocks "github.com/kolide/launcher/ee/agent/types/mocks"
 	"github.com/kolide/launcher/pkg/log/multislogger"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -149,7 +153,11 @@ func TestConstructKATCTables(t *testing.T) {
 		t.Run(tt.testCaseName, func(t *testing.T) {
 			t.Parallel()
 
-			plugins := ConstructKATCTables(tt.katcConfig, multislogger.NewNopLogger())
+			mockFlags := typesmocks.NewFlags(t)
+			mockFlags.On("TableGenerateTimeout").Return(4 * time.Minute).Maybe()
+			mockFlags.On("RegisterChangeObserver", mock.Anything, keys.TableGenerateTimeout).Return().Maybe()
+
+			plugins := ConstructKATCTables(tt.katcConfig, mockFlags, multislogger.NewNopLogger())
 			require.Equal(t, tt.expectedPluginCount, len(plugins), "unexpected number of plugins")
 		})
 	}

diff --git a/ee/tables/airport/table_darwin.go b/ee/tables/airport/table_darwin.go
@@ -12,6 +12,7 @@ import (
 	"log/slog"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/dataflattentable"
@@ -32,7 +33,7 @@ type Table struct {
 
 const tableName = "kolide_airport_util"
 
-func TablePlugin(slogger *slog.Logger) *table.Plugin {
+func TablePlugin(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := dataflattentable.Columns(
 		table.TextColumn("option"),
 	)
@@ -42,7 +43,7 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 		slogger: slogger.With("name", tableName),
 	}
 
-	return tablewrapper.New(slogger, t.name, columns, t.generate)
+	return tablewrapper.New(flags, slogger, t.name, columns, t.generate)
 }
 
 type airportExecutor struct {

diff --git a/ee/tables/app-icons/app_icons_darwin.go b/ee/tables/app-icons/app_icons_darwin.go
@@ -35,6 +35,7 @@ import (
 	"image/png"
 	"unsafe"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/tables/tablewrapper"
 	"github.com/kolide/launcher/pkg/traces"
 	"github.com/nfnt/resize"
@@ -45,13 +46,13 @@ import (
 
 var crcTable = crc64.MakeTable(crc64.ECMA)
 
-func AppIcons(slogger *slog.Logger) *table.Plugin {
+func AppIcons(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := []table.ColumnDefinition{
 		table.TextColumn("path"),
 		table.TextColumn("icon"),
 		table.TextColumn("hash"),
 	}
-	return tablewrapper.New(slogger, "kolide_app_icons", columns, generateAppIcons)
+	return tablewrapper.New(flags, slogger, "kolide_app_icons", columns, generateAppIcons)
 }
 
 func generateAppIcons(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/apple_silicon_security_policy/table.go b/ee/tables/apple_silicon_security_policy/table.go
@@ -9,6 +9,7 @@ import (
 	"log/slog"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/dataflattentable"
@@ -24,7 +25,7 @@ type Table struct {
 	slogger *slog.Logger
 }
 
-func TablePlugin(slogger *slog.Logger) *table.Plugin {
+func TablePlugin(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := dataflattentable.Columns()
 
 	tableName := "kolide_apple_silicon_security_policy"
@@ -33,7 +34,7 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 		slogger: slogger.With("table", tableName),
 	}
 
-	return tablewrapper.New(slogger, tableName, columns, t.generate)
+	return tablewrapper.New(flags, slogger, tableName, columns, t.generate)
 }
 
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/crowdstrike/falcon_kernel_check/table.go b/ee/tables/crowdstrike/falcon_kernel_check/table.go
@@ -9,6 +9,7 @@ import (
 	"log/slog"
 	"regexp"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/tables/tablehelpers"
 	"github.com/kolide/launcher/ee/tables/tablewrapper"
@@ -20,7 +21,7 @@ type Table struct {
 	slogger *slog.Logger
 }
 
-func TablePlugin(slogger *slog.Logger) *table.Plugin {
+func TablePlugin(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := []table.ColumnDefinition{
 		table.TextColumn("kernel"),
 		table.IntegerColumn("supported"),
@@ -33,7 +34,7 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 		slogger: slogger.With("table", tableName),
 	}
 
-	return tablewrapper.New(slogger, tableName, columns, t.generate)
+	return tablewrapper.New(flags, slogger, tableName, columns, t.generate)
 }
 
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/crowdstrike/falconctl/table.go b/ee/tables/crowdstrike/falconctl/table.go
@@ -10,6 +10,7 @@ import (
 	"log/slog"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/dataflattentable"
@@ -48,7 +49,7 @@ type falconctlOptionsTable struct {
 	execFunc  execFunc
 }
 
-func NewFalconctlOptionTable(slogger *slog.Logger) *table.Plugin {
+func NewFalconctlOptionTable(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := dataflattentable.Columns(
 		table.TextColumn("options"),
 	)
@@ -59,7 +60,7 @@ func NewFalconctlOptionTable(slogger *slog.Logger) *table.Plugin {
 		execFunc:  tablehelpers.RunSimple,
 	}
 
-	return tablewrapper.New(slogger, t.tableName, columns, t.generate)
+	return tablewrapper.New(flags, slogger, t.tableName, columns, t.generate)
 }
 
 func (t *falconctlOptionsTable) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/cryptoinfotable/table.go b/ee/tables/cryptoinfotable/table.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/cryptoinfo"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/dataflattentable"
@@ -23,7 +24,7 @@ type Table struct {
 	slogger *slog.Logger
 }
 
-func TablePlugin(slogger *slog.Logger) *table.Plugin {
+func TablePlugin(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := dataflattentable.Columns(
 		table.TextColumn("passphrase"),
 		table.TextColumn("path"),
@@ -33,7 +34,7 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 		slogger: slogger.With("table", "kolide_cryptinfo"),
 	}
 
-	return tablewrapper.New(slogger, "kolide_cryptinfo", columns, t.generate)
+	return tablewrapper.New(flags, slogger, "kolide_cryptinfo", columns, t.generate)
 }
 
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/cryptsetup/table.go b/ee/tables/cryptsetup/table.go
@@ -9,6 +9,7 @@ import (
 	"log/slog"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/dataflattentable"
@@ -25,7 +26,7 @@ type Table struct {
 	name    string
 }
 
-func TablePlugin(slogger *slog.Logger) *table.Plugin {
+func TablePlugin(flags types.Flags, slogger *slog.Logger) *table.Plugin {
 	columns := dataflattentable.Columns(
 		table.TextColumn("name"),
 	)
@@ -35,7 +36,7 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 		name:    "kolide_cryptsetup_status",
 	}
 
-	return tablewrapper.New(slogger, t.name, columns, t.generate)
+	return tablewrapper.New(flags, slogger, t.name, columns, t.generate)
 }
 
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {

diff --git a/ee/tables/dataflattentable/exec.go b/ee/tables/dataflattentable/exec.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/tables/tablehelpers"
@@ -26,7 +27,7 @@ func WithKVSeparator(separator string) ExecTableOpt {
 	}
 }
 
-func TablePluginExec(slogger *slog.Logger, tableName string, dataSourceType DataSourceType, cmdGen allowedcmd.AllowedCommand, execArgs []string, opts ...ExecTableOpt) *table.Plugin {
+func TablePluginExec(flags types.Flags, slogger *slog.Logger, tableName string, dataSourceType DataSourceType, cmdGen allowedcmd.AllowedCommand, execArgs []string, opts ...ExecTableOpt) *table.Plugin {
 	columns := Columns()
 
 	t := &Table{
@@ -43,7 +44,7 @@ func TablePluginExec(slogger *slog.Logger, tableName string, dataSourceType Data
 
 	t.flattenBytesFunc = dataSourceType.FlattenBytesFunc(t.keyValueSeparator)
 
-	return tablewrapper.New(slogger, t.tableName, columns, t.generateExec)
+	return tablewrapper.New(flags, slogger, t.tableName, columns, t.generateExec)
 }
 
 func (t *Table) generateExec(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {