From 1debf88307951b23f06be36ca4448fc77b127fca Mon Sep 17 00:00:00 2001
From: David Gannon <19214156+dgannon991@users.noreply.github.com>
Date: Thu, 14 Nov 2024 21:10:23 +0000
Subject: [PATCH] Added optional cancelable context

Signed-off-by: David Gannon <19214156+dgannon991@users.noreply.github.com>
---
 action/operation_configs.go |  6 ++++++
 driver/docker/docker.go     | 33 +++++++++++++++++++--------------
 driver/driver.go            |  2 ++
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/action/operation_configs.go b/action/operation_configs.go
index 3dd48a56..6c839b5f 100644
--- a/action/operation_configs.go
+++ b/action/operation_configs.go
@@ -1,6 +1,7 @@
 package action
 
 import (
+	"context"
 	"os"
 
 	"github.com/cnabio/cnab-go/driver"
@@ -33,5 +34,10 @@ func (cfgs OperationConfigs) ApplyConfig(op *driver.Operation) error {
 		op.Err = os.Stderr
 	}
 
+	// if no context were provided
+	if op.CTX == nil {
+		op.CTX = context.Background()
+	}
+
 	return nil
 }
diff --git a/driver/docker/docker.go b/driver/docker/docker.go
index 33edcd2f..83f35201 100644
--- a/driver/docker/docker.go
+++ b/driver/docker/docker.go
@@ -182,7 +182,6 @@ func (d *Driver) initializeDockerCli() (command.Cli, error) {
 }
 
 func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
-	ctx := context.Background()
 
 	cli, err := d.initializeDockerCli()
 	if err != nil {
@@ -193,12 +192,12 @@ func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
 		return driver.OperationResult{}, nil
 	}
 	if d.config["PULL_ALWAYS"] == "1" {
-		if err := pullImage(ctx, cli, op.Image.Image); err != nil {
+		if err := pullImage(op.CTX, cli, op.Image.Image); err != nil {
 			return driver.OperationResult{}, err
 		}
 	}
 
-	ii, err := d.inspectImage(ctx, op.Image)
+	ii, err := d.inspectImage(op.CTX, op.Image)
 	if err != nil {
 		return driver.OperationResult{}, err
 	}
@@ -212,13 +211,13 @@ func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
 		return driver.OperationResult{}, err
 	}
 
-	resp, err := cli.Client().ContainerCreate(ctx, &d.containerCfg, &d.containerHostCfg, nil, nil, "")
+	resp, err := cli.Client().ContainerCreate(op.CTX, &d.containerCfg, &d.containerHostCfg, nil, nil, "")
 	if err != nil {
 		return driver.OperationResult{}, fmt.Errorf("cannot create container: %v", err)
 	}
 
 	if d.config["CLEANUP_CONTAINERS"] == "true" {
-		defer cli.Client().ContainerRemove(ctx, resp.ID, container.RemoveOptions{})
+		defer cli.Client().ContainerRemove(op.CTX, resp.ID, container.RemoveOptions{})
 	}
 
 	containerUID := getContainerUserID(ii.Config.User)
@@ -231,12 +230,12 @@ func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
 	}
 	// This copies the tar to the root of the container. The tar has been assembled using the
 	// path from the given file, starting at the /.
-	err = cli.Client().CopyToContainer(ctx, resp.ID, "/", tarContent, options)
+	err = cli.Client().CopyToContainer(op.CTX, resp.ID, "/", tarContent, options)
 	if err != nil {
 		return driver.OperationResult{}, fmt.Errorf("error copying to / in container: %s", err)
 	}
 
-	attach, err := cli.Client().ContainerAttach(ctx, resp.ID, container.AttachOptions{
+	attach, err := cli.Client().ContainerAttach(op.CTX, resp.ID, container.AttachOptions{
 		Stream: true,
 		Stdout: true,
 		Stderr: true,
@@ -269,28 +268,34 @@ func (d *Driver) exec(op *driver.Operation) (driver.OperationResult, error) {
 		}
 	}()
 
-	if err = cli.Client().ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil {
+	if err = cli.Client().ContainerStart(op.CTX, resp.ID, container.StartOptions{}); err != nil {
 		return driver.OperationResult{}, fmt.Errorf("cannot start container: %v", err)
 	}
-	statusc, errc := cli.Client().ContainerWait(ctx, resp.ID, container.WaitConditionNotRunning)
+	statusc, errc := cli.Client().ContainerWait(op.CTX, resp.ID, container.WaitConditionNotRunning)
 	select {
+	case <-op.CTX.Done():
+		err = cli.Client().ContainerStop(context.Background(), resp.ID, container.StopOptions{})
+		if err != nil {
+			return driver.OperationResult{}, err
+		}
+		return driver.OperationResult{}, op.CTX.Err()
 	case err := <-errc:
 		if err != nil {
-			opResult, fetchErr := d.fetchOutputs(ctx, resp.ID, op)
+			opResult, fetchErr := d.fetchOutputs(op.CTX, resp.ID, op)
 			return opResult, containerError("error in container", err, fetchErr)
 		}
 	case s := <-statusc:
 		if s.StatusCode == 0 {
-			return d.fetchOutputs(ctx, resp.ID, op)
+			return d.fetchOutputs(op.CTX, resp.ID, op)
 		}
 		if s.Error != nil {
-			opResult, fetchErr := d.fetchOutputs(ctx, resp.ID, op)
+			opResult, fetchErr := d.fetchOutputs(op.CTX, resp.ID, op)
 			return opResult, containerError(fmt.Sprintf("container exit code: %d, message", s.StatusCode), err, fetchErr)
 		}
-		opResult, fetchErr := d.fetchOutputs(ctx, resp.ID, op)
+		opResult, fetchErr := d.fetchOutputs(op.CTX, resp.ID, op)
 		return opResult, containerError(fmt.Sprintf("container exit code: %d, message", s.StatusCode), err, fetchErr)
 	}
-	opResult, fetchErr := d.fetchOutputs(ctx, resp.ID, op)
+	opResult, fetchErr := d.fetchOutputs(op.CTX, resp.ID, op)
 	if fetchErr != nil {
 		return opResult, fmt.Errorf("fetching outputs failed: %s", fetchErr)
 	}
diff --git a/driver/driver.go b/driver/driver.go
index 48e38ff9..9585b2df 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -1,6 +1,7 @@
 package driver
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -40,6 +41,7 @@ type Operation struct {
 	Err io.Writer `json:"-"`
 	// Bundle represents the bundle information for use by the operation
 	Bundle *bundle.Bundle
+	CTX    context.Context
 }
 
 // ResolvedCred is a credential that has been resolved and is ready for injection into the runtime.