diff --git a/internal/appsec/dyngo/operation.go b/internal/appsec/dyngo/operation.go index b5c0a16831..f3c58d4209 100644 --- a/internal/appsec/dyngo/operation.go +++ b/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(¤t.eventRegister, op, args) + for current := op.unwrap().parent; current != nil; current = current.unwrap().parent { + emitEvent(¤t.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(¤t.eventRegister, op, results) + var current Operation = op + for ; current != nil; current = current.unwrap().parent { + emitEvent(¤t.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(¤t.dataBroadcaster, data) + for current := op; current != nil; current = current.unwrap().parent { + emitData(¤t.unwrap().dataBroadcaster, data) } } diff --git a/internal/appsec/dyngo/operation_test.go b/internal/appsec/dyngo/operation_test.go index e5ce3f1c24..35a666a98a 100644 --- a/internal/appsec/dyngo/operation_test.go +++ b/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)) +} diff --git a/internal/appsec/emitter/graphqlsec/request.go b/internal/appsec/emitter/graphqlsec/request.go index 70e1f6dff7..20ae575660 100644 --- a/internal/appsec/emitter/graphqlsec/request.go +++ b/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 diff --git a/internal/appsec/emitter/grpcsec/grpc.go b/internal/appsec/emitter/grpcsec/grpc.go index e6b28e3124..2495e8bbd5 100644 --- a/internal/appsec/emitter/grpcsec/grpc.go +++ b/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) + } } diff --git a/internal/appsec/emitter/httpsec/http.go b/internal/appsec/emitter/httpsec/http.go index 07f6dd9ba8..41ebfa7e23 100644 --- a/internal/appsec/emitter/httpsec/http.go +++ b/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 = `