Add context-aware versions of several functions and some methods

Introduce versions of a number of functions and some methods which take
a context.Context, to least try to provide proper cancellation support.

When copier.Get() returns an error, also write a non-NUL string after
whatever it's already written, so that a tar.Reader that reads what it
output will also flag the malformed header with an unexpected error.

In main(), connect the global default context to use a timeout set by an
experimental timeout flag, and to be canceled on the interrupt signal.

Signed-off-by: Nalin Dahyabhai <nalin@redhat.com>

Author: nalind

add.go

@@ -26,9 +26,10 @@ import (
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/sirupsen/logrus"
-	"github.com/tonistiigi/dchapes-mode"
+	mode "github.com/tonistiigi/dchapes-mode"
 	"go.podman.io/buildah/copier"
 	"go.podman.io/buildah/define"
+	"go.podman.io/buildah/internal/ctxreader"
 	"go.podman.io/buildah/internal/tmpdir"
 	"go.podman.io/buildah/internal/urlsource"
 	"go.podman.io/buildah/pkg/chrootuser"
@@ -129,7 +130,7 @@ type AddAndCopyOptions struct {
 }
 
 // getURL writes a tar archive containing the named content
-func getURL(src string, chown *idtools.IDPair, mountpoint, renameTarget string, writer io.Writer, chmod string, srcDigest digest.Digest, certPath string, insecureSkipTLSVerify types.OptionalBool, timestamp *time.Time) error {
+func getURL(ctx context.Context, src string, chown *idtools.IDPair, mountpoint, renameTarget string, writer io.Writer, chmod string, srcDigest digest.Digest, certPath string, insecureSkipTLSVerify types.OptionalBool, timestamp *time.Time) error {
 	url, err := url.Parse(src)
 	if err != nil {
 		return err
@@ -152,7 +153,11 @@ func getURL(src string, chown *idtools.IDPair, mountpoint, renameTarget string,
 		Proxy:           http.ProxyFromEnvironment,
 	}
 	httpClient := &http.Client{Transport: tr}
-	response, err := httpClient.Get(src)
+	req, err := http.NewRequestWithContext(ctx, "GET", src, nil)
+	if err != nil {
+		return err
+	}
+	response, err := httpClient.Do(req)
 	if err != nil {
 		return err
 	}
@@ -184,7 +189,7 @@ func getURL(src string, chown *idtools.IDPair, mountpoint, renameTarget string,
 	}
 	// Figure out the size of the content.
 	size := response.ContentLength
-	var responseBody io.Reader = response.Body
+	responseBody := ctxreader.NewCancelableReader(ctx, response.Body)
 	if size < 0 {
 		// Create a temporary file and copy the content to it, so that
 		// we can figure out how much content there is.
@@ -307,10 +312,15 @@ func getParentsPrefixToRemoveAndParentsToSkip(pattern string, contextDir string)
 	return prefix, out
 }
 
-// Add copies the contents of the specified sources into the container's root
+// Add() calls AddContext() with context.Background().
+func (b *Builder) Add(destination string, extract bool, options AddAndCopyOptions, sources ...string) error {
+	return b.AddContext(context.Background(), destination, extract, options, sources...)
+}
+
+// AddContext copies the contents of the specified sources into the container's root
 // filesystem, optionally extracting contents of local files that look like
 // non-empty archives.
-func (b *Builder) Add(destination string, extract bool, options AddAndCopyOptions, sources ...string) error {
+func (b *Builder) AddContext(ctx context.Context, destination string, extract bool, options AddAndCopyOptions, sources ...string) error {
 	mountPoint, err := b.Mount(b.MountLabel)
 	if err != nil {
 		return err
@@ -374,7 +384,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 			DisallowWildcard:   options.AllowWildcard == types.OptionalBoolFalse,
 			AllowEmptyWildcard: options.AllowEmptyWildcard == types.OptionalBoolTrue,
 		}
-		localSourceStats, err = copier.Stat(contextDir, contextDir, statOptions, localSources)
+		localSourceStats, err = copier.StatContext(ctx, contextDir, contextDir, statOptions, localSources)
 		if err != nil {
 			return fmt.Errorf("checking on sources under %q: %w", contextDir, err)
 		}
@@ -466,7 +476,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 	statOptions := copier.StatOptions{
 		CheckForArchives: extract,
 	}
-	destStats, err := copier.Stat(mountPoint, filepath.Join(mountPoint, b.WorkDir()), statOptions, []string{extractDirectory})
+	destStats, err := copier.StatContext(ctx, mountPoint, filepath.Join(mountPoint, b.WorkDir()), statOptions, []string{extractDirectory})
 	if err != nil {
 		return fmt.Errorf("checking on destination %v: %w", extractDirectory, err)
 	}
@@ -559,7 +569,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 		if !strings.HasPrefix(putDirAbs, stagingDirAbs+string(os.PathSeparator)) && putDirAbs != stagingDirAbs {
 			return fmt.Errorf("destination path %q escapes staging directory", destination)
 		}
-		if err := copier.Mkdir(putRoot, putDirAbs, mkdirOptions); err != nil {
+		if err := copier.MkdirContext(ctx, putRoot, putDirAbs, mkdirOptions); err != nil {
 			return fmt.Errorf("ensuring target directory exists: %w", err)
 		}
 		tempPath := putDir
@@ -570,7 +580,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 			tempPath = filepath.Dir(tempPath)
 		}
 	} else {
-		if err := copier.Mkdir(mountPoint, extractDirectory, mkdirOptions); err != nil {
+		if err := copier.MkdirContext(ctx, mountPoint, extractDirectory, mkdirOptions); err != nil {
 			return fmt.Errorf("ensuring target directory exists: %w", err)
 		}
 
@@ -599,7 +609,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 					defer wg.Done()
 					defer pipeWriter.Close()
 					var cloneDir, subdir string
-					cloneDir, subdir, getErr = define.TempDirForURL(tmpdir.GetTempDir(), "", src)
+					cloneDir, subdir, getErr = define.TempDirForURLContext(ctx, tmpdir.GetTempDir(), "", src)
 					if getErr != nil {
 						return
 					}
@@ -621,12 +631,12 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 					}
 					writer := io.WriteCloser(pipeWriter)
 					repositoryDir := filepath.Join(cloneDir, subdir)
-					getErr = copier.Get(repositoryDir, repositoryDir, getOptions, []string{"."}, writer)
+					getErr = copier.GetContext(ctx, repositoryDir, repositoryDir, getOptions, []string{"."}, writer)
 				}()
 			} else {
 				go func() {
-					getErr = retry.IfNecessary(context.TODO(), func() error {
-						return getURL(src, chownFiles, mountPoint, renameTarget, pipeWriter, options.Chmod, srcDigest, options.CertPath, options.InsecureSkipTLSVerify, options.Timestamp)
+					getErr = retry.IfNecessary(ctx, func() error {
+						return getURL(ctx, src, chownFiles, mountPoint, renameTarget, pipeWriter, options.Chmod, srcDigest, options.CertPath, options.InsecureSkipTLSVerify, options.Timestamp)
 					}, &retry.Options{
 						MaxRetry: options.MaxRetries,
 						Delay:    options.RetryDelay,
@@ -656,7 +666,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 						IgnoreDevices: userns.RunningInUserNS(),
 						Timestamp:     options.Timestamp,
 					}
-					putErr = copier.Put(putRoot, putDir, putOptions, io.TeeReader(pipeReader, hasher))
+					putErr = copier.PutContext(ctx, putRoot, putDir, putOptions, io.TeeReader(pipeReader, hasher))
 				}
 				hashCloser.Close()
 				pipeReader.Close()
@@ -788,7 +798,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 					AllowEmptyWildcard: options.AllowEmptyWildcard == types.OptionalBoolTrue,
 					NoDerefSymlinks:    options.FollowSymlink == types.OptionalBoolFalse,
 				}
-				getErr = copier.Get(contextDir, contextDir, getOptions, []string{globbedToGlobbable(globbed)}, writer)
+				getErr = copier.GetContext(ctx, contextDir, contextDir, getOptions, []string{globbedToGlobbable(globbed)}, writer)
 				closeErr = writer.Close()
 				if renameTarget != "" && renamedItems > 1 {
 					renameErr = fmt.Errorf("internal error: renamed %d items when we expected to only rename 1", renamedItems)
@@ -820,7 +830,7 @@ func (b *Builder) Add(destination string, extract bool, options AddAndCopyOption
 						IgnoreDevices:   userns.RunningInUserNS(),
 						Timestamp:       options.Timestamp,
 					}
-					putErr = copier.Put(putRoot, putDir, putOptions, io.TeeReader(pipeReader, hasher))
+					putErr = copier.PutContext(ctx, putRoot, putDir, putOptions, io.TeeReader(pipeReader, hasher))
 				}
 				hashCloser.Close()
 				pipeReader.Close()

buildah_test.go

@@ -1,7 +1,6 @@
 package buildah
 
 import (
-	"context"
 	"flag"
 	"os"
 	"testing"
@@ -46,7 +45,7 @@ func TestOpenBuilderCommonBuildOpts(t *testing.T) {
 	// or builder must enable sometime of locking mechanism i.e if
 	// routine is creating Builder other's must wait for it.
 	// Tracked here: https://github.com/containers/buildah/issues/5967
-	ctx := context.TODO()
+	ctx := t.Context()
 	store, err := storage.GetStore(types.StoreOptions{
 		RunRoot:         t.TempDir(),
 		GraphRoot:       t.TempDir(),

chroot/run_common.go

@@ -4,6 +4,7 @@ package chroot
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -51,10 +52,15 @@ type runUsingChrootExecSubprocOptions struct {
 	NoPivot    bool
 }
 
-// RunUsingChroot runs a chrooted process, using some of the settings from the
+// RunUsingChroot() calls RunUsingChrootContext() with context.Background().
+func RunUsingChroot(spec *specs.Spec, bundlePath, homeDir string, stdin io.Reader, stdout, stderr io.Writer, noPivot bool) (err error) {
+	return RunUsingChrootContext(context.Background(), spec, bundlePath, homeDir, stdin, stdout, stderr, noPivot)
+}
+
+// RunUsingChrootContext runs a chrooted process, using some of the settings from the
 // passed-in spec, and using the specified bundlePath to hold temporary files,
 // directories, and mountpoints.
-func RunUsingChroot(spec *specs.Spec, bundlePath, homeDir string, stdin io.Reader, stdout, stderr io.Writer, noPivot bool) (err error) {
+func RunUsingChrootContext(ctx context.Context, spec *specs.Spec, bundlePath, homeDir string, stdin io.Reader, stdout, stderr io.Writer, noPivot bool) (err error) {
 	var confwg sync.WaitGroup
 	var homeFound bool
 	for _, env := range spec.Process.Env {
@@ -127,6 +133,7 @@ func RunUsingChroot(spec *specs.Spec, bundlePath, homeDir string, stdin io.Reade
 
 	// Start the grandparent subprocess.
 	cmd := unshare.Command(runUsingChrootCommand)
+	cmd.Cmd = reexec.CommandContext(ctx, runUsingChrootCommand) // TODO: add an unshare.CommandContext()
 	setPdeathsig(cmd.Cmd)
 	cmd.Stdin, cmd.Stdout, cmd.Stderr = stdin, stdout, stderr
 	cmd.Dir = "/"

chroot/unsupported.go

@@ -3,6 +3,7 @@
 package chroot
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -11,5 +12,10 @@ import (
 
 // RunUsingChroot is not supported.
 func RunUsingChroot(spec *specs.Spec, bundlePath, homeDir string, stdin io.Reader, stdout, stderr io.Writer) (err error) {
+	return RunUsingChrootContext(context.Background(), spec, bundlePath, homeDir, stdin, stdout, stderr)
+}
+
+// RunUsingChrootContext is not supported.
+func RunUsingChrootContext(ctx context.Context, spec *specs.Spec, bundlePath, homeDir string, stdin io.Reader, stdout, stderr io.Writer) (err error) {
 	return fmt.Errorf("--isolation chroot is not supported on this platform")
 }

cmd/buildah/addcopy.go

@@ -322,7 +322,7 @@ func addAndCopyCmd(c *cobra.Command, args []string, verb string, iopts addCopyRe
 	}
 
 	extractLocalArchives := verb == "ADD"
-	err = builder.Add(dest, extractLocalArchives, options, args...)
+	err = builder.AddContext(getContext(), dest, extractLocalArchives, options, args...)
 	if err != nil {
 		return fmt.Errorf("adding content to container %q: %w", builder.Container, err)
 	}

cmd/buildah/common.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/signal"
+	"sync"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -157,11 +159,31 @@ func openImage(ctx context.Context, sc *types.SystemContext, store storage.Store
 	return builder, nil
 }
 
-// getContext returns a context.TODO
+// getContext returns a context that may have a timeout, and which cancels if
+// it receives os.Interrupt
 func getContext() context.Context {
-	return context.TODO()
+	ctx, _ := getContextWithCancel()
+	return ctx
 }
 
+func getContextCancel() context.CancelFunc {
+	_, cancel := getContextWithCancel()
+	return cancel
+}
+
+var getContextWithCancel = sync.OnceValues(func() (context.Context, context.CancelFunc) {
+	var ctx context.Context
+	var cancel1, cancel2 func()
+	if rootCmd.PersistentFlags().Changed("experimental-timeout") {
+		ctx, cancel1 = context.WithTimeout(context.Background(), globalFlagResults.ExperimentalTimeout)
+	} else {
+		ctx = context.Background()
+		cancel1 = func() {}
+	}
+	ctx, cancel2 = signal.NotifyContext(ctx, os.Interrupt)
+	return ctx, func() { cancel1(); cancel2() }
+})
+
 func getUserFlags() pflag.FlagSet {
 	fs := pflag.FlagSet{}
 	fs.String("user", "", "`user[:group]` to run the command as")

cmd/buildah/from.go

@@ -129,7 +129,7 @@ func onBuild(builder *buildah.Builder, quiet bool) error {
 				dest = args[size-1]
 				args = args[:size-1]
 			}
-			if err := builder.Add(dest, command == "ADD", buildah.AddAndCopyOptions{}, args...); err != nil {
+			if err := builder.AddContext(getContext(), dest, command == "ADD", buildah.AddAndCopyOptions{}, args...); err != nil {
 				return err
 			}
 		case "ANNOTATION":

cmd/buildah/images.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -133,7 +132,7 @@ func imagesCmd(c *cobra.Command, args []string, iopts *imageResults) error {
 		return err
 	}
 
-	ctx := context.Background()
+	ctx := getContext()
 
 	options := &libimage.ListImagesOptions{}
 	if len(iopts.filter) > 0 {

cmd/buildah/main.go

@@ -9,6 +9,7 @@ import (
 	"runtime/pprof"
 	"strings"
 	"syscall"
+	"time"
 
 	ispecs "github.com/opencontainers/image-spec/specs-go"
 	rspecs "github.com/opencontainers/runtime-spec/specs-go"
@@ -58,6 +59,7 @@ type globalFlags struct {
 	MemoryProfile              string
 	UserShortNameAliasConfPath string
 	CgroupManager              string
+	ExperimentalTimeout        time.Duration
 }
 
 var rootCmd = &cobra.Command{
@@ -107,6 +109,7 @@ func mainInit() {
 	rootCmd.CompletionOptions.HiddenDefaultCmd = true
 	// rootCmd.TraverseChildren = true
 	rootCmd.Version = fmt.Sprintf("%s (image-spec %s, runtime-spec %s)", define.Version, ispecs.Version, rspecs.Version)
+	rootCmd.PersistentFlags().DurationVar(&globalFlagResults.ExperimentalTimeout, "experimental-timeout", 0, "global timeout") // maybe rename this at some point?
 	rootCmd.PersistentFlags().BoolVar(&globalFlagResults.Debug, "debug", false, "print debugging information")
 	// TODO Need to allow for environment variable
 	rootCmd.PersistentFlags().StringVar(&globalFlagResults.RegistriesConf, "registries-conf", "", "path to registries.conf file (not usually used)")
@@ -138,6 +141,9 @@ func mainInit() {
 	if err := rootCmd.PersistentFlags().MarkHidden("memory-profile"); err != nil {
 		logrus.Fatalf("unable to mark memory-profile flag as hidden: %v", err)
 	}
+	if err := rootCmd.PersistentFlags().MarkHidden("experimental-timeout"); err != nil {
+		logrus.Fatalf("unable to mark experimental-timeout flag as hidden: %v", err)
+	}
 	rootCmd.AddGroup(commandGroups...)
 }
 
@@ -293,7 +299,9 @@ func main() {
 	// Hard code TMPDIR functions to use $TMPDIR or /var/tmp
 	os.Setenv("TMPDIR", parse.GetTempDir())
 
-	if err := rootCmd.Execute(); err != nil {
+	err := rootCmd.Execute()
+	getContextCancel()()
+	if err != nil {
 		if logrus.IsLevelEnabled(logrus.TraceLevel) {
 			fmt.Fprintf(os.Stderr, "Error: %+v\n", err)
 		} else {

cmd/buildah/manifest.go

@@ -741,7 +741,7 @@ func manifestRmCmd(c *cobra.Command, args []string) error {
 		Filters:        []string{"readonly=false"},
 		LookupManifest: true,
 	}
-	rmiReports, rmiErrors := runtime.RemoveImages(context.Background(), args, options)
+	rmiReports, rmiErrors := runtime.RemoveImages(getContext(), args, options)
 	for _, r := range rmiReports {
 		for _, u := range r.Untagged {
 			fmt.Printf("untagged: %s\n", u)