feat: add keepAlive flag and fix worker flag read (americanexpress#47)

grinish21 · Jul 7, 2022 · 549081d · 549081d
1 parent e81d2cf
commit 549081d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/pkg/core/const.go b/pkg/core/const.go
@@ -42,7 +42,7 @@ func (i *arrayFlags) Set(value string) error {
 	return nil
 }
 
-//Define our static CLI flags
+// Define our static CLI flags
 var (
 	userHomeDir, _                = os.UserHomeDir()
 	levelOptions                  = utils.GetDisplayList(cfgreader.Settings.GetLevelNames())
@@ -72,4 +72,5 @@ var (
 	ptrDisplayConfidenceThreshold = flag.String("display-confidence", cfgreader.Settings.TranslateLevelID(cfgreader.Settings.DisplayConfidenceThreshold), "Lowest confidence level to display "+levelOptions)
 	ptrFailConfidenceThreshold    = flag.String("fail-confidence", cfgreader.Settings.TranslateLevelID(cfgreader.Settings.FailThreshold), "Lowest confidence level at which to fail "+levelOptions)
 	ptrModuleConfigFile           = flag.String("module-config-file", "", "Path to file with per module config settings")
+	ptrDisableHttpKeepAlives      = flag.Bool("disable-keep-alives", false, "To disable keep-alives when running as http Server. By default, keep-alives are always enabled")
 )
diff --git a/pkg/core/core.go b/pkg/core/core.go
@@ -122,6 +122,9 @@ func (eb *EarlybirdCfg) StartHTTP(ptr PTRHTTPConfig) {
 		Handler:      r,
 	}
 
+	// To control whether HTTP keep-alives are enabled or not.
+	srv.SetKeepAlivesEnabled(!*ptrDisableHttpKeepAlives)
+
 	if *ptr.HTTPS != "" {
 		srv.Addr = *ptr.HTTPS
 		err := http2.ConfigureServer(srv, &http2.Server{})

diff --git a/pkg/scan/scan.go b/pkg/scan/scan.go
@@ -79,7 +79,7 @@ func SearchFiles(cfg *cfgReader.EarlybirdConfig, files []File, compressPaths []s
 func scanPool(cfg *cfgReader.EarlybirdConfig, wg *sync.WaitGroup, jobMutex *sync.Mutex, jobs chan WorkJob, hits chan<- Hit) {
 	//Create duplicate map
 	dupeMap := make(map[string]bool) //HASH:true
-	for w := 1; w <= 100; w++ {
+	for w := 1; w <= cfg.WorkerCount; w++ {
 		wg.Add(1)
 		go func(w int) {
 			for j := range jobs {