From 8d2dd624e447c4d24125937c01664c8f7c780425 Mon Sep 17 00:00:00 2001
From: Dustin Decker <dustin@trufflesec.com>
Date: Sun, 3 Apr 2022 21:13:39 -0700
Subject: [PATCH] fetcher controls fetch interval, prevent fetching loop

---
 main.go                | 32 +++++++++++++++++---------------
 pkg/engine/engine.go   |  6 +++---
 pkg/updater/updater.go | 12 +++++++++++-
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/main.go b/main.go
index 690af03c..4d6b0558 100644
--- a/main.go
+++ b/main.go
@@ -14,15 +14,14 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/felixge/fgprof"
+	"github.com/gorilla/mux"
 	"github.com/jpillora/overseer"
+	"github.com/sirupsen/logrus"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/updater"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/version"
 	"gopkg.in/alecthomas/kingpin.v2"
 
-	"github.com/felixge/fgprof"
-
-	"github.com/gorilla/mux"
-	"github.com/sirupsen/logrus"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/decoders"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/engine"
@@ -91,6 +90,16 @@ func init() {
 	}
 
 	cmd = kingpin.MustParse(cli.Parse(os.Args[1:]))
+
+	if *jsonOut {
+		logrus.SetFormatter(&logrus.JSONFormatter{})
+	}
+	if *debug {
+		logrus.SetLevel(logrus.DebugLevel)
+		logrus.Debugf("running version %s", version.BuildVersion)
+	} else {
+		logrus.SetLevel(logrus.InfoLevel)
+	}
 }
 
 func main() {
@@ -112,9 +121,11 @@ func main() {
 }
 
 func run(state overseer.State) {
-	if *versionFlag {
+	if *debug || *versionFlag {
 		fmt.Println("trufflehog " + version.BuildVersion)
-		return
+		if *versionFlag {
+			return
+		}
 	}
 
 	// When setting a base commit, chunks must be scanned in order.
@@ -122,15 +133,6 @@ func run(state overseer.State) {
 		*concurrency = 1
 	}
 
-	if *jsonOut {
-		logrus.SetFormatter(&logrus.JSONFormatter{})
-	}
-	if *debug {
-		logrus.SetLevel(logrus.DebugLevel)
-	} else {
-		logrus.SetLevel(logrus.InfoLevel)
-	}
-
 	if *debug {
 		go func() {
 			router := mux.NewRouter()
diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
index 566f3ffa..b450fb9f 100644
--- a/pkg/engine/engine.go
+++ b/pkg/engine/engine.go
@@ -3,8 +3,6 @@ package engine
 import (
 	"bytes"
 	"context"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"runtime"
 	"strings"
 	"sync"
@@ -15,6 +13,8 @@ import (
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/decoders"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 )
 
@@ -73,7 +73,7 @@ func Start(ctx context.Context, options ...EngineOption) *Engine {
 		logrus.Warn("No concurrency specified, defaulting to ", numCPU)
 		e.concurrency = numCPU
 	}
-	logrus.Debugf("running with  up to %d workers", e.concurrency)
+	logrus.Debugf("running with up to %d workers", e.concurrency)
 
 	var workerWg sync.WaitGroup
 	for i := 0; i < e.concurrency; i++ {
diff --git a/pkg/updater/updater.go b/pkg/updater/updater.go
index c41face4..c6ed9188 100644
--- a/pkg/updater/updater.go
+++ b/pkg/updater/updater.go
@@ -29,6 +29,7 @@ func Fetcher(version string) fetcher.Interface {
 type OSS struct {
 	Interval       time.Duration
 	CurrentVersion string
+	Updated        bool
 }
 
 // Init validates the provided config
@@ -49,7 +50,10 @@ type FormData struct {
 
 // Fetch binary from URL via OSS client
 func (g *OSS) Fetch() (io.Reader, error) {
-	log.Debug("fetching trufflehog update")
+	if g.Updated {
+		select {} // block until exit
+	}
+	g.Updated = true
 
 	zone, _ := time.Now().Zone()
 	data := &FormData{
@@ -71,6 +75,12 @@ func (g *OSS) Fetch() (io.Reader, error) {
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode == http.StatusNoContent {
+		return nil, nil
+	}
+
+	log.Debug("fetching trufflehog update")
+
 	newBinBytes, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
 		return nil, err