From dc24ad26392bd035ec87ad37913665163d6ac6ae Mon Sep 17 00:00:00 2001
From: Joona Hoikkala <joohoi@users.noreply.github.com>
Date: Thu, 24 Sep 2020 12:04:31 +0300
Subject: [PATCH] Make SIGINT more responsive, and handle zombied TCP
 connections properly (#302)

---
 CHANGELOG.md         | 2 ++
 main.go              | 2 +-
 pkg/ffuf/config.go   | 4 +++-
 pkg/ffuf/job.go      | 7 ++++++-
 pkg/runner/simple.go | 7 ++++++-
 5 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 97a7385..761d23f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
 
   - Changed
     - Pre-flight errors are now displayed also after the usage text to prevent the need to scroll through backlog.
+    - Cancelling via SIGINT (Ctrl-C) is now more responsive
+    - Fixed issue where a thread would hang due to TCP errors
     - The `-w` flag now accepts comma delimited values in the form of `file1:W1,file2:W2`.
 
 - v1.1.0
diff --git a/main.go b/main.go
index cc0a4d9..2191716 100644
--- a/main.go
+++ b/main.go
@@ -82,7 +82,7 @@ func (m *wordlistFlag) Set(value string) error {
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	conf := ffuf.NewConfig(ctx)
+	conf := ffuf.NewConfig(ctx, cancel)
 	opts := cliOptions{}
 	var ignored bool
 	flag.BoolVar(&conf.IgnoreWordlistComments, "ic", false, "Ignore wordlist comments")
diff --git a/pkg/ffuf/config.go b/pkg/ffuf/config.go
index 96b2016..8cbaf09 100644
--- a/pkg/ffuf/config.go
+++ b/pkg/ffuf/config.go
@@ -35,6 +35,7 @@ type Config struct {
 	Matchers               map[string]FilterProvider `json:"matchers"`
 	Threads                int                       `json:"threads"`
 	Context                context.Context           `json:"-"`
+	Cancel                 context.CancelFunc        `json:"-"`
 	ProxyURL               string                    `json:"proxyurl"`
 	ReplayProxyURL         string                    `json:"replayproxyurl"`
 	CommandLine            string                    `json:"cmdline"`
@@ -52,9 +53,10 @@ type InputProviderConfig struct {
 	Value   string `json:"value"`
 }
 
-func NewConfig(ctx context.Context) Config {
+func NewConfig(ctx context.Context, cancel context.CancelFunc) Config {
 	var conf Config
 	conf.Context = ctx
+	conf.Cancel = cancel
 	conf.Headers = make(map[string]string)
 	conf.Method = "GET"
 	conf.Url = ""
diff --git a/pkg/ffuf/job.go b/pkg/ffuf/job.go
index 53e56b4..79871d3 100644
--- a/pkg/ffuf/job.go
+++ b/pkg/ffuf/job.go
@@ -146,7 +146,11 @@ func (j *Job) sleepIfNeeded() {
 		}
 		sleepDuration = sleepDuration * time.Millisecond
 	}
-	time.Sleep(sleepDuration)
+	// makes the sleep cancellable by context
+	select {
+	case <-j.Config.Context.Done(): // cancelled
+	case <-time.After(sleepDuration): // sleep
+	}
 }
 
 func (j *Job) startExecution() {
@@ -428,6 +432,7 @@ func (j *Job) CheckStop() {
 //Stop the execution of the Job
 func (j *Job) Stop() {
 	j.Running = false
+	j.Config.Cancel()
 	return
 }
 
diff --git a/pkg/runner/simple.go b/pkg/runner/simple.go
index fccf601..3b7cbd9 100644
--- a/pkg/runner/simple.go
+++ b/pkg/runner/simple.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/textproto"
@@ -51,6 +52,10 @@ func NewSimpleRunner(conf *ffuf.Config, replay bool) ffuf.RunnerProvider {
 			MaxIdleConns:        1000,
 			MaxIdleConnsPerHost: 500,
 			MaxConnsPerHost:     500,
+			DialContext: (&net.Dialer{
+				Timeout: time.Duration(time.Duration(conf.Timeout) * time.Second),
+			}).DialContext,
+			TLSHandshakeTimeout: time.Duration(time.Duration(conf.Timeout) * time.Second),
 			TLSClientConfig: &tls.Config{
 				InsecureSkipVerify: true,
 				Renegotiation:      tls.RenegotiateOnceAsClient,
@@ -92,7 +97,7 @@ func (r *SimpleRunner) Execute(req *ffuf.Request) (ffuf.Response, error) {
 	var err error
 	var rawreq []byte
 	data := bytes.NewReader(req.Data)
-	httpreq, err = http.NewRequest(req.Method, req.Url, data)
+	httpreq, err = http.NewRequestWithContext(r.config.Context, req.Method, req.Url, data)
 	if err != nil {
 		return ffuf.Response{}, err
 	}