Pass context to drivers during bundle operation

action/action.go

@@ -1,6 +1,7 @@
 package action
 
 import (
+	"context"
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
@@ -78,7 +79,7 @@ func (a Action) Run(c claim.Claim, creds credentials.Set, opCfgs ...OperationCon
 	}
 
 	var opErr *multierror.Error
-	opResult, err := a.Driver.Run(op)
+	opResult, err := a.Driver.Run(context.TODO(), op)
 	if err != nil {
 		opErr = multierror.Append(opErr, err)
 	}

action/action_test.go

@@ -1,6 +1,7 @@
 package action
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"io/ioutil"
@@ -32,7 +33,7 @@ type mockDriver struct {
 func (d *mockDriver) Handles(imageType string) bool {
 	return d.shouldHandle
 }
-func (d *mockDriver) Run(op *driver.Operation) (driver.OperationResult, error) {
+func (d *mockDriver) Run(context context.Context, op *driver.Operation) (driver.OperationResult, error) {
 	d.Operation = op
 	return d.Result, d.Error
 }

driver/command/command.go

@@ -2,6 +2,7 @@ package command
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -21,8 +22,8 @@ type Driver struct {
 }
 
 // Run executes the command
-func (d *Driver) Run(op *driver.Operation) (driver.OperationResult, error) {
-	return d.exec(op)
+func (d *Driver) Run(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
+	return d.exec(ctx, op)
 }
 
 // Handles executes the driver with `--handles` and parses the results
@@ -45,7 +46,7 @@ func (d *Driver) cliName() string {
 	return "cnab-" + strings.ToLower(d.Name)
 }
 
-func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
+func (d *Driver) exec(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
 	// We need to do two things here: We need to make it easier for the
 	// command to access data, and we need to make it easy for the command
 	// to pass that data on to the image it invokes. So we do some data
@@ -81,7 +82,7 @@ func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
 	}
 
 	args := []string{}
-	cmd := exec.Command(d.cliName(), args...)
+	cmd := exec.CommandContext(ctx, d.cliName(), args...)
 	cmd.Dir, err = os.Getwd()
 	if err != nil {
 		return driver.OperationResult{}, err

driver/command/command_nix_test.go

@@ -3,9 +3,13 @@
 package command
 
 import (
+	"bytes"
+	"context"
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/stretchr/testify/assert"
 
 	"github.com/cnabio/cnab-go/bundle"
@@ -58,7 +62,7 @@ func TestCommandDriverOutputs(t *testing.T) {
 				},
 			},
 		}
-		opResult, err := cmddriver.Run(&op)
+		opResult, err := cmddriver.Run(context.Background(), &op)
 		if err != nil {
 			t.Fatalf("Driver Run failed %v", err)
 		}
@@ -113,7 +117,7 @@ func TestCommandDriverOutputs(t *testing.T) {
 				},
 			},
 		}
-		_, err := cmddriver.Run(&op)
+		_, err := cmddriver.Run(context.Background(), &op)
 		assert.NoError(t, err)
 	}
 	CreateAndRunTestCommandDriver(t, name, content, testfunc)
@@ -163,7 +167,7 @@ func TestCommandDriverOutputs(t *testing.T) {
 				},
 			},
 		}
-		opResult, err := cmddriver.Run(&op)
+		opResult, err := cmddriver.Run(context.Background(), &op)
 		if err != nil {
 			t.Fatalf("Driver Run failed %v", err)
 		}
@@ -174,3 +178,40 @@ func TestCommandDriverOutputs(t *testing.T) {
 	}
 	CreateAndRunTestCommandDriver(t, name, content, testfunc)
 }
