appsec: never start a new WAF Context when one is in context

internal/appsec/dyngo/operation.go

@@ -89,7 +89,7 @@ func SwapRootOperation(new Operation) {
 // bubble-up the operation stack, which allows listening to future events that
 // might happen in the operation lifetime.
 type operation struct {
-	parent *operation
+	parent Operation
 	eventRegister
 	dataBroadcaster
 
@@ -146,11 +146,8 @@ func NewOperation(parent Operation) Operation {
 			parent = *ptr
 		}
 	}
-	var parentOp *operation
-	if parent != nil {
-		parentOp = parent.unwrap()
-	}
-	return &operation{parent: parentOp}
+
+	return &operation{parent: parent}
 }
 
 // FromContext looks into the given context (or the GLS if orchestrion is enabled) for a parent Operation and returns it.
@@ -164,13 +161,33 @@ func FromContext(ctx context.Context) (Operation, bool) {
 	return op, ok
 }
 
+// FindOperation looks into the current operation tree for the first operation matching the given type.
+// It has a hardcoded limit of 32 levels of depth even looking for the operation in the parent tree
+func FindOperation[T any, O interface {
+	Operation
+	*T
+}](ctx context.Context) (*T, bool) {
+	op, found := FromContext(ctx)
+	if !found {
+		return nil, false
+	}
+
+	for current := op; current != nil; current = current.unwrap().parent {
+		if o, ok := current.(O); ok {
+			return o, true
+		}
+	}
+
+	return nil, false
+}
+
 // StartOperation starts a new operation along with its arguments and emits a
 // start event with the operation arguments.
 func StartOperation[O Operation, E ArgOf[O]](op O, args E) {
 	// Bubble-up the start event starting from the parent operation as you can't
 	// listen for your own start event
-	for current := op.unwrap().parent; current != nil; current = current.parent {
-		emitEvent(&current.eventRegister, op, args)
+	for current := op.unwrap().parent; current != nil; current = current.unwrap().parent {
+		emitEvent(&current.unwrap().eventRegister, op, args)
 	}
 }
 
@@ -205,8 +222,9 @@ func FinishOperation[O Operation, E ResultOf[O]](op O, results E) {
 		return
 	}
 
-	for current := o; current != nil; current = current.parent {
-		emitEvent(&current.eventRegister, op, results)
+	var current Operation = op
+	for ; current != nil; current = current.unwrap().parent {
+		emitEvent(&current.unwrap().eventRegister, op, results)
 	}
 }
 
@@ -274,8 +292,8 @@ func EmitData[T any](op Operation, data T) {
 	// Bubble up the data to the stack of operations. Contrary to events,
 	// we also send the data to ourselves since SDK operations are leaf operations
 	// that both emit and listen for data (errors).
-	for current := o; current != nil; current = current.parent {
-		emitData(&current.dataBroadcaster, data)
+	for current := op; current != nil; current = current.unwrap().parent {
+		emitData(&current.unwrap().dataBroadcaster, data)
 	}
 }
 

internal/appsec/dyngo/operation_test.go

@@ -6,6 +6,7 @@
 package dyngo_test
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -17,6 +18,8 @@ import (
 	"sync/atomic"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 
 	"github.com/stretchr/testify/require"
@@ -859,3 +862,51 @@ func BenchmarkGoAssumptions(b *testing.B) {
 		})
 	})
 }
+
+func testFindOperation[T any, O interface {
+	dyngo.Operation
+	*T
+}](arg dyngo.Operation, expectOp dyngo.Operation, expectFound bool) func(t *testing.T) {
+	return func(t *testing.T) {
+		t.Helper()
+
+		op, found := dyngo.FindOperation[T, O](dyngo.RegisterOperation(context.Background(), arg))
+		assert.Equal(t, expectFound, found, "FindOperation() found = %v, want %v", found, expectFound)
+		if !found {
+			return
+		}
+
+		if expectOp == nil {
+			assert.Nil(t, op, "FindOperation() op = %v, want %v", op, expectOp)
+		} else {
+			assert.EqualValues(t, expectOp, op, "FindOperation() op = %v, want %v", op, expectOp)
+		}
+	}
+}
+
+func TestFindOperation(t *testing.T) {
+	root := dyngo.NewOperation(nil)
+	type Op1 struct{ dyngo.Operation }
+	type Op2 struct{ dyngo.Operation }
+	type Op3 struct{ dyngo.Operation }
+
+	var op1 dyngo.Operation = &Op1{root}
+	var op2 dyngo.Operation = &Op2{root}
+	var op3 dyngo.Operation = &Op3{root}
+	var parentOp1 dyngo.Operation = &Op1{dyngo.NewOperation(op3)}
+	var parentOp2 dyngo.Operation = &Op2{dyngo.NewOperation(op1)}
+	var parentOp3 dyngo.Operation = &Op3{dyngo.NewOperation(op2)}
+	var gpOp1 dyngo.Operation = &Op1{dyngo.NewOperation(parentOp1)}
+	var gpOp2 dyngo.Operation = &Op2{dyngo.NewOperation(parentOp3)}
+	var gpOp3 dyngo.Operation = &Op3{dyngo.NewOperation(parentOp2)}
+
+	t.Run("no-parent", testFindOperation[Op1](root, nil, false))
+	t.Run("found", testFindOperation[Op1](op1, op1, true))
+	t.Run("not-found", testFindOperation[Op1](op2, nil, false))
+	t.Run("found-parent", testFindOperation[Op1](parentOp2, op1, true))
+	t.Run("found-parent-2", testFindOperation[Op2](gpOp3, parentOp2, true))
+	t.Run("not-found-parent", testFindOperation[Op1](parentOp3, nil, false))
+	t.Run("found-grandparent", testFindOperation[Op1](gpOp3, op1, true))
+	t.Run("found-grandparent-2", testFindOperation[Op3](gpOp1, op3, true))
+	t.Run("not-found-grandparent", testFindOperation[Op1](gpOp2, nil, false))
+}

