Set context timeouts on all exec functions (#688)

Add context timeouts to all table execs
kolide · Dec 22, 2020 · b77a613 · b77a613
1 parent ae4210c
commit b77a613
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 14 deletions.
diff --git a/pkg/osquery/table/mdfind_darwin.go b/pkg/osquery/table/mdfind_darwin.go
@@ -12,7 +12,7 @@ import (
 )
 
 func mdfind(args ...string) ([]string, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
 	path := "/usr/bin/mdfind"

diff --git a/pkg/osquery/table/mdm.go b/pkg/osquery/table/mdm.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"os/exec"
 	"strconv"
+	"time"
 
 	"github.com/go-kit/kit/log"
 	"github.com/groob/plist"
@@ -31,13 +32,13 @@ func MDMInfo(logger log.Logger) *table.Plugin {
 }
 
 func generateMDMInfo(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
-	profiles, err := getMDMProfile()
+	profiles, err := getMDMProfile(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	depEnrolled, userApproved := "unknown", "unknown"
-	status, err := getMDMProfileStatus()
+	status, err := getMDMProfileStatus(ctx)
 	if err == nil { // only supported on 10.13.4+
 		depEnrolled = strconv.FormatBool(status.DEPEnrolled)
 		userApproved = strconv.FormatBool(status.UserApproved)
@@ -84,8 +85,11 @@ func generateMDMInfo(ctx context.Context, queryContext table.QueryContext) ([]ma
 	return results, nil
 }
 
-func getMDMProfile() (*profilesOutput, error) {
-	cmd := exec.Command("/usr/bin/profiles", "-L", "-o", "stdout-xml")
+func getMDMProfile(ctx context.Context) (*profilesOutput, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	cmd := exec.CommandContext(ctx, "/usr/bin/profiles", "-L", "-o", "stdout-xml")
 	out, err := cmd.Output()
 	if err != nil {
 		return nil, errors.Wrap(err, "calling /usr/bin/profiles to get MDM profile payload")
@@ -124,8 +128,11 @@ type payloadContent struct {
 	SignMessage             bool
 }
 
-func getMDMProfileStatus() (profileStatus, error) {
-	cmd := exec.Command("/usr/bin/profiles", "status", "-type", "enrollment")
+func getMDMProfileStatus(ctx context.Context) (profileStatus, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	cmd := exec.CommandContext(ctx, "/usr/bin/profiles", "status", "-type", "enrollment")
 	out, err := cmd.Output()
 	if err != nil {
 		return profileStatus{}, errors.Wrap(err, "calling /usr/bin/profiles to get MDM profile status")

diff --git a/pkg/osquery/table/mdm_test.go b/pkg/osquery/table/mdm_test.go
@@ -3,6 +3,7 @@
 package table
 
 import (
+	"context"
 	"testing"
 
 	"github.com/kolide/kit/env"
@@ -11,9 +12,9 @@ import (
 
 func TestMDMProfileStatus(t *testing.T) {
 	if env.Bool("SKIP_TEST_MDM", true) {
-		t.Skip("No docker")
+		t.Skip("Skipping MDM Test")
 	}
 
-	_, err := getMDMProfileStatus()
+	_, err := getMDMProfileStatus(context.TODO())
 	require.Nil(t, err)
 }
diff --git a/pkg/osquery/table/touchid_system_darwin.go b/pkg/osquery/table/touchid_system_darwin.go
@@ -6,6 +6,7 @@ import (
 	"os/exec"
 	"regexp"
 	"strings"
+	"time"
 
 	"github.com/go-kit/kit/log"
 	"github.com/pkg/errors"
@@ -44,6 +45,9 @@ type touchIDSystemConfig struct {
 
 // TouchIDSystemConfigGenerate will be called whenever the table is queried.
 func (t *touchIDSystemConfigTable) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
 	var results []map[string]string
 	var touchIDCompatible, secureEnclaveCPU, touchIDEnabled, touchIDUnlock string
 

diff --git a/pkg/osquery/table/touchid_user_darwin.go b/pkg/osquery/table/touchid_user_darwin.go
@@ -8,6 +8,7 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
@@ -141,6 +142,9 @@ func (t *touchIDUserConfigTable) generate(ctx context.Context, queryContext tabl
 
 // runCommand runs a given command and arguments as the supplied user
 func runCommandContext(ctx context.Context, uid int, cmd string, args ...string) (string, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
 	// Set up the command
 	var stdout bytes.Buffer
 	c := exec.CommandContext(ctx, cmd, args...)

diff --git a/pkg/osquery/tables/dataflattentable/exec.go b/pkg/osquery/tables/dataflattentable/exec.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"os/exec"
 	"strings"
+	"time"
 
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
@@ -57,6 +58,9 @@ func (t *Table) generateExec(ctx context.Context, queryContext table.QueryContex
 }
 
 func (t *Table) exec(ctx context.Context) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(ctx, 50*time.Second)
+	defer cancel()
+
 	var stdout bytes.Buffer
 	var stderr bytes.Buffer
 

diff --git a/pkg/osquery/tables/filevault/filevault.go b/pkg/osquery/tables/filevault/filevault.go
@@ -36,22 +36,28 @@ func TablePlugin(client *osquery.ExtensionManagerClient, logger log.Logger) *tab
 }
 
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	var results []map[string]string
 
 	// Read the system's fdesetup configuration
 	var stdout bytes.Buffer
+
 	cmd := exec.CommandContext(ctx, fdesetupPath, "status")
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
 	cmd.Stdout = &stdout
+
 	if err := cmd.Run(); err != nil {
 		return nil, errors.Wrap(err, "calling fdesetup")
 	}
+
 	output := string(stdout.Bytes())
 	status := strings.TrimSuffix(output, "\n")
+
 	result := map[string]string{
 		"status": status,
 	}
+
 	results = append(results, result)
 
 	return results, nil

diff --git a/pkg/osquery/tables/ioreg/ioreg.go b/pkg/osquery/tables/ioreg/ioreg.go
@@ -157,14 +157,14 @@ func (t *Table) flattenOutput(dataQuery string, systemOutput []byte) ([]dataflat
 }
 
 func (t *Table) execIoreg(ctx context.Context, args []string) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
 	var stdout bytes.Buffer
 	var stderr bytes.Buffer
 
 	args = append(args, "-a")
 
-	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
-
 	cmd := exec.CommandContext(ctx, ioregPath, args...)
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr

diff --git a/pkg/osquery/tables/systemprofiler/systemprofiler.go b/pkg/osquery/tables/systemprofiler/systemprofiler.go
@@ -41,6 +41,7 @@ import (
 	"context"
 	"os/exec"
 	"strings"
+	"time"
 
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
@@ -197,6 +198,13 @@ func (t *Table) getRowsFromOutput(dataQuery, detailLevel string, systemProfilerO
 }
 
 func (t *Table) execSystemProfiler(ctx context.Context, detailLevel string, subcommands []string) ([]byte, error) {
+	timeout := 30 * time.Second
+	if detailLevel == "full" {
+		timeout = 5 * time.Minute
+	}
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	var stdout bytes.Buffer
 	var stderr bytes.Buffer