+
+func TestCommandDriverCancellation(t *testing.T) {
+	content := `#!/bin/sh
+		echo command executed
+	`
+	name := "test-command.sh"
+	output := bytes.Buffer{}
+	testfunc := func(t *testing.T, cmddriver *Driver) {
+		if !cmddriver.CheckDriverExists() {
+			t.Fatalf("Expected driver %s to exist Driver Name %s ", name, cmddriver.Name)
+		}
+		op := driver.Operation{
+			Action:       "install",
+			Installation: "test",
+			Parameters:   map[string]interface{}{},
+			Image: bundle.InvocationImage{
+				BaseImage: bundle.BaseImage{
+					Image:     "cnab/helloworld:latest",
+					ImageType: "docker",
+				},
+			},
+			Revision:    "01DDY0MT808KX0GGZ6SMXN4TW",
+			Environment: map[string]string{},
+			Files: map[string]string{
+				"/cnab/app/image-map.json": "{}",
+			},
+			Out:    &output,
+			Bundle: &bundle.Bundle{Name: "mybun"},
+		}
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := cmddriver.Run(ctx, &op)
+		require.EqualError(t, err, "Start of driver (test-command.sh) failed: context canceled")
+		assert.NotContains(t, output.String(), "command executed")
+	}
+	CreateAndRunTestCommandDriver(t, name, content, testfunc)
+}

driver/command/command_test.go

@@ -28,6 +28,7 @@ func TestCheckDriverExists(t *testing.T) {
 	}
 	CreateAndRunTestCommandDriver(t, name, "", testfunc)
 }
