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>

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)
-			return nil
+			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