Simplify temp dir cleaning (#2133)

* Simplify temp dir cleaning

* rename vars

* add test

* update test

main.go

@@ -9,7 +9,6 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
-	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/felixge/fgprof"
@@ -180,29 +179,6 @@ func init() {
 	}
 }
 
-// Encloses tempdir cleanup in a function so it can be pushed
-// to a goroutine
-func runCleanup(ctx context.Context, execName string) {
-	// Every 15 minutes, attempt to remove dirs
-	pid := os.Getpid()
-	// Inital orphaned dir cleanup when the scanner is invoked
-	err := cleantemp.CleanTempDir(ctx, execName, pid)
-	if err != nil {
-		ctx.Logger().Error(err, "Error cleaning up orphaned directories ")
-	}
-
-	ticker := time.NewTicker(15 * time.Second)
-	defer ticker.Stop()
-
-	for range ticker.C {
-		err := cleantemp.CleanTempDir(ctx, execName, pid)
-		if err != nil {
-			ctx.Logger().Error(err, "Error cleaning up orphaned directories ")
-		}
-	}
-
-}
-
 func main() {
 	// setup logger
 	logFormat := log.WithConsoleSink
@@ -243,9 +219,7 @@ func main() {
 
 	ctx := context.Background()
 
-	var execName = "trufflehog"
-
-	go runCleanup(ctx, execName)
+	go cleantemp.RunCleanupLoop(ctx)
 
 }
 

pkg/cleantemp/cleantemp.go

@@ -8,6 +8,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/mitchellh/go-ps"
 
@@ -35,7 +36,13 @@ type CleanTemp interface {
 }
 
 // Deletes orphaned temp directories that do not contain running PID values
-func CleanTempDir(ctx logContext.Context, dirName string, pid int) error {
+func CleanTempDir(ctx logContext.Context) error {
+	executablePath, err := os.Executable()
+	if err != nil {
+		executablePath = "trufflehog"
+	}
+	execName := filepath.Base(executablePath)
+
 	// Finds other trufflehog PIDs that may be running
 	var pids []string
 	procs, err := ps.Processes()
@@ -44,38 +51,35 @@ func CleanTempDir(ctx logContext.Context, dirName string, pid int) error {
 	}
 
 	for _, proc := range procs {
-		if strings.Contains(proc.Executable(), dirName) {
+		if proc.Executable() == execName {
 			pids = append(pids, strconv.Itoa(proc.Pid()))
 		}
 	}
 
 	tempDir := os.TempDir()
-	files, err := os.ReadDir(tempDir)
+	dirs, err := os.ReadDir(tempDir)
 	if err != nil {
 		return fmt.Errorf("Error reading temp dir: %w", err)
 	}
 
-	// Current PID
-	pidStr := strconv.Itoa(pid)
-
 	pattern := `^trufflehog-\d+-\d+$`
 	re := regexp.MustCompile(pattern)
 
-	for _, file := range files {
-		// Make sure we don't delete the working dir of the current PID
-		if file.IsDir() && re.MatchString(file.Name()) && !strings.Contains(file.Name(), pidStr) {
+	for _, dir := range dirs {
+		// Ensure that all directories match the pattern
+		if re.MatchString(dir.Name()) {
 			// Mark these directories initially as ones that should be deleted
 			shouldDelete := true
 			// If they match any live PIDs, mark as should not delete
 			for _, pidval := range pids {
-				if strings.Contains(file.Name(), pidval) {
+				if strings.Contains(dir.Name(), fmt.Sprintf("-%s-", pidval)) {
 					shouldDelete = false
 					// break out so we can still delete directories even if no other Trufflehog processes are running
 					break
 				}
 			}
 			if shouldDelete {
-				dirPath := filepath.Join(tempDir, file.Name())
+				dirPath := filepath.Join(tempDir, dir.Name())
 				if err := os.RemoveAll(dirPath); err != nil {
 					return fmt.Errorf("Error deleting temp directory: %s", dirPath)
 				}
@@ -85,3 +89,22 @@ func CleanTempDir(ctx logContext.Context, dirName string, pid int) error {
 	}
 	return nil
 }
+
+// RunCleanupLoop runs a loop that cleans up orphaned directories every 15 seconds
+func RunCleanupLoop(ctx logContext.Context) {
+	err := CleanTempDir(ctx)
+	if err != nil {
+		ctx.Logger().Error(err, "Error cleaning up orphaned directories ")
+	}
+
+	ticker := time.NewTicker(15 * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		err := CleanTempDir(ctx)
+		if err != nil {
+			ctx.Logger().Error(err, "Error cleaning up orphaned directories ")
+		}
+	}
+
+}

pkg/cleantemp/cleantemp_test.go

@@ -0,0 +1,30 @@
+package cleantemp
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/mitchellh/go-ps"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExecName(t *testing.T) {
+	executablePath, err := os.Executable()
+	assert.Nil(t, err)
+	execName := filepath.Base(executablePath)
+	assert.Equal(t, "cleantemp.test", execName)
+
+	procs, err := ps.Processes()
+	assert.Nil(t, err)
+	assert.NotEmpty(t, procs)
+
+	found := false
+	for _, proc := range procs {
+		if proc.Executable() == execName {
+			found = true
+		}
+	}
+
+	assert.True(t, found)
+}