+
 func CreateAndRunTestCommandDriver(t *testing.T, name string, content string, testfunc func(t *testing.T, d *Driver)) {
 	cmddriver := &Driver{Name: name}
 	dirname, err := ioutil.TempDir("", "cnab")

driver/debug/debug.go

@@ -1,6 +1,7 @@
 package debug
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 
@@ -15,18 +16,24 @@ type Driver struct {
 }
 
 // Run executes the operation on the Debug driver
-func (d *Driver) Run(op *driver.Operation) (driver.OperationResult, error) {
-	data, err := json.MarshalIndent(op, "", "  ")
-	if err != nil {
-		return driver.OperationResult{}, err
-	}
+func (d *Driver) Run(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
+	select {
+	case <-ctx.Done():
+		return driver.OperationResult{}, ctx.Err()
+	default:
+
+		data, err := json.MarshalIndent(op, "", "  ")
+		if err != nil {
+			return driver.OperationResult{}, err
+		}
 
-	result := driver.OperationResult{}
-	result.Logs.Write(data)
+		result := driver.OperationResult{}
+		result.Logs.Write(data)
 
-	fmt.Fprintln(op.Out, result.Logs.String())
+		fmt.Fprintln(op.Out, result.Logs.String())
 
-	return result, nil
+		return result, nil
+	}
 }
 
 // Handles always returns true, effectively claiming to work for any image type

driver/debug/debug_test.go

@@ -1,6 +1,8 @@
 package debug
 
 import (
+	"bytes"
+	"context"
 	"io/ioutil"
 	"testing"
 
@@ -25,17 +27,31 @@ func TestDebugDriver_Run(t *testing.T) {
 	is := assert.New(t)
 	is.NotNil(d)
 
-	op := &driver.Operation{
+	op := driver.Operation{
 		Installation: "test",
 		Image: bundle.InvocationImage{
 			BaseImage: bundle.BaseImage{
 				Image:     "test:1.2.3",
 				ImageType: "oci",
 			},
 		},
-		Out: ioutil.Discard,
 	}
 
-	_, err := d.Run(op)
-	is.NoError(err)
+	t.Run("success", func(t *testing.T) {
+		op.Out = ioutil.Discard
+
+		_, err := d.Run(context.Background(), &op)
+		is.NoError(err)
+	})
+
+	t.Run("cancelled", func(t *testing.T) {
+		output := bytes.Buffer{}
+		op.Out = &output
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := d.Run(ctx, &op)
+		is.Empty(output.String())
+		is.EqualError(err, "context canceled")
+	})
 }

driver/docker/docker.go

@@ -39,8 +39,8 @@ type Driver struct {
 }
 
 // Run executes the Docker driver
-func (d *Driver) Run(op *driver.Operation) (driver.OperationResult, error) {
-	return d.exec(op)
+func (d *Driver) Run(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
+	return d.exec(ctx, op)
 }
 
 // Handles indicates that the Docker driver supports "docker" and "oci"
@@ -171,9 +171,7 @@ func (d *Driver) initializeDockerCli() (command.Cli, error) {
 	return cli, nil
 }
 
-func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
-	ctx := context.Background()
-
+func (d *Driver) exec(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
 	cli, err := d.initializeDockerCli()
 	if err != nil {
 		return driver.OperationResult{}, err

driver/docker/docker_integration_test.go

@@ -4,9 +4,12 @@ package docker
 
 import (
 	"bytes"
+	"context"
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/stretchr/testify/assert"
 
 	"github.com/cnabio/cnab-go/bundle"
@@ -65,8 +68,8 @@ func TestDriver_Run(t *testing.T) {
 	}
 
 	docker := &Driver{}
-	docker.SetContainerOut(op.Out) // Docker driver writes container stdout to driver.containerOut.
-	opResult, err := docker.Run(op)
+	docker.SetContainerOut(&output) // Docker driver writes container stdout to driver.containerOut.
+	opResult, err := docker.Run(context.Background(), op)
 
 	assert.NoError(t, err)
 	assert.Equal(t, "Install action\nAction install complete for example\n", output.String())
@@ -75,4 +78,66 @@ func TestDriver_Run(t *testing.T) {
 		"output1": "SOME INSTALL CONTENT 1\n",
 		"output2": "SOME INSTALL CONTENT 2\n",
 	}, opResult.Outputs)
+	assert.Contains(t, output.String(), "Install action", "expected text from the install action to be captured")
+}
+
+func TestDriver_RunCancelled(t *testing.T) {
+	imageFromEnv, ok := os.LookupEnv("DOCKER_INTEGRATION_TEST_IMAGE")
+	var image bundle.InvocationImage
+
+	if ok {
+		image = bundle.InvocationImage{
+			BaseImage: bundle.BaseImage{
+				Image: imageFromEnv,
+			},
+		}
+	} else {
+		image = bundle.InvocationImage{
+			BaseImage: bundle.BaseImage{
+				Image:  "pvtlmc/example-outputs",
+				Digest: "sha256:568461508c8d220742add8abd226b33534d4269868df4b3178fae1cba3818a6e",
+			},
+		}
+	}
+
+	op := &driver.Operation{
+		Installation: "example",
+		Action:       "install",
+		Image:        image,
+		Outputs: map[string]string{
+			"/cnab/app/outputs/output1": "output1",
+			"/cnab/app/outputs/output2": "output2",
+		},
+		Bundle: &bundle.Bundle{
+			Definitions: definition.Definitions{
+				"output1": &definition.Schema{},
+				"output2": &definition.Schema{},
+			},
+			Outputs: map[string]bundle.Output{
+				"output1": {
+					Definition: "output1",
+				},
+				"output2": {
+					Definition: "output2",
+				},
+			},
+		},
+	}
+
+	var output bytes.Buffer
+	op.Out = &output
+	op.Environment = map[string]string{
+		"CNAB_ACTION":            op.Action,
+		"CNAB_INSTALLATION_NAME": op.Installation,
+	}
+
+	docker := &Driver{}
+	docker.SetContainerOut(op.Out) // Docker driver writes container stdout to driver.containerOut.
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err := docker.Run(ctx, op)
+	require.Error(t, err, "expected an error")
+	assert.Contains(t, err.Error(), "context canceled")
+	assert.Empty(t, output.String(), "expected the driver to not output anything")
 }

driver/driver.go

@@ -2,6 +2,7 @@ package driver
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 
@@ -90,7 +91,7 @@ func (r *OperationResult) SetDefaultOutputValues(op Operation) error {
 // Driver is capable of running a invocation image
 type Driver interface {
 	// Run executes the operation inside of the invocation image
-	Run(*Operation) (OperationResult, error)
+	Run(context.Context, *Operation) (OperationResult, error)
 	// Handles receives an ImageType* and answers whether this driver supports that type
 	Handles(string) bool
 }

driver/kubernetes/kubernetes.go

@@ -1,6 +1,7 @@
 package kubernetes
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -131,7 +132,8 @@ func (k *Driver) setClient(conf *rest.Config) error {
 }
 
 // Run executes the operation inside of the invocation image.
-func (k *Driver) Run(op *driver.Operation) (driver.OperationResult, error) {
+func (k *Driver) Run(ctx context.Context, op *driver.Operation) (driver.OperationResult, error) {
+	// TODO: use passed context to handle cancellation and timeouts https://github.com/cnabio/cnab-go/issues/220
 	if k.Namespace == "" {
 		return driver.OperationResult{}, fmt.Errorf("KUBE_NAMESPACE is required")
 	}