internal/appsec/emitter/graphqlsec/request.go

@@ -22,6 +22,9 @@ type (
 		dyngo.Operation
 		// used in case we don't have a parent operation
 		*waf.ContextOperation
+
+		// wafContextOwner indicates if the waf.ContextOperation was started by us or not and if we need to close it.
+		wafContextOwner bool
 	}
 
 	// RequestOperationArgs describes arguments passed to a GraphQL request.
@@ -43,7 +46,7 @@ type (
 // the operation stack.
 func (op *RequestOperation) Finish(span trace.TagSetter, res RequestOperationRes) {
 	dyngo.FinishOperation(op, res)
-	if op.ContextOperation != nil {
+	if op.wafContextOwner {
 		op.ContextOperation.Finish(span)
 	}
 }
@@ -56,13 +59,15 @@ func (RequestOperationRes) IsResultOf(*RequestOperation) {}
 // 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) {
-	parent, ok := dyngo.FromContext(ctx)
-	op := &RequestOperation{}
-	if !ok { // Usually we can find the HTTP Handler Operation as the parent but it's technically optional
-		op.ContextOperation, ctx = waf.StartContextOperation(ctx)
-		op.Operation = dyngo.NewOperation(op.ContextOperation)
-	} else {
-		op.Operation = dyngo.NewOperation(parent)
+	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)
+	}
+
+	op := &RequestOperation{
+		Operation:        dyngo.NewOperation(wafOp),
+		ContextOperation: wafOp,
+		wafContextOwner:  !found, // If we started the parent operation, we finish it, otherwise we don't
 	}
 
 	return dyngo.StartAndRegisterOperation(ctx, op, args), op

internal/appsec/emitter/grpcsec/grpc.go

@@ -42,6 +42,9 @@ type (
 	HandlerOperation struct {
 		dyngo.Operation
 		*waf.ContextOperation
+
+		// wafContextOwner indicates if the waf.ContextOperation was started by us or not and if we need to close it.
+		wafContextOwner bool
 	}
 
 	// HandlerOperationArgs is the grpc handler arguments.
@@ -75,10 +78,14 @@ func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}
 // 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]) {
-	wafOp, ctx := waf.StartContextOperation(ctx)
+	wafOp, found := dyngo.FindOperation[waf.ContextOperation](ctx)
+	if !found {
+		wafOp, ctx = waf.StartContextOperation(ctx)
+	}
 	op := &HandlerOperation{
 		Operation:        dyngo.NewOperation(wafOp),
 		ContextOperation: wafOp,
+		wafContextOwner:  !found, // If the parent is not found, we need to close the WAF context.
 	}
 
 	var block atomic.Pointer[actions.BlockGRPC]
@@ -112,5 +119,7 @@ func MonitorResponseMessage(ctx context.Context, msg any) error {
 // finish event up in the operation stack.
 func (op *HandlerOperation) Finish(span trace.TagSetter, res HandlerOperationRes) {
 	dyngo.FinishOperation(op, res)
-	op.ContextOperation.Finish(span)
+	if op.wafContextOwner {
+		op.ContextOperation.Finish(span)
+	}
 }

internal/appsec/emitter/httpsec/http.go

@@ -15,7 +15,6 @@ import (
 	// Blank import needed to use embed for the default blocked response payloads
 	_ "embed"
 	"net/http"
-	"sync"
 	"sync/atomic"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
@@ -31,7 +30,9 @@ type (
 	HandlerOperation struct {
 		dyngo.Operation
 		*waf.ContextOperation
-		mu sync.RWMutex
+
+		// wafContextOwner indicates if the waf.ContextOperation was started by us or not and if we need to close it.
+		wafContextOwner bool
 	}
 
 	// HandlerOperationArgs is the HTTP handler operation arguments.
@@ -57,10 +58,15 @@ func (HandlerOperationArgs) IsArgOf(*HandlerOperation)   {}
 func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}
 
 func StartOperation(ctx context.Context, args HandlerOperationArgs) (*HandlerOperation, *atomic.Pointer[actions.BlockHTTP], context.Context) {
-	wafOp, ctx := waf.StartContextOperation(ctx)
+	wafOp, found := dyngo.FindOperation[waf.ContextOperation](ctx)
+	if !found {
+		wafOp, ctx = waf.StartContextOperation(ctx)
+	}
+
 	op := &HandlerOperation{
 		Operation:        dyngo.NewOperation(wafOp),
 		ContextOperation: wafOp,
+		wafContextOwner:  !found, // If we started the parent operation, we finish it, otherwise we don't
 	}
 
 	// We need to use an atomic pointer to store the action because the action may be created asynchronously in the future
@@ -75,7 +81,9 @@ 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) {
 	dyngo.FinishOperation(op, res)
-	op.ContextOperation.Finish(span)
+	if op.wafContextOwner {
+		op.ContextOperation.Finish(span)
+	}
 }
 
 const monitorBodyErrorLog = `