appsec: stop storing span tags, directly call span.SetTag

eliottness · eliottness · commit 2ac229bb8c92 · 2024-12-17T19:33:01.000+01:00
Signed-off-by: Eliott Bouhana &lt;eliott.bouhana@datadoghq.com&gt;
diff --git a/contrib/99designs/gqlgen/tracer.go b/contrib/99designs/gqlgen/tracer.go
@@ -103,7 +103,7 @@ func (t *gqlTracer) Validate(_ graphql.ExecutableSchema) error {
 func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.OperationHandler) graphql.ResponseHandler {
 	opCtx := graphql.GetOperationContext(ctx)
 	span, ctx := t.createRootSpan(ctx, opCtx)
-	ctx, req := graphqlsec.StartRequestOperation(ctx, graphqlsec.RequestOperationArgs{
+	ctx, req := graphqlsec.StartRequestOperation(ctx, span, graphqlsec.RequestOperationArgs{
 		RawQuery:      opCtx.RawQuery,
 		OperationName: opCtx.OperationName,
 		Variables:     opCtx.Variables,
@@ -137,7 +137,7 @@ func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.Operati
 		}
 
 		query.Finish(executionOperationRes)
-		req.Finish(span, requestOperationRes)
+		req.Finish(requestOperationRes)
 		return response
 	}
 }
diff --git a/contrib/google.golang.org/grpc/appsec.go b/contrib/google.golang.org/grpc/appsec.go
@@ -44,7 +44,7 @@ func appsecUnaryHandlerMiddleware(method string, span ddtrace.Span, handler grpc
 			remoteAddr = p.Addr.String()
 		}
 
-		ctx, op, blockAtomic := grpcsec.StartHandlerOperation(ctx, grpcsec.HandlerOperationArgs{
+		ctx, op, blockAtomic := grpcsec.StartHandlerOperation(ctx, span, grpcsec.HandlerOperationArgs{
 			Method:     method,
 			Metadata:   md,
 			RemoteAddr: remoteAddr,
@@ -55,7 +55,7 @@ func appsecUnaryHandlerMiddleware(method string, span ddtrace.Span, handler grpc
 			if statusErr, ok := rpcErr.(interface{ GRPCStatus() *status.Status }); ok && !applyAction(blockAtomic, &rpcErr) {
 				statusCode = int(statusErr.GRPCStatus().Code())
 			}
-			op.Finish(span, grpcsec.HandlerOperationRes{StatusCode: statusCode})
+			op.Finish(grpcsec.HandlerOperationRes{StatusCode: statusCode})
 			applyAction(blockAtomic, &rpcErr)
 		}()
 
@@ -90,7 +90,7 @@ func appsecStreamHandlerMiddleware(method string, span ddtrace.Span, handler grp
 		}
 
 		// Create the handler operation and listen to blocking gRPC actions to detect a blocking condition
-		ctx, op, blockAtomic := grpcsec.StartHandlerOperation(ctx, grpcsec.HandlerOperationArgs{
+		ctx, op, blockAtomic := grpcsec.StartHandlerOperation(ctx, span, grpcsec.HandlerOperationArgs{
 			Method:     method,
 			Metadata:   md,
 			RemoteAddr: remoteAddr,
@@ -104,7 +104,7 @@ func appsecStreamHandlerMiddleware(method string, span ddtrace.Span, handler grp
 				statusCode = int(res.Status())
 			}
 
-			op.Finish(span, grpcsec.HandlerOperationRes{StatusCode: statusCode})
+			op.Finish(grpcsec.HandlerOperationRes{StatusCode: statusCode})
 			applyAction(blockAtomic, &rpcErr)
 		}()
 
diff --git a/contrib/graph-gophers/graphql-go/graphql.go b/contrib/graph-gophers/graphql-go/graphql.go
@@ -70,7 +70,7 @@ func (t *Tracer) TraceQuery(ctx context.Context, queryString, operationName stri
 	}
 	span, ctx := ddtracer.StartSpanFromContext(ctx, t.cfg.querySpanName, opts...)
 
-	ctx, request := graphqlsec.StartRequestOperation(ctx, graphqlsec.RequestOperationArgs{
+	ctx, request := graphqlsec.StartRequestOperation(ctx, span, graphqlsec.RequestOperationArgs{
 		RawQuery:      queryString,
 		OperationName: operationName,
 		Variables:     variables,
@@ -92,7 +92,7 @@ func (t *Tracer) TraceQuery(ctx context.Context, queryString, operationName stri
 			err = fmt.Errorf("%s (and %d more errors)", errs[0], n-1)
 		}
 		defer span.Finish(ddtracer.WithError(err))
-		defer request.Finish(span, graphqlsec.RequestOperationRes{Error: err})
+		defer request.Finish(graphqlsec.RequestOperationRes{Error: err})
 		query.Finish(graphqlsec.ExecutionOperationRes{Error: err})
 	}
 }
diff --git a/contrib/graphql-go/graphql/graphql.go b/contrib/graphql-go/graphql/graphql.go
@@ -72,7 +72,7 @@ type contextData struct {
 // finish closes the top-level request operation, as well as the server span.
 func (c *contextData) finish(data any, err error) {
 	defer c.serverSpan.Finish(tracer.WithError(err))
-	c.requestOp.Finish(c.serverSpan, graphqlsec.RequestOperationRes{Data: data, Error: err})
+	c.requestOp.Finish(graphqlsec.RequestOperationRes{Data: data, Error: err})
 }
 
 var extensionName = reflect.TypeOf((*datadogExtension)(nil)).Elem().Name()
@@ -97,7 +97,7 @@ func (i datadogExtension) Init(ctx context.Context, params *graphql.Params) cont
 		tracer.Tag(ext.Component, componentName),
 		tracer.Measured(),
 	)
-	ctx, request := graphqlsec.StartRequestOperation(ctx, graphqlsec.RequestOperationArgs{
+	ctx, request := graphqlsec.StartRequestOperation(ctx, span, graphqlsec.RequestOperationArgs{
 		RawQuery:      params.RequestString,
 		Variables:     params.VariableValues,
 		OperationName: params.OperationName,
diff --git a/internal/appsec/emitter/graphqlsec/request.go b/internal/appsec/emitter/graphqlsec/request.go
@@ -44,10 +44,10 @@ type (
 
 // Finish the GraphQL query operation, along with the given results, and emit a finish event up in
 // the operation stack.
-func (op *RequestOperation) Finish(span trace.TagSetter, res RequestOperationRes) {
+func (op *RequestOperation) Finish(res RequestOperationRes) {
 	dyngo.FinishOperation(op, res)
 	if op.wafContextOwner {
-		op.ContextOperation.Finish(span)
+		op.ContextOperation.Finish()
 	}
 }
 
@@ -58,10 +58,10 @@ func (RequestOperationRes) IsResultOf(*RequestOperation) {}
 // emits a start event up in the operation stack. The operation is usually linked to tge global root
 // operation. The operation is tracked on the returned context, and can be extracted later on using
 // FromContext.
-func StartRequestOperation(ctx context.Context, args RequestOperationArgs) (context.Context, *RequestOperation) {
+func StartRequestOperation(ctx context.Context, span trace.TagSetter, args RequestOperationArgs) (context.Context, *RequestOperation) {
 	wafOp, found := dyngo.FindOperation[waf.ContextOperation](ctx)
 	if !found { // Usually we can find the HTTP Handler Operation as the parent, but it's technically optional
-		wafOp, ctx = waf.StartContextOperation(ctx)
+		wafOp, ctx = waf.StartContextOperation(ctx, span)
 	}
 
 	op := &RequestOperation{
diff --git a/internal/appsec/emitter/grpcsec/grpc.go b/internal/appsec/emitter/grpcsec/grpc.go
@@ -77,10 +77,10 @@ func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}
 // given arguments and parent operation, and emits a start event up in the
 // operation stack. When parent is nil, the operation is linked to the global
 // root operation.
-func StartHandlerOperation(ctx context.Context, args HandlerOperationArgs) (context.Context, *HandlerOperation, *atomic.Pointer[actions.BlockGRPC]) {
+func StartHandlerOperation(ctx context.Context, span trace.TagSetter, args HandlerOperationArgs) (context.Context, *HandlerOperation, *atomic.Pointer[actions.BlockGRPC]) {
 	wafOp, found := dyngo.FindOperation[waf.ContextOperation](ctx)
 	if !found {
-		wafOp, ctx = waf.StartContextOperation(ctx)
+		wafOp, ctx = waf.StartContextOperation(ctx, span)
 	}
 	op := &HandlerOperation{
 		Operation:        dyngo.NewOperation(wafOp),
@@ -117,9 +117,9 @@ func MonitorResponseMessage(ctx context.Context, msg any) error {
 
 // Finish the gRPC handler operation, along with the given results, and emit a
 // finish event up in the operation stack.
-func (op *HandlerOperation) Finish(span trace.TagSetter, res HandlerOperationRes) {
+func (op *HandlerOperation) Finish(res HandlerOperationRes) {
 	dyngo.FinishOperation(op, res)
 	if op.wafContextOwner {
-		op.ContextOperation.Finish(span)
+		op.ContextOperation.Finish()
 	}
 }
diff --git a/internal/appsec/emitter/httpsec/http.go b/internal/appsec/emitter/httpsec/http.go
@@ -19,6 +19,7 @@ import (
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/trace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/waf"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/waf/actions"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/waf/addresses"
@@ -57,10 +58,10 @@ type (
 func (HandlerOperationArgs) IsArgOf(*HandlerOperation)   {}
 func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}
 
-func StartOperation(ctx context.Context, args HandlerOperationArgs) (*HandlerOperation, *atomic.Pointer[actions.BlockHTTP], context.Context) {
+func StartOperation(ctx context.Context, args HandlerOperationArgs, span trace.TagSetter) (*HandlerOperation, *atomic.Pointer[actions.BlockHTTP], context.Context) {
 	wafOp, found := dyngo.FindOperation[waf.ContextOperation](ctx)
 	if !found {
-		wafOp, ctx = waf.StartContextOperation(ctx)
+		wafOp, ctx = waf.StartContextOperation(ctx, span)
 	}
 
 	op := &HandlerOperation{
@@ -79,10 +80,10 @@ func StartOperation(ctx context.Context, args HandlerOperationArgs) (*HandlerOpe
 }
 
 // Finish the HTTP handler operation and its children operations and write everything to the service entry span.
-func (op *HandlerOperation) Finish(res HandlerOperationRes, span ddtrace.Span) {
+func (op *HandlerOperation) Finish(res HandlerOperationRes) {
 	dyngo.FinishOperation(op, res)
 	if op.wafContextOwner {
-		op.ContextOperation.Finish(span)
+		op.ContextOperation.Finish()
 	}
 }
 
@@ -142,7 +143,7 @@ func BeforeHandle(
 		Cookies:     makeCookies(r.Cookies()),
 		QueryParams: r.URL.Query(),
 		PathParams:  pathParams,
-	})
+	}, span)
 	tr := r.WithContext(ctx)
 	var blocked atomic.Bool
 
@@ -154,7 +155,7 @@ func BeforeHandle(
 		op.Finish(HandlerOperationRes{
 			Headers:    opts.ResponseHeaderCopier(w),
 			StatusCode: statusCode,
-		}, span)
+		})
 
 		if blockPtr := blockAtomic.Swap(nil); blockPtr != nil {
 			blockPtr.Handler.ServeHTTP(w, tr)
diff --git a/internal/appsec/emitter/trace/service_entry_span.go b/internal/appsec/emitter/trace/service_entry_span.go
@@ -18,9 +18,9 @@ type (
 	// ServiceEntrySpanOperation is a dyngo.Operation that holds a the first span of a service. Usually a http or grpc span.
 	ServiceEntrySpanOperation struct {
 		dyngo.Operation
-		tags     map[string]any
-		jsonTags map[string]any
-		mu       sync.Mutex
+		jsonTags  map[string]any
+		tagSetter TagSetter
+		mu        sync.Mutex
 	}
 
 	// ServiceEntrySpanArgs is the arguments for a ServiceEntrySpanOperation
@@ -52,7 +52,7 @@ func (ServiceEntrySpanArgs) IsArgOf(*ServiceEntrySpanOperation) {}
 func (op *ServiceEntrySpanOperation) SetTag(key string, value any) {
 	op.mu.Lock()
 	defer op.mu.Unlock()
-	op.tags[key] = value
+	op.tagSetter.SetTag(key, value)
 }
 
 // SetSerializableTag adds the key/value pair to the tags to add to the service entry span.
@@ -76,7 +76,7 @@ func (op *ServiceEntrySpanOperation) SetSerializableTags(tags map[string]any) {
 func (op *ServiceEntrySpanOperation) setSerializableTag(key string, value any) {
 	switch value.(type) {
 	case string, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, bool:
-		op.tags[key] = value
+		op.tagSetter.SetTag(key, value)
 	default:
 		op.jsonTags[key] = value
 	}
@@ -87,7 +87,7 @@ func (op *ServiceEntrySpanOperation) SetTags(tags map[string]any) {
 	op.mu.Lock()
 	defer op.mu.Unlock()
 	for k, v := range tags {
-		op.tags[k] = v
+		op.tagSetter.SetTag(k, v)
 	}
 }
 
@@ -96,7 +96,7 @@ func (op *ServiceEntrySpanOperation) SetStringTags(tags map[string]string) {
 	op.mu.Lock()
 	defer op.mu.Unlock()
 	for k, v := range tags {
-		op.tags[k] = v
+		op.tagSetter.SetTag(k, v)
 	}
 }
 
@@ -126,32 +126,30 @@ func (op *ServiceEntrySpanOperation) OnSpanTagEvent(tag SpanTag) {
 	op.SetTag(tag.Key, tag.Value)
 }
 
-func StartServiceEntrySpanOperation(ctx context.Context) (*ServiceEntrySpanOperation, context.Context) {
+func StartServiceEntrySpanOperation(ctx context.Context, span TagSetter) (*ServiceEntrySpanOperation, context.Context) {
 	parent, _ := dyngo.FromContext(ctx)
 	op := &ServiceEntrySpanOperation{
 		Operation: dyngo.NewOperation(parent),
-		tags:      make(map[string]any),
-		jsonTags:  make(map[string]any),
+		jsonTags:  make(map[string]any, 2),
+		tagSetter: span,
 	}
 	return op, dyngo.StartAndRegisterOperation(ctx, op, ServiceEntrySpanArgs{})
 }
 
-func (op *ServiceEntrySpanOperation) Finish(span TagSetter) {
+func (op *ServiceEntrySpanOperation) Finish() {
+	span := op.tagSetter
 	if _, ok := span.(*NoopTagSetter); ok { // If the span is a NoopTagSetter or is nil, we don't need to set any tags
 		return
 	}
 
 	op.mu.Lock()
 	defer op.mu.Unlock()
 
-	for k, v := range op.tags {
-		span.SetTag(k, v)
-	}
-
 	for k, v := range op.jsonTags {
 		strValue, err := json.Marshal(v)
 		if err != nil {
 			log.Debug("appsec: failed to marshal tag %s: %v", k, err)
+			continue
 		}
 		span.SetTag(k, string(strValue))
 	}
diff --git a/internal/appsec/emitter/waf/context.go b/internal/appsec/emitter/waf/context.go
@@ -61,18 +61,18 @@ type (
 func (ContextArgs) IsArgOf(*ContextOperation)   {}
 func (ContextRes) IsResultOf(*ContextOperation) {}
 
-func StartContextOperation(ctx context.Context) (*ContextOperation, context.Context) {
-	entrySpanOp, ctx := trace.StartServiceEntrySpanOperation(ctx)
+func StartContextOperation(ctx context.Context, span trace.TagSetter) (*ContextOperation, context.Context) {
+	entrySpanOp, ctx := trace.StartServiceEntrySpanOperation(ctx, span)
 	op := &ContextOperation{
 		Operation:                 dyngo.NewOperation(entrySpanOp),
 		ServiceEntrySpanOperation: entrySpanOp,
 	}
 	return op, dyngo.StartAndRegisterOperation(ctx, op, ContextArgs{})
 }
 
-func (op *ContextOperation) Finish(span trace.TagSetter) {
+func (op *ContextOperation) Finish() {
 	dyngo.FinishOperation(op, ContextRes{})
-	op.ServiceEntrySpanOperation.Finish(span)
+	op.ServiceEntrySpanOperation.Finish()
 }
 
 func (op *ContextOperation) SwapContext(ctx *waf.Context) *waf.Context {

Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ func (t *gqlTracer) Validate(_ graphql.ExecutableSchema) error {`
`103`	`103`	`func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.OperationHandler) graphql.ResponseHandler {`
`104`	`104`	`opCtx := graphql.GetOperationContext(ctx)`
`105`	`105`	`span, ctx := t.createRootSpan(ctx, opCtx)`
`106`		`- ctx, req := graphqlsec.StartRequestOperation(ctx, graphqlsec.RequestOperationArgs{`
	`106`	`+ ctx, req := graphqlsec.StartRequestOperation(ctx, span, graphqlsec.RequestOperationArgs{`
`107`	`107`	`RawQuery: opCtx.RawQuery,`
`108`	`108`	`OperationName: opCtx.OperationName,`
`109`	`109`	`Variables: opCtx.Variables,`
`@@ -137,7 +137,7 @@ func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.Operati`
`137`	`137`	`}`
`138`	`138`
`139`	`139`	`query.Finish(executionOperationRes)`
`140`		`- req.Finish(span, requestOperationRes)`
	`140`	`+ req.Finish(requestOperationRes)`
`141`	`141`	`return response`
`142`	`142`	`}`
`143`	`143`	`}`
Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,7 @@ func (t *Tracer) TraceQuery(ctx context.Context, queryString, operationName stri`
`70`	`70`	`}`
`71`	`71`	`span, ctx := ddtracer.StartSpanFromContext(ctx, t.cfg.querySpanName, opts...)`
`72`	`72`
`73`		`- ctx, request := graphqlsec.StartRequestOperation(ctx, graphqlsec.RequestOperationArgs{`
	`73`	`+ ctx, request := graphqlsec.StartRequestOperation(ctx, span, graphqlsec.RequestOperationArgs{`
`74`	`74`	`RawQuery: queryString,`
`75`	`75`	`OperationName: operationName,`
`76`	`76`	`Variables: variables,`
`@@ -92,7 +92,7 @@ func (t *Tracer) TraceQuery(ctx context.Context, queryString, operationName stri`
`92`	`92`	`err = fmt.Errorf("%s (and %d more errors)", errs[0], n-1)`
`93`	`93`	`}`
`94`	`94`	`defer span.Finish(ddtracer.WithError(err))`
`95`		`- defer request.Finish(span, graphqlsec.RequestOperationRes{Error: err})`
	`95`	`+ defer request.Finish(graphqlsec.RequestOperationRes{Error: err})`
`96`	`96`	`query.Finish(graphqlsec.ExecutionOperationRes{Error: err})`
`97`	`97`	`}`
`98`	`98`	`}`