make AllLocations accept a context (#2518)

The previous implementation would leak a goroutine if the caller of
AllLocations stopped iterating early. Now, accept a context so that the
caller can cancel the AllLocations iterator rather than leak the
goroutine.

Signed-off-by: Will Murphy <will.murphy@anchore.com>

go.mod

@@ -73,6 +73,7 @@ require (
 	github.com/wagoodman/go-progress v0.0.0-20230925121702-07e42b3cdba0
 	github.com/xeipuuv/gojsonschema v1.2.0
 	github.com/zyedidia/generic v1.2.2-0.20230320175451-4410d2372cb1
+	go.uber.org/goleak v1.3.0
 	golang.org/x/mod v0.14.0
 	golang.org/x/net v0.20.0
 	gopkg.in/yaml.v3 v3.0.1

go.sum

@@ -854,8 +854,8 @@ go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
 go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
-go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
+go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
-go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
+go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI=
 go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ=

syft/file/cataloger/filemetadata/cataloger.go

@@ -1,6 +1,7 @@
 package filemetadata
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/dustin/go-humanize"
@@ -21,8 +22,10 @@ func NewCataloger() *Cataloger {
 func (i *Cataloger) Catalog(resolver file.Resolver, coordinates ...file.Coordinates) (map[file.Coordinates]file.Metadata, error) {
 	results := make(map[file.Coordinates]file.Metadata)
 	var locations <-chan file.Location
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
 	if len(coordinates) == 0 {
-		locations = resolver.AllLocations()
+		locations = resolver.AllLocations(ctx)
 	} else {
 		locations = func() <-chan file.Location {
 			ch := make(chan file.Location)
@@ -35,7 +38,12 @@ func (i *Cataloger) Catalog(resolver file.Resolver, coordinates ...file.Coordina
 						continue
 					}
 					for _, loc := range locs {
-						ch <- loc
+						select {
+						case <-ctx.Done():
+							return
+						case ch <- loc:
+							continue
+						}
 					}
 				}
 			}()

syft/file/cataloger/internal/all_regular_files.go

@@ -1,13 +1,17 @@
 package internal
 
 import (
+	"context"
+
 	stereoscopeFile "github.com/anchore/stereoscope/pkg/file"
 	"github.com/anchore/syft/internal/log"
 	"github.com/anchore/syft/syft/file"
 )
 
 func AllRegularFiles(resolver file.Resolver) (locations []file.Location) {
-	for location := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for location := range resolver.AllLocations(ctx) {
 		resolvedLocations, err := resolver.FilesByPath(location.RealPath)
 		if err != nil {
 			log.Warnf("unable to resolve %+v: %+v", location, err)

syft/file/mock_resolver.go

@@ -1,6 +1,7 @@
 package file
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -144,12 +145,17 @@ func (r MockResolver) RelativeFileByPath(_ Location, path string) *Location {
 	return &paths[0]
 }
 
-func (r MockResolver) AllLocations() <-chan Location {
+func (r MockResolver) AllLocations(ctx context.Context) <-chan Location {
 	results := make(chan Location)
 	go func() {
 		defer close(results)
 		for _, l := range r.locations {
-			results <- l
+			select {
+			case <-ctx.Done():
+				return
+			case results <- l:
+				continue
+			}
 		}
 	}()
 	return results

syft/file/resolver.go

@@ -1,6 +1,9 @@
 package file
 
-import "io"
+import (
+	"context"
+	"io"
+)
 
 // Resolver is an interface that encompasses how to get specific file references and file contents for a generic data source.
 type Resolver interface {
@@ -53,7 +56,7 @@ type LocationResolver interface {
 	// The implementation for this may vary, however, generally the following considerations should be made:
 	// - NO symlink resolution should be performed on results
 	// - returns locations for any file or directory
-	AllLocations() <-chan Location
+	AllLocations(ctx context.Context) <-chan Location
 }
 
 type WritableResolver interface {

syft/internal/fileresolver/container_image_all_layers.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -234,14 +235,19 @@ func (r *ContainerImageAllLayers) FilesByMIMEType(types ...string) ([]file.Locat
 	return uniqueLocations, nil
 }
 
-func (r *ContainerImageAllLayers) AllLocations() <-chan file.Location {
+func (r *ContainerImageAllLayers) AllLocations(ctx context.Context) <-chan file.Location {
 	results := make(chan file.Location)
 	go func() {
 		defer close(results)
 		for _, layerIdx := range r.layers {
 			tree := r.img.Layers[layerIdx].Tree
 			for _, ref := range tree.AllFiles(stereoscopeFile.AllTypes()...) {
-				results <- file.NewLocationFromImage(string(ref.RealPath), ref, r.img)
+				select {
+				case <-ctx.Done():
+					return
+				case results <- file.NewLocationFromImage(string(ref.RealPath), ref, r.img):
+					continue
+				}
 			}
 		}
 	}()

syft/internal/fileresolver/container_image_all_layers_test.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 	"sort"
 	"testing"
@@ -358,7 +359,9 @@ func TestAllLayersImageResolver_FilesContents_errorOnDirRequest(t *testing.T) {
 	assert.NoError(t, err)
 
 	var dirLoc *file.Location
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		entry, err := resolver.img.FileCatalog.Get(loc.Reference())
 		require.NoError(t, err)
 		if entry.Metadata.IsDir() {
@@ -517,7 +520,9 @@ func TestAllLayersResolver_AllLocations(t *testing.T) {
 	assert.NoError(t, err)
 
 	paths := strset.New()
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		paths.Add(loc.RealPath)
 	}
 	expected := []string{

syft/internal/fileresolver/container_image_squash.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -172,12 +173,17 @@ func (r *ContainerImageSquash) FileContentsByLocation(location file.Location) (i
 	return r.img.OpenReference(location.Reference())
 }
 
-func (r *ContainerImageSquash) AllLocations() <-chan file.Location {
+func (r *ContainerImageSquash) AllLocations(ctx context.Context) <-chan file.Location {
 	results := make(chan file.Location)
 	go func() {
 		defer close(results)
 		for _, ref := range r.img.SquashedTree().AllFiles(stereoscopeFile.AllTypes()...) {
-			results <- file.NewLocationFromImage(string(ref.RealPath), ref, r.img)
+			select {
+			case <-ctx.Done():
+				return
+			case results <- file.NewLocationFromImage(string(ref.RealPath), ref, r.img):
+				continue
+			}
 		}
 	}()
 	return results

syft/internal/fileresolver/container_image_squash_test.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 	"sort"
 	"testing"
@@ -343,7 +344,9 @@ func TestSquashImageResolver_FilesContents_errorOnDirRequest(t *testing.T) {
 	assert.NoError(t, err)
 
 	var dirLoc *file.Location
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		entry, err := resolver.img.FileCatalog.Get(loc.Reference())
 		require.NoError(t, err)
 		if entry.Metadata.IsDir() {
@@ -533,7 +536,9 @@ func TestSquashResolver_AllLocations(t *testing.T) {
 	assert.NoError(t, err)
 
 	paths := strset.New()
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		paths.Add(loc.RealPath)
 	}
 	expected := []string{

syft/internal/fileresolver/deferred.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 
 	"github.com/anchore/syft/internal/log"
@@ -80,13 +81,13 @@ func (d *Deferred) RelativeFileByPath(location file.Location, path string) *file
 	return r.RelativeFileByPath(location, path)
 }
 
-func (d *Deferred) AllLocations() <-chan file.Location {
+func (d *Deferred) AllLocations(ctx context.Context) <-chan file.Location {
 	r, err := d.getResolver()
 	if err != nil {
 		log.Debug("unable to get resolver: %v", err)
 		return nil
 	}
-	return r.AllLocations()
+	return r.AllLocations(ctx)
 }
 
 func (d *Deferred) FileMetadataByLocation(location file.Location) (file.Metadata, error) {

syft/internal/fileresolver/directory.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -229,12 +230,17 @@ func (r Directory) FileContentsByLocation(location file.Location) (io.ReadCloser
 	return stereoscopeFile.NewLazyReadCloser(filePath), nil
 }
 
-func (r *Directory) AllLocations() <-chan file.Location {
+func (r *Directory) AllLocations(ctx context.Context) <-chan file.Location {
 	results := make(chan file.Location)
 	go func() {
 		defer close(results)
 		for _, ref := range r.tree.AllFiles(stereoscopeFile.AllTypes()...) {
-			results <- file.NewLocationFromDirectory(r.responsePath(string(ref.RealPath)), ref)
+			select {
+			case <-ctx.Done():
+				return
+			case results <- file.NewLocationFromDirectory(r.responsePath(string(ref.RealPath)), ref):
+				continue
+			}
 		}
 	}()
 	return results

syft/internal/fileresolver/directory_test.go

@@ -4,6 +4,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 	"io/fs"
 	"os"
@@ -17,6 +18,7 @@ import (
 	"github.com/scylladb/go-set/strset"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 
 	stereoscopeFile "github.com/anchore/stereoscope/pkg/file"
 	"github.com/anchore/syft/syft/file"
@@ -1380,7 +1382,7 @@ func TestDirectoryResolver_DoNotAddVirtualPathsToTree(t *testing.T) {
 	require.NoError(t, err)
 
 	var allRealPaths []stereoscopeFile.Path
-	for l := range resolver.AllLocations() {
+	for l := range resolver.AllLocations(context.Background()) {
 		allRealPaths = append(allRealPaths, stereoscopeFile.Path(l.RealPath))
 	}
 	pathSet := stereoscopeFile.NewPathSet(allRealPaths...)
@@ -1398,11 +1400,14 @@ func TestDirectoryResolver_DoNotAddVirtualPathsToTree(t *testing.T) {
 }
 
 func TestDirectoryResolver_FilesContents_errorOnDirRequest(t *testing.T) {
+	defer goleak.VerifyNone(t)
 	resolver, err := NewFromDirectory("./test-fixtures/system_paths", "")
 	assert.NoError(t, err)
 
 	var dirLoc *file.Location
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		entry, err := resolver.index.Get(loc.Reference())
 		require.NoError(t, err)
 		if entry.Metadata.IsDir() {
@@ -1423,7 +1428,7 @@ func TestDirectoryResolver_AllLocations(t *testing.T) {
 	assert.NoError(t, err)
 
 	paths := strset.New()
-	for loc := range resolver.AllLocations() {
+	for loc := range resolver.AllLocations(context.Background()) {
 		if strings.HasPrefix(loc.RealPath, "/") {
 			// ignore outside the fixture root for now
 			continue
@@ -1449,3 +1454,14 @@ func TestDirectoryResolver_AllLocations(t *testing.T) {
 
 	assert.ElementsMatchf(t, expected, pathsList, "expected all paths to be indexed, but found different paths: \n%s", cmp.Diff(expected, paths.List()))
 }
+
+func TestAllLocationsDoesNotLeakGoRoutine(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	resolver, err := NewFromDirectory("./test-fixtures/symlinks-from-image-symlinks-fixture", "")
+	require.NoError(t, err)
+	ctx, cancel := context.WithCancel(context.Background())
+	for range resolver.AllLocations(ctx) {
+		break
+	}
+	cancel()
+}

syft/internal/fileresolver/empty.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 
 	"github.com/anchore/syft/syft/file"
@@ -34,7 +35,7 @@ func (e Empty) RelativeFileByPath(_ file.Location, _ string) *file.Location {
 	return nil
 }
 
-func (e Empty) AllLocations() <-chan file.Location {
+func (e Empty) AllLocations(_ context.Context) <-chan file.Location {
 	return nil
 }
 

syft/internal/fileresolver/excluding_file.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -69,13 +70,18 @@ func (r *excluding) RelativeFileByPath(location file.Location, path string) *fil
 	return l
 }
 
-func (r *excluding) AllLocations() <-chan file.Location {
+func (r *excluding) AllLocations(ctx context.Context) <-chan file.Location {
 	c := make(chan file.Location)
 	go func() {
 		defer close(c)
-		for location := range r.delegate.AllLocations() {
+		for location := range r.delegate.AllLocations(ctx) {
 			if !locationMatches(&location, r.excludeFn) {
-				c <- location
+				select {
+				case <-ctx.Done():
+					return
+				case c <- location:
+					continue
+				}
 			}
 		}
 	}()

syft/internal/fileresolver/excluding_file_test.go

@@ -1,6 +1,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 	"strings"
 	"testing"
@@ -69,7 +70,9 @@ func TestExcludingResolver(t *testing.T) {
 
 			locations = []file.Location{}
 
-			channel := er.AllLocations()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			channel := er.AllLocations(ctx)
 			for location := range channel {
 				locations = append(locations, location)
 			}
@@ -182,13 +185,18 @@ func (r *mockResolver) RelativeFileByPath(_ file.Location, path string) *file.Lo
 	return &l
 }
 
-func (r *mockResolver) AllLocations() <-chan file.Location {
+func (r *mockResolver) AllLocations(ctx context.Context) <-chan file.Location {
 	c := make(chan file.Location)
 	go func() {
 		defer close(c)
 		locations, _ := r.getLocations()
 		for _, location := range locations {
-			c <- location
+			select {
+			case <-ctx.Done():
+				return
+			case c <- location:
+				continue
+			}
 		}
 	}()
 	return c

syft/internal/fileresolver/unindexed_directory.go

@@ -1,6 +1,8 @@
 package fileresolver
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"io"
 	"io/fs"
@@ -225,8 +227,9 @@ func (u UnindexedDirectory) RelativeFileByPath(l file.Location, p string) *file.
 
 // - NO symlink resolution should be performed on results
 // - returns locations for any file or directory
-func (u UnindexedDirectory) AllLocations() <-chan file.Location {
+func (u UnindexedDirectory) AllLocations(ctx context.Context) <-chan file.Location {
 	out := make(chan file.Location)
+	errWalkCanceled := fmt.Errorf("walk canceled")
 	go func() {
 		defer close(out)
 		err := afero.Walk(u.fs, u.absPath("."), func(p string, info fs.FileInfo, err error) error {
@@ -235,10 +238,14 @@ func (u UnindexedDirectory) AllLocations() <-chan file.Location {
 				return nil
 			}
 			p = strings.TrimPrefix(p, "/")
-			out <- file.NewLocation(p)
+			select {
+			case out <- file.NewLocation(p):
 				return nil
+			case <-ctx.Done():
+				return errWalkCanceled
+			}
 		})
-		if err != nil {
+		if err != nil && !errors.Is(err, errWalkCanceled) {
 			log.Debug(err)
 		}
 	}()

syft/internal/fileresolver/unindexed_directory_test.go

@@ -4,6 +4,7 @@
 package fileresolver
 
 import (
+	"context"
 	"io"
 	"os"
 	"path"
@@ -17,6 +18,7 @@ import (
 	"github.com/scylladb/go-set/strset"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 
 	stereoscopeFile "github.com/anchore/stereoscope/pkg/file"
 	"github.com/anchore/syft/syft/file"
@@ -538,6 +540,19 @@ func Test_UnindexedDirectoryResolver_Basic(t *testing.T) {
 	require.Len(t, locations, 5)
 }
 
+func Test_UnindexedDirectoryResolver_NoGoroutineLeak(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	wd, err := os.Getwd()
+	require.NoError(t, err)
+
+	r := NewFromUnindexedDirectory(path.Join(wd, "test-fixtures"))
+	ctx, cancel := context.WithCancel(context.Background())
+	for range r.AllLocations(ctx) {
+		break
+	}
+	cancel()
+}
+
 func Test_UnindexedDirectoryResolver_FilesByPath_relativeRoot(t *testing.T) {
 	cases := []struct {
 		name         string
@@ -1168,7 +1183,9 @@ func Test_UnindexedDirectoryResolver_resolvesLinks(t *testing.T) {
 func Test_UnindexedDirectoryResolver_DoNotAddVirtualPathsToTree(t *testing.T) {
 	resolver := NewFromUnindexedDirectory("./test-fixtures/symlinks-prune-indexing")
 
-	allLocations := resolver.AllLocations()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	allLocations := resolver.AllLocations(ctx)
 	var allRealPaths []stereoscopeFile.Path
 	for l := range allLocations {
 		allRealPaths = append(allRealPaths, stereoscopeFile.Path(l.RealPath))
@@ -1200,7 +1217,9 @@ func Test_UnindexedDirectoryResolver_AllLocations(t *testing.T) {
 	resolver := NewFromUnindexedDirectory("./test-fixtures/symlinks-from-image-symlinks-fixture")
 
 	paths := strset.New()
-	for loc := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for loc := range resolver.AllLocations(ctx) {
 		if strings.HasPrefix(loc.RealPath, "/") {
 			// ignore outside of the fixture root for now
 			continue

syft/pkg/cataloger/binary/cataloger_test.go

@@ -1,6 +1,7 @@
 package binary
 
 import (
+	"context"
 	"errors"
 	"flag"
 	"fmt"
@@ -1237,7 +1238,7 @@ func (p *panicyResolver) RelativeFileByPath(_ file.Location, _ string) *file.Loc
 	return nil
 }
 
-func (p *panicyResolver) AllLocations() <-chan file.Location {
+func (p *panicyResolver) AllLocations(_ context.Context) <-chan file.Location {
 	return nil
 }
 

syft/pkg/cataloger/internal/pkgtest/observing_resolver.go

@@ -1,6 +1,7 @@
 package pkgtest
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"sort"
@@ -209,8 +210,8 @@ func (r *ObservingResolver) FileContentsByLocation(location file.Location) (io.R
 
 // For the remaining resolver methods...
 
-func (r *ObservingResolver) AllLocations() <-chan file.Location {
+func (r *ObservingResolver) AllLocations(ctx context.Context) <-chan file.Location {
-	return r.decorated.AllLocations()
+	return r.decorated.AllLocations(ctx)
 }
 
 func (r *ObservingResolver) HasPath(s string) bool {

syft/pkg/cataloger/nix/cataloger.go

@@ -4,6 +4,7 @@ Package nix provides a concrete Cataloger implementation for packages within the
 package nix
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/bmatcuk/doublestar/v4"
@@ -31,7 +32,9 @@ func (c *StoreCataloger) Catalog(resolver file.Resolver) ([]pkg.Package, []artif
 	// we want to search for only directories, which isn't possible via the stereoscope API, so we need to apply the glob manually on all returned paths
 	var pkgs []pkg.Package
 	var filesByPath = make(map[string]*file.LocationSet)
-	for location := range resolver.AllLocations() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	for location := range resolver.AllLocations(ctx) {
 		matchesStorePath, err := doublestar.Match("**/nix/store/*", location.RealPath)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to match nix store path: %w", err)

syft/pkg/cataloger/redhat/parse_rpm_db_test.go

@@ -1,6 +1,7 @@
 package redhat
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"testing"
@@ -34,7 +35,7 @@ func (r rpmdbTestFileResolverMock) FileContentsByLocation(location file.Location
 	panic("not implemented")
 }
 
-func (r rpmdbTestFileResolverMock) AllLocations() <-chan file.Location {
+func (r rpmdbTestFileResolverMock) AllLocations(_ context.Context) <-chan file.Location {
 	panic("not implemented")
 }