appsec: stricter time budget per request and better metrics for the W…

…AF (#2613)

Signed-off-by: Eliott Bouhana <eliott.bouhana@datadoghq.com>

.gitlab-ci.yml

@@ -4,12 +4,12 @@ stages:
   - test-apps
 
 variables:
-  # This base image is created here: https://gitlab.ddbuild.io/DataDog/apm-reliability/benchmarking-platform/-/pipelines/30087677
-  BASE_CI_IMAGE: 486234852809.dkr.ecr.us-east-1.amazonaws.com/ci/benchmarking-platform:dd-trace-go-30087677
+  # This base image is created here: https://gitlab.ddbuild.io/DataDog/apm-reliability/benchmarking-platform/-/pipelines/30723596
+  BASE_CI_IMAGE: 486234852809.dkr.ecr.us-east-1.amazonaws.com/ci/benchmarking-platform:dd-trace-go-30723596
   INDEX_FILE: index.txt
   KUBERNETES_SERVICE_ACCOUNT_OVERWRITE: dd-trace-go
   FF_USE_LEGACY_KUBERNETES_EXECUTION_STRATEGY: "true"
-  BENCHMARK_TARGETS: "BenchmarkStartRequestSpan|BenchmarkHttpServeTrace|BenchmarkTracerAddSpans|BenchmarkStartSpan|BenchmarkSingleSpanRetention|BenchmarkOTelApiWithCustomTags|BenchmarkInjectW3C|BenchmarkExtractW3C|BenchmarkPartialFlushing|BenchmarkGraphQL"
+  BENCHMARK_TARGETS: "BenchmarkStartRequestSpan|BenchmarkHttpServeTrace|BenchmarkTracerAddSpans|BenchmarkStartSpan|BenchmarkSingleSpanRetention|BenchmarkOTelApiWithCustomTags|BenchmarkInjectW3C|BenchmarkExtractW3C|BenchmarkPartialFlushing|BenchmarkGraphQL|BenchmarkSampleWAFContext"
 
 include:
   - ".gitlab/benchmarks.yml"

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/DataDog/datadog-agent/pkg/obfuscate v0.48.0
 	github.com/DataDog/datadog-agent/pkg/remoteconfig/state v0.48.1
 	github.com/DataDog/datadog-go/v5 v5.3.0
-	github.com/DataDog/go-libddwaf/v2 v2.3.2
+	github.com/DataDog/go-libddwaf/v2 v2.4.2
 	github.com/DataDog/gostackparse v0.7.0
 	github.com/DataDog/sketches-go v1.4.2
 	github.com/IBM/sarama v1.40.0

go.sum

@@ -633,8 +633,8 @@ github.com/DataDog/datadog-agent/pkg/remoteconfig/state v0.48.1/go.mod h1:Vc+snp
 github.com/DataDog/datadog-go v3.2.0+incompatible/go.mod h1:LButxg5PwREeZtORoXG3tL4fMGNddJ+vMq1mwgfaqoQ=
 github.com/DataDog/datadog-go/v5 v5.3.0 h1:2q2qjFOb3RwAZNU+ez27ZVDwErJv5/VpbBPprz7Z+s8=
 github.com/DataDog/datadog-go/v5 v5.3.0/go.mod h1:XRDJk1pTc00gm+ZDiBKsjh7oOOtJfYfglVCmFb8C2+Q=
-github.com/DataDog/go-libddwaf/v2 v2.3.2 h1:pdi9xjWW57IpOpTeOyPuNveEDFLmmInsHDeuZk3TY34=
-github.com/DataDog/go-libddwaf/v2 v2.3.2/go.mod h1:gsCdoijYQfj8ce/T2bEDNPZFIYnmHluAgVDpuQOWMZE=
+github.com/DataDog/go-libddwaf/v2 v2.4.2 h1:ilquGKUmN9/Ty0sIxiEyznVRxP3hKfmH15Y1SMq5gjA=
+github.com/DataDog/go-libddwaf/v2 v2.4.2/go.mod h1:gsCdoijYQfj8ce/T2bEDNPZFIYnmHluAgVDpuQOWMZE=
 github.com/DataDog/go-tuf v1.0.2-0.5.2 h1:EeZr937eKAWPxJ26IykAdWA4A0jQXJgkhUjqEI/w7+I=
 github.com/DataDog/go-tuf v1.0.2-0.5.2/go.mod h1:zBcq6f654iVqmkk8n2Cx81E1JnNTMOAx1UEO/wZR+P0=
 github.com/DataDog/gostackparse v0.7.0 h1:i7dLkXHvYzHV308hnkvVGDL3BR4FWl7IsXNPz/IGQh4=

internal/apps/go.mod

@@ -9,7 +9,7 @@ require (
 
 require (
 	github.com/DataDog/appsec-internal-go v1.5.0 // indirect
-	github.com/DataDog/go-libddwaf/v2 v2.3.2 // indirect
+	github.com/DataDog/go-libddwaf/v2 v2.4.2 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/ebitengine/purego v0.6.0-alpha.5 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect

internal/apps/go.sum

@@ -6,8 +6,8 @@ github.com/DataDog/datadog-agent/pkg/remoteconfig/state v0.48.1 h1:5nE6N3JSs2IG3
 github.com/DataDog/datadog-agent/pkg/remoteconfig/state v0.48.1/go.mod h1:Vc+snp0Bey4MrrJyiV2tVxxJb6BmLomPvN1RgAvjGaQ=
 github.com/DataDog/datadog-go/v5 v5.3.0 h1:2q2qjFOb3RwAZNU+ez27ZVDwErJv5/VpbBPprz7Z+s8=
 github.com/DataDog/datadog-go/v5 v5.3.0/go.mod h1:XRDJk1pTc00gm+ZDiBKsjh7oOOtJfYfglVCmFb8C2+Q=
-github.com/DataDog/go-libddwaf/v2 v2.3.2 h1:pdi9xjWW57IpOpTeOyPuNveEDFLmmInsHDeuZk3TY34=
-github.com/DataDog/go-libddwaf/v2 v2.3.2/go.mod h1:gsCdoijYQfj8ce/T2bEDNPZFIYnmHluAgVDpuQOWMZE=
+github.com/DataDog/go-libddwaf/v2 v2.4.2 h1:ilquGKUmN9/Ty0sIxiEyznVRxP3hKfmH15Y1SMq5gjA=
+github.com/DataDog/go-libddwaf/v2 v2.4.2/go.mod h1:gsCdoijYQfj8ce/T2bEDNPZFIYnmHluAgVDpuQOWMZE=
 github.com/DataDog/go-tuf v1.0.2-0.5.2 h1:EeZr937eKAWPxJ26IykAdWA4A0jQXJgkhUjqEI/w7+I=
 github.com/DataDog/go-tuf v1.0.2-0.5.2/go.mod h1:zBcq6f654iVqmkk8n2Cx81E1JnNTMOAx1UEO/wZR+P0=
 github.com/DataDog/gostackparse v0.7.0 h1:i7dLkXHvYzHV308hnkvVGDL3BR4FWl7IsXNPz/IGQh4=

internal/appsec/listener/graphqlsec/graphql.go

@@ -73,7 +73,7 @@ func newWafEventListener(wafHandle *waf.Handle, cfg *config.Config, limiter limi
 // NewWAFEventListener returns the WAF event listener to register in order
 // to enable it.
 func (l *wafEventListener) onEvent(request *types.RequestOperation, _ types.RequestOperationArgs) {
-	wafCtx := waf.NewContext(l.wafHandle)
+	wafCtx := waf.NewContextWithBudget(l.wafHandle, l.config.WAFTimeout)
 	if wafCtx == nil {
 		return
 	}
@@ -113,10 +113,7 @@ func (l *wafEventListener) onEvent(request *types.RequestOperation, _ types.Requ
 	dyngo.OnFinish(request, func(request *types.RequestOperation, res types.RequestOperationRes) {
 		defer wafCtx.Close()
 
-		overall, internal := wafCtx.TotalRuntime()
-		nbTimeouts := wafCtx.TotalTimeouts()
-		shared.AddWAFMonitoringTags(request, l.wafDiags.Version, overall, internal, nbTimeouts)
-
+		shared.AddWAFMonitoringTags(request, l.wafDiags.Version, wafCtx.Stats().Metrics())
 		trace.SetEventSpanTags(request, request.Events())
 	})
 }

internal/appsec/listener/grpcsec/grpc.go

@@ -97,7 +97,7 @@ func (l *wafEventListener) onEvent(op *types.HandlerOperation, handlerArgs types
 		mu     sync.Mutex // events mutex
 	)
 
-	wafCtx := waf.NewContext(l.wafHandle)
+	wafCtx := waf.NewContextWithBudget(l.wafHandle, l.config.WAFTimeout)
 	if wafCtx == nil {
 		// The WAF event listener got concurrently released
 		return
@@ -186,8 +186,7 @@ func (l *wafEventListener) onEvent(op *types.HandlerOperation, handlerArgs types
 	// When the gRPC handler finishes
 	dyngo.OnFinish(op, func(op *types.HandlerOperation, _ types.HandlerOperationRes) {
 		defer wafCtx.Close()
-		overallRuntimeNs, internalRuntimeNs := wafCtx.TotalRuntime()
-		shared.AddWAFMonitoringTags(op, l.wafDiags.Version, overallRuntimeNs, internalRuntimeNs, wafCtx.TotalTimeouts())
+		shared.AddWAFMonitoringTags(op, l.wafDiags.Version, wafCtx.Stats().Metrics())
 
 		// Log the following metrics once per instantiation of a WAF handle
 		l.once.Do(func() {

internal/appsec/listener/httpsec/http.go

@@ -95,7 +95,7 @@ func newWafEventListener(wafHandle *waf.Handle, actions sharedsec.Actions, cfg *
 
 // NewWAFEventListener returns the WAF event listener to register in order to enable it.
 func (l *wafEventListener) onEvent(op *types.Operation, args types.HandlerOperationArgs) {
-	wafCtx := waf.NewContext(l.wafHandle)
+	wafCtx := waf.NewContextWithBudget(l.wafHandle, l.config.WAFTimeout)
 	if wafCtx == nil {
 		// The WAF event listener got concurrently released
 		return
@@ -196,8 +196,7 @@ func (l *wafEventListener) onEvent(op *types.Operation, args types.HandlerOperat
 		wafResult := shared.RunWAF(wafCtx, waf.RunAddressData{Persistent: values}, l.config.WAFTimeout)
 
 		// Add WAF metrics.
-		overallRuntimeNs, internalRuntimeNs := wafCtx.TotalRuntime()
-		shared.AddWAFMonitoringTags(op, l.wafDiags.Version, overallRuntimeNs, internalRuntimeNs, wafCtx.TotalTimeouts())
+		shared.AddWAFMonitoringTags(op, l.wafDiags.Version, wafCtx.Stats().Metrics())
 
 		// Add the following metrics once per instantiation of a WAF handle
 		l.once.Do(func() {

internal/appsec/listener/sharedsec/shared.go

@@ -43,9 +43,6 @@ const (
 	eventRulesErrorsTag  = "_dd.appsec.event_rules.errors"
 	eventRulesLoadedTag  = "_dd.appsec.event_rules.loaded"
 	eventRulesFailedTag  = "_dd.appsec.event_rules.error_count"
-	wafDurationTag       = "_dd.appsec.waf.duration"
-	wafDurationExtTag    = "_dd.appsec.waf.duration_ext"
-	wafTimeoutTag        = "_dd.appsec.waf.timeouts"
 	wafVersionTag        = "_dd.appsec.waf.version"
 )
 
@@ -70,12 +67,14 @@ func AddRulesMonitoringTags(th trace.TagSetter, wafDiags *waf.Diagnostics) {
 }
 
 // Add the tags related to the monitoring of the WAF
-func AddWAFMonitoringTags(th trace.TagSetter, rulesVersion string, overallRuntimeNs, internalRuntimeNs, timeouts uint64) {
+func AddWAFMonitoringTags(th trace.TagSetter, rulesVersion string, stats map[string]any) {
 	// Rules version is set for every request to help the backend associate WAF duration metrics with rule version
 	th.SetTag(eventRulesVersionTag, rulesVersion)
-	th.SetTag(wafTimeoutTag, timeouts)
-	th.SetTag(wafDurationTag, float64(internalRuntimeNs)/1e3)   // ns to us
-	th.SetTag(wafDurationExtTag, float64(overallRuntimeNs)/1e3) // ns to us
+
+	// Report the stats sent by the WAF
+	for k, v := range stats {
+		th.SetTag(k, v)
+	}
 }
 
 // ProcessActions sends the relevant actions to the operation's data listener.

internal/appsec/listener/sharedsec/shared_test.go

@@ -13,6 +13,12 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/trace"
 )
 
+const (
+	wafDurationTag    = "_dd.appsec.waf.duration"
+	wafDurationExtTag = "_dd.appsec.waf.duration_ext"
+	wafTimeoutTag     = "_dd.appsec.waf.timeouts"
+)
+
 // Test that internal functions used to set span tags use the correct types
 func TestTagsTypes(t *testing.T) {
 	th := trace.NewTagsHolder()
@@ -26,13 +32,22 @@ func TestTagsTypes(t *testing.T) {
 	}
 
 	AddRulesMonitoringTags(&th, &wafDiags)
-	AddWAFMonitoringTags(&th, "1.2.3", 2, 1, 3)
+
+	stats := map[string]any{
+		wafDurationTag:                     10,
+		wafDurationExtTag:                  20,
+		wafTimeoutTag:                      0,
+		"_dd.appsec.waf.truncations.depth": []int{1, 2, 3},
+		"_dd.appsec.waf.run":               12000,
+	}
+
+	AddWAFMonitoringTags(&th, "1.2.3", stats)
 
 	tags := th.Tags()
 	_, ok := tags[eventRulesErrorsTag].(string)
 	require.True(t, ok)
 
-	for _, tag := range []string{eventRulesLoadedTag, eventRulesFailedTag, wafDurationTag, wafDurationExtTag, wafVersionTag} {
+	for _, tag := range []string{eventRulesLoadedTag, eventRulesFailedTag, wafDurationTag, wafDurationExtTag, wafVersionTag, wafTimeoutTag} {
 		require.Contains(t, tags, tag)
 	}
 }

internal/appsec/waf_test.go

@@ -6,6 +6,8 @@
 package appsec_test
 
 import (
+	"encoding/json"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener/httpsec"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -491,3 +493,76 @@ func TestAPISecurity(t *testing.T) {
 		require.Nil(t, spans[0].Tag("_dd.appsec.s.req.body"))
 	})
 }
+
+// BenchmarkSampleWAFContext benchmarks the creation of a WAF context and running the WAF on a request/response pair
+// This is a basic sample of what could happen in a real-world scenario.
+func BenchmarkSampleWAFContext(b *testing.B) {
+	rules, err := internal.DefaultRuleset()
+	if err != nil {
+		b.Fatalf("error loading rules: %v", err)
+	}
+
+	var parsedRuleset map[string]any
+	err = json.Unmarshal(rules, &parsedRuleset)
+	if err != nil {
+		b.Fatalf("error parsing rules: %v", err)
+	}
+
+	handle, err := waf.NewHandle(parsedRuleset, internal.DefaultObfuscatorKeyRegex, internal.DefaultObfuscatorValueRegex)
+	for i := 0; i < b.N; i++ {
+		ctx := waf.NewContext(handle)
+		if ctx == nil {
+			b.Fatal("nil context")
+		}
+
+		// Request WAF Run
+		_, err := ctx.Run(
+			waf.RunAddressData{
+				Persistent: map[string]any{
+					httpsec.HTTPClientIPAddr:        "1.1.1.1",
+					httpsec.ServerRequestMethodAddr: "GET",
+					httpsec.ServerRequestRawURIAddr: "/",
+					httpsec.ServerRequestHeadersNoCookiesAddr: map[string][]string{
+						"host":            {"example.com"},
+						"content-length":  {"0"},
+						"Accept":          {"application/json"},
+						"User-Agent":      {"curl/7.64.1"},
+						"Accept-Encoding": {"gzip"},
+						"Connection":      {"close"},
+					},
+					httpsec.ServerRequestCookiesAddr: map[string][]string{
+						"cookie": {"session=1234"},
+					},
+					httpsec.ServerRequestQueryAddr: map[string][]string{
+						"query": {"value"},
+					},
+					httpsec.ServerRequestPathParamsAddr: map[string]string{
+						"param": "value",
+					},
+				},
+			}, 0)
+
+		if err != nil {
+			b.Fatalf("error running waf: %v", err)
+		}
+
+		// Response WAF Run
+		_, err = ctx.Run(
+			waf.RunAddressData{
+				Persistent: map[string]any{
+					httpsec.ServerResponseHeadersNoCookiesAddr: map[string][]string{
+						"content-type":   {"application/json"},
+						"content-length": {"0"},
+						"Connection":     {"close"},
+					},
+					httpsec.ServerResponseStatusAddr: 200,
+				},
+			}, 0)
+
+		if err != nil {
+			b.Fatalf("error running waf: %v", err)
+		}
+
+		ctx.Close()
+	}
+}