From 1e9e09152b0b40f614fe892c8a1d937a888e6995 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Mazeau?= <francois.mazeau@datadoghq.com>
Date: Wed, 6 Dec 2023 12:24:01 +0100
Subject: [PATCH] internal/appsec: serialize API Sec schemas

---
 internal/appsec/listener/httpsec/http.go     | 16 ++++++++++------
 internal/appsec/listener/sharedsec/shared.go | 13 +++++++++----
 internal/appsec/waf_test.go                  |  1 -
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/internal/appsec/listener/httpsec/http.go b/internal/appsec/listener/httpsec/http.go
index 7ef739912c..d12ec5e2f9 100644
--- a/internal/appsec/listener/httpsec/http.go
+++ b/internal/appsec/listener/httpsec/http.go
@@ -124,9 +124,7 @@ func NewWAFEventListener(handle *waf.Handle, actions emitter.Actions, addresses
 		}
 
 		wafResult := listener.RunWAF(wafCtx, runData, timeout)
-		if wafResult.HasDerivatives() {
-			listener.AddTags(op, wafResult.Derivatives)
-		}
+		listener.AddAPISecurityTags(op, wafResult.Derivatives)
 		if wafResult.HasActions() || wafResult.HasEvents() {
 			interrupt := listener.ProcessActions(op, actions, wafResult.Actions)
 			listener.AddSecurityEvents(op, limiter, wafResult.Events)
@@ -139,7 +137,13 @@ func NewWAFEventListener(handle *waf.Handle, actions emitter.Actions, addresses
 
 		if _, ok := addresses[ServerRequestBodyAddr]; ok {
 			op.On(httpsec.OnSDKBodyOperationStart(func(sdkBodyOp *httpsec.SDKBodyOperation, args httpsec.SDKBodyOperationArgs) {
+				runData.Persistent = make(map[string]any, 2)
+				runData.Persistent[ServerRequestBodyAddr] = args.Body
+				if extractSchemas {
+					runData.Persistent["waf.context.processor"] = map[string]any{"extract-schema": true}
+				}
 				wafResult := listener.RunWAF(wafCtx, waf.RunAddressData{Persistent: map[string]any{ServerRequestBodyAddr: args.Body}}, timeout)
+				listener.AddAPISecurityTags(op, wafResult.Derivatives)
 				if wafResult.HasActions() || wafResult.HasEvents() {
 					listener.ProcessHTTPSDKAction(sdkBodyOp, actions, wafResult.Actions)
 					listener.AddSecurityEvents(op, limiter, wafResult.Events)
@@ -184,13 +188,13 @@ func NewWAFEventListener(handle *waf.Handle, actions emitter.Actions, addresses
 				log.Debug("appsec: attack detected by the waf")
 				listener.AddSecurityEvents(op, limiter, wafResult.Events)
 			}
-			if wafResult.HasDerivatives() {
-				listener.AddTags(op, wafResult.Derivatives)
-			}
+			listener.AddAPISecurityTags(op, wafResult.Derivatives)
 		}))
 	})
 }
 
+// canExtractSchemas checks that API Security is enabled and that sampling rate
+// allows extracting schemas
 func canExtractSchemas(cfg *internal.APISecConfig) bool {
 	return cfg != nil && cfg.Enabled && cfg.SampleRate >= rand.Float64()
 }
diff --git a/internal/appsec/listener/sharedsec/shared.go b/internal/appsec/listener/sharedsec/shared.go
index f5ab7dfaef..f653404b03 100644
--- a/internal/appsec/listener/sharedsec/shared.go
+++ b/internal/appsec/listener/sharedsec/shared.go
@@ -82,10 +82,15 @@ func AddWAFMonitoringTags(th tagsHolder, rulesVersion string, overallRuntimeNs,
 	th.AddTag(wafDurationExtTag, float64(overallRuntimeNs)/1e3) // ns to us
 }
 
-// AddTags adds arbitrary tags to the provided tags holder
-func AddTags(th tagsHolder, tags map[string]any) {
-	for k, v := range tags {
-		th.AddTag(k, v)
+// AddAPISecurityTags serializes the WAF derivatives and adds them to the tags
+func AddAPISecurityTags(th tagsHolder, derivatives map[string]any) {
+	for k, v := range derivatives {
+		schema, err := json.Marshal(v)
+		if err != nil {
+			log.Debug("appsec: could not serialize API Security schema for %s", k)
+			continue
+		}
+		th.AddTag(k, schema)
 	}
 }
 
diff --git a/internal/appsec/waf_test.go b/internal/appsec/waf_test.go
index b520bf9723..a6b29f8fb6 100644
--- a/internal/appsec/waf_test.go
+++ b/internal/appsec/waf_test.go
@@ -438,7 +438,6 @@ func TestAPISecurity(t *testing.T) {
 
 	req, err := http.NewRequest("POST", srv.URL+"/apisec?vin=AAAAAAAAAAAAAAAAA", nil)
 	require.NoError(t, err)
-	req.Header.Set("User-Agent", "dd-test-scanner-log")
 
 	t.Run("enabled", func(t *testing.T) {
 		t.Setenv("DD_EXPERIMENTAL_API_SECURITY_ENABLED", "true")