Merge branch 'main' into katie.hockman/128-default

DataDog · Nov 10, 2023 · 07e257e · 07e257e
2 parents 0342acd + 9d39644
commit 07e257e
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 41 deletions.
diff --git a/contrib/google.golang.org/grpc/appsec_test.go b/contrib/google.golang.org/grpc/appsec_test.go
@@ -30,15 +30,21 @@ func TestAppSec(t *testing.T) {
 		t.Skip("appsec disabled")
 	}
 
-	rig, err := newRig(false)
-	require.NoError(t, err)
-	defer rig.Close()
+	setup := func() (FixtureClient, mocktracer.Tracer, func()) {
+		rig, err := newRig(false)
+		require.NoError(t, err)
+
+		mt := mocktracer.Start()
 
-	client := rig.client
+		return rig.client, mt, func() {
+			rig.Close()
+			mt.Stop()
+		}
+	}
 
 	t.Run("unary", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		// Send a XSS attack in the payload along with the canary value in the RPC metadata
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log"))
@@ -58,8 +64,8 @@ func TestAppSec(t *testing.T) {
 	})
 
 	t.Run("stream", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		// Send a XSS attack in the payload along with the canary value in the RPC metadata
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log"))
@@ -110,15 +116,21 @@ func TestBlocking(t *testing.T) {
 		t.Skip("appsec disabled")
 	}
 
-	rig, err := newRig(false)
-	require.NoError(t, err)
-	defer rig.Close()
+	setup := func() (FixtureClient, mocktracer.Tracer, func()) {
+		rig, err := newRig(false)
+		require.NoError(t, err)
 
-	client := rig.client
+		mt := mocktracer.Start()
+
+		return rig.client, mt, func() {
+			rig.Close()
+			mt.Stop()
+		}
+	}
 
 	t.Run("unary-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		// Send a XSS attack in the payload along with the canary value in the RPC metadata
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log", "x-client-ip", "1.2.3.4"))
@@ -136,8 +148,8 @@ func TestBlocking(t *testing.T) {
 	})
 
 	t.Run("unary-no-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, _, cleanup := setup()
+		defer cleanup()
 
 		// Send a XSS attack in the payload along with the canary value in the RPC metadata
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log", "x-client-ip", "1.2.3.5"))
@@ -148,8 +160,8 @@ func TestBlocking(t *testing.T) {
 	})
 
 	t.Run("stream-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log", "x-client-ip", "1.2.3.4"))
 		stream, err := client.StreamPing(ctx)
@@ -168,8 +180,8 @@ func TestBlocking(t *testing.T) {
 	})
 
 	t.Run("stream-no-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, _, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("dd-canary", "dd-test-scanner-log", "x-client-ip", "1.2.3.5"))
 		stream, err := client.StreamPing(ctx)
@@ -197,14 +209,21 @@ func TestUserBlocking(t *testing.T) {
 		t.Skip("appsec disabled")
 	}
 
-	rig, err := newAppsecRig(false)
-	require.NoError(t, err)
-	defer rig.Close()
-	client := rig.client
+	setup := func() (FixtureClient, mocktracer.Tracer, func()) {
+		rig, err := newAppsecRig(false)
+		require.NoError(t, err)
 
-	t.Run("unary-block", func(t *testing.T) {
 		mt := mocktracer.Start()
-		defer mt.Stop()
+
+		return rig.client, mt, func() {
+			rig.Close()
+			mt.Stop()
+		}
+	}
+
+	t.Run("unary-block", func(t *testing.T) {
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		// Send a XSS attack in the payload along with the canary value in the RPC metadata
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "blocked-user-1"))
@@ -223,8 +242,8 @@ func TestUserBlocking(t *testing.T) {
 	})
 
 	t.Run("unary-no-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, _, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "legit user"))
 		reply, err := client.Ping(ctx, &FixtureRequest{Name: "<script>alert('xss');</script>"})
@@ -236,8 +255,8 @@ func TestUserBlocking(t *testing.T) {
 	// This test checks that IP blocking happens BEFORE user blocking, since user blocking needs the request handler
 	// to be invoked while IP blocking doesn't
 	t.Run("unary-mixed-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "blocked-user-1", "x-forwarded-for", "1.2.3.4"))
 		reply, err := client.Ping(ctx, &FixtureRequest{})
@@ -253,8 +272,8 @@ func TestUserBlocking(t *testing.T) {
 	})
 
 	t.Run("stream-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "blocked-user-1"))
 		stream, err := client.StreamPing(ctx)
@@ -273,8 +292,8 @@ func TestUserBlocking(t *testing.T) {
 	})
 
 	t.Run("stream-no-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, _, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "legit user"))
 		stream, err := client.StreamPing(ctx)
@@ -293,8 +312,8 @@ func TestUserBlocking(t *testing.T) {
 	// This test checks that IP blocking happens BEFORE user blocking, since user blocking needs the request handler
 	// to be invoked while IP blocking doesn't
 	t.Run("stream-mixed-block", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
+		client, mt, cleanup := setup()
+		defer cleanup()
 
 		ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("user-id", "blocked-user-1", "x-forwarded-for", "1.2.3.4"))
 		stream, err := client.StreamPing(ctx)

diff --git a/contrib/google.golang.org/grpc/grpc_test.go b/contrib/google.golang.org/grpc/grpc_test.go
@@ -39,11 +39,6 @@ import (
 func TestUnary(t *testing.T) {
 	assert := assert.New(t)
 
-	rig, err := newRig(true, WithServiceName("grpc"), WithRequestTags())
-	require.NoError(t, err, "error setting up rig")
-	defer rig.Close()
-	client := rig.client
-
 	for name, tt := range map[string]struct {
 		message     string
 		error       bool
@@ -67,6 +62,11 @@ func TestUnary(t *testing.T) {
 		},
 	} {
 		t.Run(name, func(t *testing.T) {
+			rig, err := newRig(true, WithServiceName("grpc"), WithRequestTags())
+			require.NoError(t, err, "error setting up rig")
+			defer rig.Close()
+			client := rig.client
+
 			mt := mocktracer.Start()
 			defer mt.Stop()