refactor: add context in packager (#2597)

## Description

This change adds contexts to packager functions to enable them to be
cancelled.

## Related Issue

Relates to #2594

## Checklist before merging

- [x] Test, docs, adr added or updated as needed
- [x] [Contributor Guide
Steps](https://github.com/defenseunicorns/zarf/blob/main/.github/CONTRIBUTING.md#developer-workflow)
followed

Co-authored-by: Austin Abro <37223396+AustinAbro321@users.noreply.github.com>

src/cmd/dev.go

@@ -204,7 +204,7 @@ var devFindImagesCmd = &cobra.Command{
 	Args:    cobra.MaximumNArgs(1),
 	Short:   lang.CmdDevFindImagesShort,
 	Long:    lang.CmdDevFindImagesLong,
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.CreateOpts.BaseDir = common.SetBaseDirectory(args)
 
 		v := common.GetViper()
@@ -216,7 +216,7 @@ var devFindImagesCmd = &cobra.Command{
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		if _, err := pkgClient.FindImages(); err != nil {
+		if _, err := pkgClient.FindImages(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdDevFindImagesErr, err.Error())
 		}
 	},
@@ -249,12 +249,12 @@ var devLintCmd = &cobra.Command{
 	Aliases: []string{"l"},
 	Short:   lang.CmdDevLintShort,
 	Long:    lang.CmdDevLintLong,
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.CreateOpts.BaseDir = common.SetBaseDirectory(args)
 		v := common.GetViper()
 		pkgConfig.CreateOpts.SetVariables = helpers.TransformAndMergeMap(
 			v.GetStringMapString(common.VPkgCreateSet), pkgConfig.CreateOpts.SetVariables, strings.ToUpper)
-		validator, err := lint.Validate(pkgConfig.CreateOpts)
+		validator, err := lint.Validate(cmd.Context(), pkgConfig.CreateOpts)
 		if err != nil {
 			message.Fatal(err, err.Error())
 		}

src/cmd/initialize.go

@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -49,7 +50,7 @@ var initCmd = &cobra.Command{
 
 		// Try to use an init-package in the executable directory if none exist in current working directory
 		var err error
-		if pkgConfig.PkgOpts.PackageSource, err = findInitPackage(initPackageName); err != nil {
+		if pkgConfig.PkgOpts.PackageSource, err = findInitPackage(cmd.Context(), initPackageName); err != nil {
 			message.Fatal(err, err.Error())
 		}
 
@@ -74,7 +75,7 @@ var initCmd = &cobra.Command{
 	},
 }
 
-func findInitPackage(initPackageName string) (string, error) {
+func findInitPackage(ctx context.Context, initPackageName string) (string, error) {
 	// First, look for the init package in the current working directory
 	if !helpers.InvalidPath(initPackageName) {
 		return initPackageName, nil
@@ -103,7 +104,7 @@ func findInitPackage(initPackageName string) (string, error) {
 	}
 
 	// Finally, if the init-package doesn't exist in the cache directory, suggest downloading it
-	downloadCacheTarget, err := downloadInitPackage(config.GetAbsCachePath())
+	downloadCacheTarget, err := downloadInitPackage(ctx, config.GetAbsCachePath())
 	if err != nil {
 		if errors.Is(err, lang.ErrInitNotFound) {
 			message.Fatal(err, err.Error())
@@ -114,7 +115,7 @@ func findInitPackage(initPackageName string) (string, error) {
 	return downloadCacheTarget, nil
 }
 
-func downloadInitPackage(cacheDirectory string) (string, error) {
+func downloadInitPackage(ctx context.Context, cacheDirectory string) (string, error) {
 	if config.CommonOptions.Confirm {
 		return "", lang.ErrInitNotFound
 	}
@@ -144,7 +145,7 @@ func downloadInitPackage(cacheDirectory string) (string, error) {
 			return "", err
 		}
 		source := &sources.OCISource{Remote: remote}
-		return source.Collect(cacheDirectory)
+		return source.Collect(ctx, cacheDirectory)
 	}
 	// Otherwise, exit and tell the user to manually download the init-package
 	return "", errors.New(lang.CmdInitPullErrManual)

src/cmd/package.go

@@ -39,7 +39,7 @@ var packageCreateCmd = &cobra.Command{
 	Args:    cobra.MaximumNArgs(1),
 	Short:   lang.CmdPackageCreateShort,
 	Long:    lang.CmdPackageCreateLong,
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.CreateOpts.BaseDir = common.SetBaseDirectory(args)
 
 		var isCleanPathRegex = regexp.MustCompile(`^[a-zA-Z0-9\_\-\/\.\~\\:]+$`)
@@ -55,7 +55,7 @@ var packageCreateCmd = &cobra.Command{
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		if err := pkgClient.Create(); err != nil {
+		if err := pkgClient.Create(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdPackageCreateErr, err.Error())
 		}
 	},
@@ -112,15 +112,15 @@ var packageInspectCmd = &cobra.Command{
 	Short:   lang.CmdPackageInspectShort,
 	Long:    lang.CmdPackageInspectLong,
 	Args:    cobra.MaximumNArgs(1),
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.PkgOpts.PackageSource = choosePackage(args)
 
 		src := identifyAndFallbackToClusterSource()
 
 		pkgClient := packager.NewOrDie(&pkgConfig, packager.WithSource(src))
 		defer pkgClient.ClearTempPaths()
 
-		if err := pkgClient.Inspect(); err != nil {
+		if err := pkgClient.Inspect(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdPackageInspectErr, err.Error())
 		}
 	},
@@ -190,7 +190,7 @@ var packagePublishCmd = &cobra.Command{
 	Short:   lang.CmdPackagePublishShort,
 	Example: lang.CmdPackagePublishExample,
 	Args:    cobra.ExactArgs(2),
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.PkgOpts.PackageSource = args[0]
 
 		if !helpers.IsOCIURL(args[1]) {
@@ -216,7 +216,7 @@ var packagePublishCmd = &cobra.Command{
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		if err := pkgClient.Publish(); err != nil {
+		if err := pkgClient.Publish(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdPackagePublishErr, err.Error())
 		}
 	},
@@ -227,13 +227,13 @@ var packagePullCmd = &cobra.Command{
 	Short:   lang.CmdPackagePullShort,
 	Example: lang.CmdPackagePullExample,
 	Args:    cobra.ExactArgs(1),
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.PkgOpts.PackageSource = args[0]
 
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		if err := pkgClient.Pull(); err != nil {
+		if err := pkgClient.Pull(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdPackagePullErr, err.Error())
 		}
 	},

src/cmd/tools/zarf.go

@@ -186,7 +186,7 @@ var clearCacheCmd = &cobra.Command{
 var downloadInitCmd = &cobra.Command{
 	Use:   "download-init",
 	Short: lang.CmdToolsDownloadInitShort,
-	Run: func(_ *cobra.Command, _ []string) {
+	Run: func(cmd *cobra.Command, _ []string) {
 		url := zoci.GetInitPackageURL(config.CLIVersion)
 
 		remote, err := zoci.NewRemote(url, oci.PlatformForArch(config.GetArch()))
@@ -196,7 +196,7 @@ var downloadInitCmd = &cobra.Command{
 
 		source := &sources.OCISource{Remote: remote}
 
-		_, err = source.Collect(outputDirectory)
+		_, err = source.Collect(cmd.Context(), outputDirectory)
 		if err != nil {
 			message.Fatalf(err, lang.CmdToolsDownloadInitErr, err.Error())
 		}

src/pkg/packager/composer/list.go

@@ -120,7 +120,7 @@ func (ic *ImportChain) append(c types.ZarfComponent, index int, originalPackageN
 
 // NewImportChain creates a new import chain from a component
 // Returning the chain on error so we can have additional information to use during lint
-func NewImportChain(head types.ZarfComponent, index int, originalPackageName, arch, flavor string) (*ImportChain, error) {
+func NewImportChain(ctx context.Context, head types.ZarfComponent, index int, originalPackageName, arch, flavor string) (*ImportChain, error) {
 	ic := &ImportChain{}
 	if arch == "" {
 		return ic, fmt.Errorf("cannot build import chain: architecture must be provided")
@@ -179,11 +179,11 @@ func NewImportChain(head types.ZarfComponent, index int, originalPackageName, ar
 			}
 		} else if isRemote {
 			importURL = node.Import.URL
-			remote, err := ic.getRemote(node.Import.URL)
+			remote, err := ic.getRemote(ctx, node.Import.URL)
 			if err != nil {
 				return ic, err
 			}
-			pkg, err = remote.FetchZarfYAML(context.TODO())
+			pkg, err = remote.FetchZarfYAML(ctx)
 			if err != nil {
 				return ic, err
 			}
@@ -274,15 +274,15 @@ func (ic *ImportChain) Migrate(build types.ZarfBuildData) (warnings []string) {
 
 // Compose merges the import chain into a single component
 // fixing paths, overriding metadata, etc
-func (ic *ImportChain) Compose() (composed *types.ZarfComponent, err error) {
+func (ic *ImportChain) Compose(ctx context.Context) (composed *types.ZarfComponent, err error) {
 	composed = &ic.tail.ZarfComponent
 
 	if ic.tail.prev == nil {
 		// only had one component in the import chain
 		return composed, nil
 	}
 
-	if err := ic.fetchOCISkeleton(); err != nil {
+	if err := ic.fetchOCISkeleton(ctx); err != nil {
 		return nil, err
 	}
 

src/pkg/packager/composer/list_test.go

@@ -5,6 +5,7 @@
 package composer
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -48,7 +49,7 @@ func TestNewImportChain(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			_, err := NewImportChain(tt.head, 0, testPackageName, tt.arch, tt.flavor)
+			_, err := NewImportChain(context.Background(), tt.head, 0, testPackageName, tt.arch, tt.flavor)
 			require.ErrorContains(t, err, tt.expectedErr)
 		})
 	}
@@ -239,7 +240,7 @@ func TestCompose(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			composed, err := tt.ic.Compose()
+			composed, err := tt.ic.Compose(context.Background())
 			require.NoError(t, err)
 			require.EqualValues(t, &tt.expectedComposed, composed)
 		})

src/pkg/packager/composer/oci.go

@@ -23,7 +23,7 @@ import (
 	ocistore "oras.land/oras-go/v2/content/oci"
 )
 
-func (ic *ImportChain) getRemote(url string) (*zoci.Remote, error) {
+func (ic *ImportChain) getRemote(ctx context.Context, url string) (*zoci.Remote, error) {
 	if ic.remote != nil {
 		return ic.remote, nil
 	}
@@ -32,7 +32,7 @@ func (ic *ImportChain) getRemote(url string) (*zoci.Remote, error) {
 	if err != nil {
 		return nil, err
 	}
-	_, err = ic.remote.ResolveRoot(context.TODO())
+	_, err = ic.remote.ResolveRoot(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("published skeleton package for %q does not exist: %w", url, err)
 	}
@@ -45,17 +45,16 @@ func (ic *ImportChain) ContainsOCIImport() bool {
 	return ic.tail.prev != nil && ic.tail.prev.Import.URL != ""
 }
 
-func (ic *ImportChain) fetchOCISkeleton() error {
+func (ic *ImportChain) fetchOCISkeleton(ctx context.Context) error {
 	if !ic.ContainsOCIImport() {
 		return nil
 	}
 	node := ic.tail.prev
-	remote, err := ic.getRemote(node.Import.URL)
+	remote, err := ic.getRemote(ctx, node.Import.URL)
 	if err != nil {
 		return err
 	}
 
-	ctx := context.TODO()
 	manifest, err := remote.FetchRoot(ctx)
 	if err != nil {
 		return err
@@ -91,7 +90,6 @@ func (ic *ImportChain) fetchOCISkeleton() error {
 			return err
 		}
 
-		ctx := context.TODO()
 		// ensure the tarball is in the cache
 		exists, err := store.Exists(ctx, componentDesc)
 		if err != nil {

src/pkg/packager/create.go

@@ -5,6 +5,7 @@
 package packager
 
 import (
+	"context"
 	"fmt"
 	"os"
 
@@ -16,7 +17,7 @@ import (
 )
 
 // Create generates a Zarf package tarball for a given PackageConfig and optional base directory.
-func (p *Packager) Create() (err error) {
+func (p *Packager) Create(ctx context.Context) (err error) {
 	cwd, err := os.Getwd()
 	if err != nil {
 		return err
@@ -34,7 +35,7 @@ func (p *Packager) Create() (err error) {
 		return err
 	}
 
-	p.cfg.Pkg, p.warnings, err = pc.LoadPackageDefinition(p.layout)
+	p.cfg.Pkg, p.warnings, err = pc.LoadPackageDefinition(ctx, p.layout)
 	if err != nil {
 		return err
 	}
@@ -43,7 +44,7 @@ func (p *Packager) Create() (err error) {
 		return fmt.Errorf("package creation canceled")
 	}
 
-	if err := pc.Assemble(p.layout, p.cfg.Pkg.Components, p.cfg.Pkg.Metadata.Architecture); err != nil {
+	if err := pc.Assemble(ctx, p.layout, p.cfg.Pkg.Components, p.cfg.Pkg.Metadata.Architecture); err != nil {
 		return err
 	}
 
@@ -52,5 +53,5 @@ func (p *Packager) Create() (err error) {
 		return err
 	}
 
-	return pc.Output(p.layout, &p.cfg.Pkg)
+	return pc.Output(ctx, p.layout, &p.cfg.Pkg)
 }

src/pkg/packager/creator/compose.go

@@ -5,13 +5,15 @@
 package creator
 
 import (
+	"context"
+
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 	"github.com/defenseunicorns/zarf/src/pkg/packager/composer"
 	"github.com/defenseunicorns/zarf/src/types"
 )
 
 // ComposeComponents composes components and their dependencies into a single Zarf package using an import chain.
-func ComposeComponents(pkg types.ZarfPackage, flavor string) (types.ZarfPackage, []string, error) {
+func ComposeComponents(ctx context.Context, pkg types.ZarfPackage, flavor string) (types.ZarfPackage, []string, error) {
 	components := []types.ZarfComponent{}
 	warnings := []string{}
 
@@ -31,7 +33,7 @@ func ComposeComponents(pkg types.ZarfPackage, flavor string) (types.ZarfPackage,
 		component.Only.Flavor = ""
 
 		// build the import chain
-		chain, err := composer.NewImportChain(component, i, pkg.Metadata.Name, arch, flavor)
+		chain, err := composer.NewImportChain(ctx, component, i, pkg.Metadata.Name, arch, flavor)
 		if err != nil {
 			return types.ZarfPackage{}, nil, err
 		}
@@ -42,7 +44,7 @@ func ComposeComponents(pkg types.ZarfPackage, flavor string) (types.ZarfPackage,
 		warnings = append(warnings, warning...)
 
 		// get the composed component
-		composed, err := chain.Compose()
+		composed, err := chain.Compose(ctx)
 		if err != nil {
 			return types.ZarfPackage{}, nil, err
 		}

src/pkg/packager/creator/compose_test.go

@@ -5,6 +5,7 @@
 package creator
 
 import (
+	"context"
 	"testing"
 
 	"github.com/defenseunicorns/zarf/src/types"
@@ -160,7 +161,7 @@ func TestComposeComponents(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			pkg, _, err := ComposeComponents(tt.pkg, tt.flavor)
+			pkg, _, err := ComposeComponents(context.Background(), tt.pkg, tt.flavor)
 
 			if tt.expectedErr == "" {
 				require.NoError(t, err)

src/pkg/packager/creator/creator.go

@@ -5,13 +5,15 @@
 package creator
 
 import (
+	"context"
+
 	"github.com/defenseunicorns/zarf/src/pkg/layout"
 	"github.com/defenseunicorns/zarf/src/types"
 )
 
 // Creator is an interface for creating Zarf packages.
 type Creator interface {
-	LoadPackageDefinition(src *layout.PackagePaths) (pkg types.ZarfPackage, warnings []string, err error)
-	Assemble(dst *layout.PackagePaths, components []types.ZarfComponent, arch string) error
-	Output(dst *layout.PackagePaths, pkg *types.ZarfPackage) error
+	LoadPackageDefinition(ctx context.Context, src *layout.PackagePaths) (pkg types.ZarfPackage, warnings []string, err error)
+	Assemble(ctx context.Context, dst *layout.PackagePaths, components []types.ZarfComponent, arch string) error
+	Output(ctx context.Context, dst *layout.PackagePaths, pkg *types.ZarfPackage) error
 }