use mutex to lock goroutine

joshsagredo · Jan 3, 2022 · 6ea74a1 · 6ea74a1
1 parent 470f014
commit 6ea74a1
Showing 1 changed file with 9 additions and 18 deletions.
diff --git a/internal/raw/raw_test.go b/internal/raw/raw_test.go
@@ -4,11 +4,13 @@ import (
 	"context"
 	"github.com/stretchr/testify/assert"
 	"math/rand"
+	"sync"
 	"testing"
 	"time"
 )
 
 func TestStartFlooding(t *testing.T) {
+	var m sync.Mutex
 	srcIps := getIps()
 	srcPorts := getPorts()
 	macAddrs := getMacAddrs()
@@ -19,39 +21,28 @@ func TestStartFlooding(t *testing.T) {
 		srcIp, dstIp                    string
 		srcMacAddr, dstMacAddr          []byte
 	}{
-		{"10byte_syn", "syn", 10, srcPorts[rand.Intn(len(srcPorts))],
-			443, 10, srcIps[rand.Intn(len(srcIps))], "213.238.175.187",
+		{"500byte_syn", "syn", 500, srcPorts[rand.Intn(len(srcPorts))], 443, 100,
+			srcIps[rand.Intn(len(srcIps))], "213.238.175.187",
 			macAddrs[rand.Intn(len(macAddrs))], macAddrs[rand.Intn(len(macAddrs))]},
-		{
-			"10byte_ack", "ack", 10, srcPorts[rand.Intn(len(srcPorts))],
-			443, 10, srcIps[rand.Intn(len(srcIps))], "213.238.175.187",
-			macAddrs[rand.Intn(len(macAddrs))], macAddrs[rand.Intn(len(macAddrs))],
-		},
-		{
-			"10byte_synack", "synAck", 10, srcPorts[rand.Intn(len(srcPorts))],
-			443, 10, srcIps[rand.Intn(len(srcIps))], "213.238.175.187",
-			macAddrs[rand.Intn(len(macAddrs))], macAddrs[rand.Intn(len(macAddrs))],
-		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			ctx, cancel := context.WithTimeout(context.Background(), time.Duration(tc.floodMilliSeconds)*time.Millisecond)
 			defer cancel()
-
+			t.Logf("starting flood, caseName=%s, floodType=%s, floodMilliSeconds=%d\n", tc.name, tc.floodType, tc.floodMilliSeconds)
 			go func() {
-				t.Logf("starting flood, caseName=%s, floodType=%s, floodMilliSeconds=%d\n", tc.name, tc.floodType, tc.floodMilliSeconds)
+				m.Lock()
+				defer m.Unlock()
 				err := StartFlooding(tc.dstIp, tc.dstPort, tc.payloadLength, tc.floodType)
 				assert.Nil(t, err)
-				t.Logf("ending flood, caseName=%s, floodType=%s, floodMilliSeconds=%d\n", tc.name, tc.floodType, tc.floodMilliSeconds)
-				cancel()
 			}()
 
 			select {
-			case <-ctx.Done():
-				t.Logf("context closed, caseName=%s, floodType=%s, floodMilliSeconds=%d\n", tc.name, tc.floodType, tc.floodMilliSeconds)
 			case <-time.After(120 * time.Second):
 				t.Log("overslept")
+			case <-ctx.Done():
+				t.Logf("ending flood, caseName=%s, floodType=%s, floodMilliSeconds=%d\n", tc.name, tc.floodType, tc.floodMilliSeconds)
 			}
 		})
 	}