From 5d369f897110bcb1613f11bc004917ef44e86584 Mon Sep 17 00:00:00 2001
From: Fabricio <fabriciogiac@gmail.com>
Date: Sat, 17 Nov 2018 07:49:18 -0200
Subject: [PATCH] fix concurrency problem with captureID

---
 main.go      |  4 ++++
 main_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/main.go b/main.go
index e5da958..bbefd63 100644
--- a/main.go
+++ b/main.go
@@ -12,6 +12,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/googollee/go-socket.io"
 )
@@ -88,6 +89,7 @@ func dashboardItemInfoHandler() http.Handler {
 func proxyHandler(config Config) http.Handler {
 	url, _ := url.Parse(config.TargetURL)
 	captureID := 0
+	mux := sync.Mutex{}
 	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		req.Host = url.Host
 		req.URL.Host = url.Host
@@ -104,6 +106,7 @@ func proxyHandler(config Config) http.Handler {
 			if err != nil {
 				return fmt.Errorf("could not dump response: %v", err)
 			}
+			mux.Lock()
 			captureID++
 			capture := Capture{
 				ID:       captureID,
@@ -115,6 +118,7 @@ func proxyHandler(config Config) http.Handler {
 			}
 			captures.Add(capture)
 			captures.RemoveLastAfterReaching(config.MaxCaptures)
+			mux.Unlock()
 			emitToDashboard(captures)
 			return nil
 		}
diff --git a/main_test.go b/main_test.go
index 58095fe..25425e5 100644
--- a/main_test.go
+++ b/main_test.go
@@ -6,10 +6,14 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math/rand"
 	"net/http"
 	"net/http/httptest"
+	"sort"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 )
 
 // Test the reverse proxy handler
@@ -152,6 +156,48 @@ func TestDumpResponseGzip(t *testing.T) {
 	}
 }
 
+func TestCaptureIDConcurrence(t *testing.T) {
+
+	// This test bothers me
+
+	interactions := 1000
+
+	// Startup servers
+	service := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
+		rw.WriteHeader(http.StatusOK)
+	}))
+	capture := httptest.NewServer(proxyHandler(Config{TargetURL: service.URL, MaxCaptures: interactions}))
+	defer service.Close()
+	defer capture.Close()
+
+	// Starts go routines so that captureID is incremented concurrently within proxyHandler()
+	wg := &sync.WaitGroup{}
+	wg.Add(interactions)
+	for i := 0; i < interactions; i++ {
+		go func() {
+			_, err := http.Get(capture.URL)
+			if err != nil {
+				t.Errorf("Request Failed: %v", err)
+			}
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	// Tests if captures IDs are sequential
+	ids := make([]int, len(captures))
+	for i := 0; i < len(captures); i++ {
+		ids[i] = captures[i].ID
+	}
+	sort.Ints(ids)
+	for i := 0; i < len(captures); i++ {
+		if ids[i] != i+1 {
+			t.Fatalf("Capture IDs are not sequential")
+		}
+	}
+}
+
 func gzipStr(str string) string {
 	var buff bytes.Buffer
 	g := gzip.NewWriter(&buff)