anchore · wagoodman · Feb 9, 2022 · Feb 9, 2022 · Feb 9, 2022
diff --git a/client.go b/client.go
@@ -16,35 +16,37 @@ import (
 	"github.com/wagoodman/go-partybus"
 )
 
-var tempDirGenerator = file.NewTempDirGenerator()
+var rootTempDirGenerator = file.NewTempDirGenerator("stereoscope")
 
 // GetImageFromSource returns an image from the explicitly provided source.
 func GetImageFromSource(ctx context.Context, imgStr string, source image.Source, registryOptions *image.RegistryOptions) (*image.Image, error) {
 	var provider image.Provider
 	log.Debugf("image: source=%+v location=%+v", source, imgStr)
 
+	tempDirGenerator := rootTempDirGenerator.NewGenerator()
+
 	switch source {
 	case image.DockerTarballSource:
 		// note: the imgStr is the path on disk to the tar file
-		provider = docker.NewProviderFromTarball(imgStr, &tempDirGenerator, nil, nil)
+		provider = docker.NewProviderFromTarball(imgStr, tempDirGenerator, nil, nil)
 	case image.DockerDaemonSource:
 		c, err := dockerClient.GetClient()
 		if err != nil {
 			return nil, err
 		}
-		provider = docker.NewProviderFromDaemon(imgStr, &tempDirGenerator, c)
+		provider = docker.NewProviderFromDaemon(imgStr, tempDirGenerator, c)
 	case image.PodmanDaemonSource:
 		c, err := podman.GetClient()
 		if err != nil {
 			return nil, err
 		}
-		provider = docker.NewProviderFromDaemon(imgStr, &tempDirGenerator, c)
+		provider = docker.NewProviderFromDaemon(imgStr, tempDirGenerator, c)
 	case image.OciDirectorySource:
-		provider = oci.NewProviderFromPath(imgStr, &tempDirGenerator)
+		provider = oci.NewProviderFromPath(imgStr, tempDirGenerator)
 	case image.OciTarballSource:
-		provider = oci.NewProviderFromTarball(imgStr, &tempDirGenerator)
+		provider = oci.NewProviderFromTarball(imgStr, tempDirGenerator)
 	case image.OciRegistrySource:
-		provider = oci.NewProviderFromRegistry(imgStr, &tempDirGenerator, registryOptions)
+		provider = oci.NewProviderFromRegistry(imgStr, tempDirGenerator, registryOptions)
 	default:
 		return nil, fmt.Errorf("unable determine image source")
 	}
@@ -80,8 +82,10 @@ func SetBus(b *partybus.Bus) {
 	bus.SetPublisher(b)
 }
 
+// Cleanup deletes all directories created by stereoscope calls. Note: please use image.Image.Cleanup() over this
+// function when possible.
 func Cleanup() {
-	if err := tempDirGenerator.Cleanup(); err != nil {
-		log.Errorf("failed to cleanup: %w", err)
+	if err := rootTempDirGenerator.Cleanup(); err != nil {
+		log.Errorf("failed to cleanup tempdir root: %w", err)
 	}
 }
diff --git a/examples/basic.go b/examples/basic.go
@@ -13,8 +13,6 @@ import (
 )
 
 func main() {
-	// note: we are writing out temp files which should be cleaned up after you're done with the image object
-	defer stereoscope.Cleanup()
 
 	// context for network requests
 	ctx, cancel := context.WithCancel(context.Background())
@@ -34,12 +32,15 @@ func main() {
 		panic(err)
 	}
 
+	// note: we are writing out temp files which should be cleaned up after you're done with the image object
+	defer image.Cleanup()
+
 	for _, layer := range image.Layers {
 		fmt.Printf("layer: %s\n", layer.Metadata.Digest)
 	}
 
-	//////////////////////////////////////////////////////////////////
-	//// Show the filetree for each layer
+	////////////////////////////////////////////////////////////////
+	// Show the filetree for each layer
 	for idx, layer := range image.Layers {
 		fmt.Printf("Walking layer: %d", idx)
 		err = layer.Tree.Walk(func(path file.Path, f filenode.FileNode) error {
@@ -53,7 +54,7 @@ func main() {
 	}
 
 	//////////////////////////////////////////////////////////////////
-	//// Show the squashed filetree for each layer
+	// Show the squashed filetree for each layer
 	for idx, layer := range image.Layers {
 		fmt.Printf("Walking squashed layer: %d", idx)
 		err = layer.SquashedTree.Walk(func(path file.Path, f filenode.FileNode) error {
@@ -67,7 +68,7 @@ func main() {
 	}
 
 	//////////////////////////////////////////////////////////////////
-	//// Show the final squashed tree
+	// Show the final squashed tree
 	fmt.Printf("Walking squashed image (same as the last layer squashed tree)")
 	err = image.SquashedTree().Walk(func(path file.Path, f filenode.FileNode) error {
 		fmt.Println("   ", path)
@@ -78,7 +79,7 @@ func main() {
 	}
 
 	//////////////////////////////////////////////////////////////////
-	//// Fetch file contents from the (squashed) image
+	// Fetch file contents from the (squashed) image
 	filePath := file.Path("/etc/group")
 	contentReader, err := image.FileContentsFromSquash(filePath)
 	if err != nil {

diff --git a/pkg/file/temp_dir_generator.go b/pkg/file/temp_dir_generator.go
@@ -1,50 +1,65 @@
 package file
 
 import (
-	"fmt"
-	"io/ioutil"
 	"os"
-	"sync"
+	"strings"
 
 	"github.com/hashicorp/go-multierror"
 )
 
 type TempDirGenerator struct {
-	tempDir []string
-	lock    *sync.Mutex
+	rootPrefix   string
+	rootLocation string
+	children     []*TempDirGenerator
 }
 
-func NewTempDirGenerator() TempDirGenerator {
-	return TempDirGenerator{
-		tempDir: make([]string, 0),
-		lock:    &sync.Mutex{},
+func NewTempDirGenerator(name string) *TempDirGenerator {
+	return &TempDirGenerator{
+		rootPrefix: name,
 	}
 }
 
-// NewTempDir creates an empty dir in the platform temp dir
-func (t *TempDirGenerator) NewTempDir() (string, error) {
-	t.lock.Lock()
-	defer t.lock.Unlock()
+func (t *TempDirGenerator) getOrCreateRootLocation() (string, error) {
+	if t.rootLocation == "" {
+		location, err := os.MkdirTemp("", t.rootPrefix+"-")
+		if err != nil {
+			return "", err
+		}
+
+		t.rootLocation = location
+	}
+	return t.rootLocation, nil
+}
 
-	dir, err := ioutil.TempDir("", "stereoscope-cache")
+// NewGenerator creates a child generator capable of making sibling temp directories.
+func (t *TempDirGenerator) NewGenerator() *TempDirGenerator {
+	gen := NewTempDirGenerator(t.rootPrefix)
+	t.children = append(t.children, gen)
+	return gen
+}
+
+// NewDirectory creates a new temp dir within the generators prefix temp dir.
+func (t *TempDirGenerator) NewDirectory(name ...string) (string, error) {
+	location, err := t.getOrCreateRootLocation()
 	if err != nil {
-		return "", fmt.Errorf("could not create temp dir: %w", err)
+		return "", err
 	}
 
-	t.tempDir = append(t.tempDir, dir)
-	return dir, nil
+	return os.MkdirTemp(location, strings.Join(name, "-")+"-")
 }
 
+// Cleanup deletes all temp dirs created by this generator and any child generator.
 func (t *TempDirGenerator) Cleanup() error {
-	t.lock.Lock()
-	defer t.lock.Unlock()
-
-	var allErrors error
-	for _, dir := range t.tempDir {
-		err := os.RemoveAll(dir)
-		if err != nil {
-			allErrors = multierror.Append(allErrors, err)
+	var allErrs error
+	for _, gen := range t.children {
+		if err := gen.Cleanup(); err != nil {
+			allErrs = multierror.Append(allErrs, err)
+		}
+	}
+	if t.rootLocation != "" {
+		if err := os.RemoveAll(t.rootLocation); err != nil {
+			allErrs = multierror.Append(allErrs, err)
 		}
 	}
-	return allErrors
+	return allErrs
 }
diff --git a/pkg/file/temp_dir_generator_test.go b/pkg/file/temp_dir_generator_test.go
@@ -0,0 +1,91 @@
+package file
+
+import (
+	"os"
+	"path"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTempDirGenerator(t *testing.T) {
+	tests := []struct {
+		name            string
+		genPrefix       string
+		names           []string
+		extraGenerators int
+	}{
+		{
+			name:      "3 temp dirs",
+			genPrefix: "a-special-prefix",
+			names: []string{
+				"a",
+				"bee",
+				"si",
+			},
+		},
+		{
+			name:      "3 temp dirs on the root generator + 2 extra generators",
+			genPrefix: "b-special-prefix",
+			names: []string{
+				"a",
+				"bee",
+				"si",
+			},
+			extraGenerators: 2,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			expectedPrefix := path.Join(os.TempDir(), test.genPrefix)
+
+			assert.True(t, !doesGlobExist(t, expectedPrefix+"*"),
+				"prefix temp dir already exists before test started")
+
+			root := NewTempDirGenerator(test.genPrefix)
+
+			for _, n := range test.names {
+				d, err := root.NewDirectory(n)
+				assert.NoError(t, err)
+				assert.True(t, doesGlobExist(t, d), "sub-temp dir does not exist (root)")
+				assert.Contains(t, d, expectedPrefix)
+				assert.NotEmpty(t, root.rootLocation)
+				assert.Contains(t, d, root.rootLocation)
+			}
+
+			assert.True(t, doesGlobExist(t, expectedPrefix+"*"), "prefix temp dir does not exist")
+
+			var gen *TempDirGenerator
+			for i := 0; i < test.extraGenerators; i++ {
+				gen = root.NewGenerator()
+				for _, n := range test.names {
+					d, err := gen.NewDirectory(n)
+					assert.NoError(t, err)
+					assert.True(t, doesGlobExist(t, d), "sub-temp dir does not exist (sub)")
+					assert.Contains(t, d, expectedPrefix)
+					assert.NotEmpty(t, gen.rootLocation)
+					assert.Contains(t, d, gen.rootLocation)
+				}
+
+			}
+
+			assert.NoError(t, root.Cleanup())
+
+			assert.True(t, !doesGlobExist(t, expectedPrefix+"*"), "cleanup did not remove prefix temp dir")
+
+		})
+	}
+}
+
+func doesGlobExist(t *testing.T, pattern string) bool {
+	t.Helper()
+	m, err := filepath.Glob(pattern)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(m) > 0 {
+		return true
+	}
+	return false
+}
diff --git a/pkg/image/docker/daemon_provider.go b/pkg/image/docker/daemon_provider.go
@@ -139,7 +139,7 @@ func (p *DaemonImageProvider) pull(ctx context.Context) error {
 
 // Provide an image object that represents the cached docker image tar fetched from a docker daemon.
 func (p *DaemonImageProvider) Provide(ctx context.Context) (*image.Image, error) {
-	imageTempDir, err := p.tmpDirGen.NewTempDir()
+	imageTempDir, err := p.tmpDirGen.NewDirectory("docker-daemon-image")
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/image/docker/tarball_provider.go b/pkg/image/docker/tarball_provider.go
@@ -93,7 +93,7 @@ func (p *TarballImageProvider) Provide(context.Context) (*image.Image, error) {
 
 	metadata = append(metadata, image.WithRepoDigests(p.repoDigests))
 
-	contentTempDir, err := p.tmpDirGen.NewTempDir()
+	contentTempDir, err := p.tmpDirGen.NewDirectory("docker-tarball-image")
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/image/image.go b/pkg/image/image.go
@@ -4,6 +4,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"io"
+	"os"
 
 	"github.com/anchore/stereoscope/internal/bus"
 	"github.com/anchore/stereoscope/internal/log"
@@ -251,10 +252,20 @@ func (i *Image) ResolveLinkByLayerSquash(ref file.Reference, layer int, options
 	return resolvedRef, err
 }
 
-// ResolveLinkByLayerSquash resolves a symlink or hardlink for the given file reference relative to the result from the image squash.
+// ResolveLinkByImageSquash resolves a symlink or hardlink for the given file reference relative to the result from the image squash.
 // If the given file reference is not a link type, or is a unresolvable (dead) link, then the given file reference is returned.
 func (i *Image) ResolveLinkByImageSquash(ref file.Reference, options ...filetree.LinkResolutionOption) (*file.Reference, error) {
 	allOptions := append([]filetree.LinkResolutionOption{filetree.FollowBasenameLinks}, options...)
 	_, resolvedRef, err := i.Layers[len(i.Layers)-1].SquashedTree.File(ref.RealPath, allOptions...)
 	return resolvedRef, err
 }
+
+// Cleanup removes all temporary files created from parsing the image. Future calls to image will not function correctly after this call.
+func (i *Image) Cleanup() error {
+	if i.contentCacheDir != "" {
+		if err := os.RemoveAll(i.contentCacheDir); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/pkg/image/oci/directory_provider.go b/pkg/image/oci/directory_provider.go
@@ -61,7 +61,7 @@ func (p *DirectoryImageProvider) Provide(context.Context) (*image.Image, error)
 		metadata = append(metadata, image.WithManifest(rawManifest))
 	}
 
-	contentTempDir, err := p.tmpDirGen.NewTempDir()
+	contentTempDir, err := p.tmpDirGen.NewDirectory("oci-dir-image")
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/image/oci/registry_provider.go b/pkg/image/oci/registry_provider.go
@@ -34,7 +34,7 @@ func NewProviderFromRegistry(imgStr string, tmpDirGen *file.TempDirGenerator, re
 func (p *RegistryImageProvider) Provide(ctx context.Context) (*image.Image, error) {
 	log.Debugf("pulling image info directly from registry image=%q", p.imageStr)
 
-	imageTempDir, err := p.tmpDirGen.NewTempDir()
+	imageTempDir, err := p.tmpDirGen.NewDirectory("oci-registry-image")
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/image/oci/tarball_provider.go b/pkg/image/oci/tarball_provider.go
@@ -32,7 +32,7 @@ func (p *TarballImageProvider) Provide(ctx context.Context) (*image.Image, error
 		return nil, fmt.Errorf("unable to open OCI tarball: %w", err)
 	}
 
-	tempDir, err := p.tmpDirGen.NewTempDir()
+	tempDir, err := p.tmpDirGen.NewDirectory("oci-tarball-image")
 	if err != nil {
 		return nil, err
 	}