Fix pre-receive hook hangs and missing logs by flushing logs on signal and using CommandContext for git commands

main.go

@@ -14,7 +14,7 @@ import (
 	"strings"
 	"sync"
 	"syscall"
-
+	"time"
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/fatih/color"
 	"github.com/felixge/fgprof"
@@ -61,6 +61,7 @@ var (
 	results             = cli.Flag("results", "Specifies which type(s) of results to output: verified (confirmed valid by API), unknown (verification failed due to error), unverified (detected but not verified), filtered_unverified (unverified but would have been filtered out). Defaults to verified,unverified,unknown.").String()
 	noColor             = cli.Flag("no-color", "Disable colorized output").Bool()
 	noColour            = cli.Flag("no-colour", "Alias for --no-color").Hidden().Bool()
+	logSync func() error //Package-level variable for sync function
 
 	allowVerificationOverlap   = cli.Flag("allow-verification-overlap", "Allow verification of similar credentials across detectors").Bool()
 	filterUnverified           = cli.Flag("filter-unverified", "Only output first unverified result per chunk per detector if there are more than one results.").Bool()
@@ -354,6 +355,7 @@ func main() {
 		logFormat = log.WithJSONSink
 	}
 	logger, sync := log.New("trufflehog", logFormat(os.Stderr, log.WithGlobalRedaction()))
+	logSync = sync //
 	// make it the default logger for contexts
 	context.SetDefaultLogger(logger)
 
@@ -413,6 +415,21 @@ func run(state overseer.State) {
 		} else {
 			logger.Info("cleaned temporary artifacts")
 		}
+
+		// Flush logs with timeout to prevent hanging
+		if logSync != nil {
+			done := make(chan struct{})
+			go func() {
+				_ = logSync()
+				close(done)
+			}()
+
+			select {
+			case <-done:
+			case <-time.After(100 * time.Millisecond):
+				logger.Info("Log flush timed out, exiting")
+			}
+		}
 		os.Exit(0)
 	}()
 

pkg/gitparse/gitparse.go

@@ -30,6 +30,9 @@ const (
 
 	// defaultMaxCommitSize is the maximum size for a commit. Larger commits will be cut off.
 	defaultMaxCommitSize int64 = 2 * 1024 * 1024 * 1024 // 2GB
+
+	// defaultWaitDelay is the default time to wait after context cancellation before forcefully killing git processes.
+	defaultWaitDelay = 5 * time.Second
 )
 
 // contentWriter defines a common interface for writing, reading, and managing diff content.
@@ -124,6 +127,7 @@ type Parser struct {
 	maxDiffSize   int64
 	maxCommitSize int64
 	dateFormat    string
+	waitDelay     time.Duration
 
 	useCustomContentWriter bool
 }
@@ -204,6 +208,14 @@ func WithMaxCommitSize(maxCommitSize int64) Option {
 	}
 }
 
+// WithWaitDelay sets the waitDelay option. This specifies how long to wait after
+// context cancellation before forcefully killing git processes.
+func WithWaitDelay(waitDelay time.Duration) Option {
+	return func(parser *Parser) {
+		parser.waitDelay = waitDelay
+	}
+}
+
 // Option is used for adding options to Config.
 type Option func(*Parser)
 
@@ -213,6 +225,7 @@ func NewParser(options ...Option) *Parser {
 		dateFormat:    defaultDateFormat,
 		maxDiffSize:   defaultMaxDiffSize,
 		maxCommitSize: defaultMaxCommitSize,
+		waitDelay:     defaultWaitDelay,
 	}
 	for _, option := range options {
 		option(parser)
@@ -253,7 +266,7 @@ func (c *Parser) RepoPath(
 		args = append(args, "--", ".", ":(exclude)"+glob)
 	}
 
-	cmd := exec.Command("git", args...)
+	cmd := exec.CommandContext(ctx, "git", args...)
 	absPath, err := filepath.Abs(source)
 	if err == nil {
 		if !isBare {
@@ -273,26 +286,27 @@ func (c *Parser) RepoPath(
 		}
 	}
 
-	return c.executeCommand(ctx, cmd, false)
+	return c.executeCommand(ctx, cmd, false, c.waitDelay)
 }
 
 // Staged parses the output of the `git diff` command for the `source` path.
 func (c *Parser) Staged(ctx context.Context, source string) (chan *Diff, error) {
 	// Provide the --cached flag to diff to get the diff of the staged changes.
 	args := []string{"-C", source, "diff", "-p", "--cached", "--full-history", "--diff-filter=AM", "--date=iso-strict"}
 
-	cmd := exec.Command("git", args...)
+	cmd := exec.CommandContext(ctx, "git", args...)
 
 	absPath, err := filepath.Abs(source)
 	if err == nil {
 		cmd.Env = append(cmd.Env, "GIT_DIR="+filepath.Join(absPath, ".git"))
 	}
 
-	return c.executeCommand(ctx, cmd, true)
+	return c.executeCommand(ctx, cmd, true, c.waitDelay)
 }
 
 // executeCommand runs an exec.Cmd, reads stdout and stderr, and waits for the Cmd to complete.
-func (c *Parser) executeCommand(ctx context.Context, cmd *exec.Cmd, isStaged bool) (chan *Diff, error) {
+// waitDelay specifies how long to wait after context cancellation before forcefully killing the process.
+func (c *Parser) executeCommand(ctx context.Context, cmd *exec.Cmd, isStaged bool, waitDelay time.Duration) (chan *Diff, error) {
 	diffChan := make(chan *Diff, 64)
 
 	stdOut, err := cmd.StdoutPipe()
@@ -304,6 +318,9 @@ func (c *Parser) executeCommand(ctx context.Context, cmd *exec.Cmd, isStaged boo
 		return diffChan, err
 	}
 
+	// Set WaitDelay to allow the command additional time to exit after context cancellation
+	cmd.WaitDelay = waitDelay
+
 	err = cmd.Start()
 	if err != nil {
 		return diffChan, err