internal/appsec/instrumentation: share SetEventSpanTags between grpcs…

…ec/httpsec
DataDog · Apr 7, 2022 · f7ee85c · f7ee85c
1 parent 2f5dda5
commit f7ee85c
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 104 deletions.
diff --git a/internal/appsec/dyngo/instrumentation/common.go b/internal/appsec/dyngo/instrumentation/common.go
@@ -8,9 +8,12 @@ package instrumentation
 
 import (
 	"encoding/json"
+	"fmt"
 	"sync"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 )
 
 // TagsHolder wraps a map holding tags. The purpose of this struct is to be used by composition in an Operation
@@ -64,3 +67,65 @@ func SetTags(span ddtrace.Span, tags map[string]interface{}) {
 		span.SetTag(k, v)
 	}
 }
+
+// SetEventSpanTags sets the security event span tags into the service entry span.
+func SetEventSpanTags(span ddtrace.Span, events []json.RawMessage) error {
+	// Set the appsec event span tag
+	val, err := makeEventTagValue(events)
+	if err != nil {
+		return err
+	}
+	span.SetTag("_dd.appsec.json", string(val))
+	// Keep this span due to the security event
+	//
+	// This is a workaround to tell the tracer that the trace was kept by AppSec.
+	// Passing any other value than `appsec.SamplerAppSec` has no effect.
+	// Customers should use `span.SetTag(ext.ManualKeep, true)` pattern
+	// to keep the trace, manually.
+	span.SetTag(ext.ManualKeep, samplernames.AppSec)
+	span.SetTag("_dd.origin", "appsec")
+	// Set the appsec.event tag needed by the appsec backend
+	span.SetTag("appsec.event", true)
+	return nil
+}
+
+// Create the value of the security event tag.
+// TODO(Julio-Guerra): a future libddwaf version should return something
+//   avoiding us the following events concatenation logic which currently
+//   involves unserializing the top-level JSON arrays to concatenate them
+//   together.
+// TODO(Julio-Guerra): avoid serializing the json in the request hot path
+func makeEventTagValue(events []json.RawMessage) (json.RawMessage, error) {
+	var v interface{}
+	if l := len(events); l == 1 {
+		// eventTag is the structure to use in the `_dd.appsec.json` span tag.
+		// In this case of 1 event, it already is an array as expected.
+		type eventTag struct {
+			Triggers json.RawMessage `json:"triggers"`
+		}
+		v = eventTag{Triggers: events[0]}
+	} else {
+		// eventTag is the structure to use in the `_dd.appsec.json` span tag.
+		// With more than one event, we need to concatenate the arrays together
+		// (ie. convert [][]json.RawMessage into []json.RawMessage).
+		type eventTag struct {
+			Triggers []json.RawMessage `json:"triggers"`
+		}
+		concatenated := make([]json.RawMessage, 0, l) // at least len(events)
+		for _, event := range events {
+			// Unmarshal the top level array
+			var tmp []json.RawMessage
+			if err := json.Unmarshal(event, &tmp); err != nil {
+				return nil, fmt.Errorf("unexpected error while unserializing the appsec event `%s`: %v", string(event), err)
+			}
+			concatenated = append(concatenated, tmp...)
+		}
+		v = eventTag{Triggers: concatenated}
+	}
+
+	tag, err := json.Marshal(v)
+	if err != nil {
+		return nil, fmt.Errorf("unexpected error while serializing the appsec event span tag: %v", err)
+	}
+	return tag, nil
+}
diff --git a/internal/appsec/dyngo/instrumentation/grpcsec/tags.go b/internal/appsec/dyngo/instrumentation/grpcsec/tags.go
@@ -7,14 +7,12 @@ package grpcsec
 
 import (
 	"encoding/json"
-	"fmt"
 	"net"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation/httpsec"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 )
 
 // SetSecurityEventTags sets the AppSec-specific span tags when a security event
@@ -26,7 +24,7 @@ func SetSecurityEventTags(span ddtrace.Span, events []json.RawMessage, addr net.
 }
 
 func setSecurityEventTags(span ddtrace.Span, events []json.RawMessage, addr net.Addr, md map[string][]string) error {
-	if err := setEventSpanTags(span, events); err != nil {
+	if err := instrumentation.SetEventSpanTags(span, events); err != nil {
 		return err
 	}
 	var ip string
@@ -44,65 +42,3 @@ func setSecurityEventTags(span ddtrace.Span, events []json.RawMessage, addr net.
 	}
 	return nil
 }
-
-// setEventSpanTags sets the security event span tags into the service entry span.
-func setEventSpanTags(span ddtrace.Span, events []json.RawMessage) error {
-	// Set the appsec event span tag
-	val, err := makeEventTagValue(events)
-	if err != nil {
-		return err
-	}
-	span.SetTag("_dd.appsec.json", string(val))
-	// Keep this span due to the security event
-	//
-	// This is a workaround to tell the tracer that the trace was kept by AppSec.
-	// Passing any other value than `appsec.SamplerAppSec` has no effect.
-	// Customers should use `span.SetTag(ext.ManualKeep, true)` pattern
-	// to keep the trace, manually.
-	span.SetTag(ext.ManualKeep, samplernames.AppSec)
-	span.SetTag("_dd.origin", "appsec")
-	// Set the appsec.event tag needed by the appsec backend
-	span.SetTag("appsec.event", true)
-	return nil
-}
-
-// Create the value of the security event tag.
-// TODO(Julio-Guerra): a future libddwaf version should return something
-//   avoiding us the following events concatenation logic which currently
-//   involves unserializing the top-level JSON arrays to concatenate them
-//   together.
-// TODO(Julio-Guerra): avoid serializing the json in the request hot path
-func makeEventTagValue(events []json.RawMessage) (json.RawMessage, error) {
-	var v interface{}
-	if l := len(events); l == 1 {
-		// eventTag is the structure to use in the `_dd.appsec.json` span tag.
-		// In this case of 1 event, it already is an array as expected.
-		type eventTag struct {
-			Triggers json.RawMessage `json:"triggers"`
-		}
-		v = eventTag{Triggers: events[0]}
-	} else {
-		// eventTag is the structure to use in the `_dd.appsec.json` span tag.
-		// With more than one event, we need to concatenate the arrays together
-		// (ie. convert [][]json.RawMessage into []json.RawMessage).
-		type eventTag struct {
-			Triggers []json.RawMessage `json:"triggers"`
-		}
-		concatenated := make([]json.RawMessage, 0, l) // at least len(events)
-		for _, event := range events {
-			// Unmarshal the top level array
-			var tmp []json.RawMessage
-			if err := json.Unmarshal(event, &tmp); err != nil {
-				return nil, fmt.Errorf("unexpected error while unserializing the appsec event `%s`: %v", string(event), err)
-			}
-			concatenated = append(concatenated, tmp...)
-		}
-		v = eventTag{Triggers: concatenated}
-	}
-
-	tag, err := json.Marshal(v)
-	if err != nil {
-		return nil, fmt.Errorf("unexpected error while serializing the appsec event span tag: %v", err)
-	}
-	return tag, nil
-}
diff --git a/internal/appsec/dyngo/instrumentation/httpsec/http.go b/internal/appsec/dyngo/instrumentation/httpsec/http.go
@@ -166,12 +166,9 @@ func fromContext(ctx context.Context) *Operation {
 
 // Finish the HTTP handler operation, along with the given results and emits a
 // finish event up in the operation stack.
-func (op *Operation) Finish(res HandlerOperationRes) json.RawMessage {
+func (op *Operation) Finish(res HandlerOperationRes) []json.RawMessage {
 	dyngo.FinishOperation(op, res)
-	if events := op.Events(); len(events) > 0 {
-		return events[0]
-	}
-	return json.RawMessage{}
+	return op.Events()
 }
 
 // StartSDKBodyOperation starts the SDKBody operation and emits a start event

diff --git a/internal/appsec/dyngo/instrumentation/httpsec/tags.go b/internal/appsec/dyngo/instrumentation/httpsec/tags.go
@@ -11,9 +11,7 @@ import (
 	"strings"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation"
 )
 
 // SetAppSecTags sets the AppSec-specific span tags that are expected to be in
@@ -23,35 +21,9 @@ func SetAppSecTags(span ddtrace.Span) {
 	span.SetTag("_dd.runtime_family", "go")
 }
 
-// setEventSpanTags sets the security event span tags into the service entry span.
-func setEventSpanTags(span ddtrace.Span, events json.RawMessage) {
-	// Set the appsec event span tag
-	// eventTag is the structure to use in the `_dd.appsec.json` span tag.
-	type eventTag struct {
-		Triggers json.RawMessage `json:"triggers"`
-	}
-	// TODO(Julio-Guerra): avoid serializing the json in the request hot path
-	event, err := json.Marshal(eventTag{Triggers: events})
-	if err != nil {
-		log.Error("appsec: unexpected error while serializing the appsec event span tag: %v", err)
-		return
-	}
-	span.SetTag("_dd.appsec.json", string(event))
-	// Keep this span due to the security event
-	//
-	// This is a workaround to tell the tracer that the trace was kept by AppSec.
-	// Passing any other value than `appsec.SamplerAppSec` has no effect.
-	// Customers should use `span.SetTag(ext.ManualKeep, true)` pattern
-	// to keep the trace, manually.
-	span.SetTag(ext.ManualKeep, samplernames.AppSec)
-	span.SetTag("_dd.origin", "appsec")
-	// Set the appsec.event tag needed by the appsec backend
-	span.SetTag("appsec.event", true)
-}
-
 // SetSecurityEventTags sets the AppSec-specific span tags when a security event occurred into the service entry span.
-func SetSecurityEventTags(span ddtrace.Span, events json.RawMessage, remoteIP string, headers, respHeaders map[string][]string) {
-	setEventSpanTags(span, events)
+func SetSecurityEventTags(span ddtrace.Span, events []json.RawMessage, remoteIP string, headers, respHeaders map[string][]string) {
+	instrumentation.SetEventSpanTags(span, events)
 	span.SetTag("network.client.ip", remoteIP)
 	for h, v := range NormalizeHTTPHeaders(headers) {
 		span.SetTag("http.request.headers."+h, v)

diff --git a/internal/appsec/waf.go b/internal/appsec/waf.go
@@ -24,6 +24,7 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation/httpsec"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/waf"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 )
 
 const (
@@ -148,7 +149,7 @@ func newHTTPWAFEventListener(handle *waf.Handle, addresses []string, timeout tim
 			addWAFDurationTags(&op.TagsHolder, float64(wafCtx.TotalRuntime()), float64(overallWAFRunDuration.Nanoseconds()))
 			once.Do(func() {
 				addRulesetInfoTags(&op.TagsHolder, handle.RulesetInfo())
-				op.AddTag(ext.ManualKeep, true)
+				op.AddTag(ext.ManualKeep, samplernames.AppSec)
 			})
 
 			// Log the attacks if any
@@ -233,7 +234,7 @@ func newGRPCWAFEventListener(handle *waf.Handle, _ []string, timeout time.Durati
 			addWAFDurationTags(&op.TagsHolder, float64(wafRunDuration), float64(wafBindingsRunDuration))
 			metricsOnce.Do(func() {
 				addRulesetInfoTags(&op.TagsHolder, handle.RulesetInfo())
-				op.AddTag(ext.ManualKeep, true)
+				op.AddTag(ext.ManualKeep, samplernames.AppSec)
 			})
 			if len(events) > 0 && limiter.Allow() {
 				op.AddSecurityEvents(events...)