From 69baa87d79d2500921d3b1cae454ff1587822df0 Mon Sep 17 00:00:00 2001 From: Eliott Bouhana Date: Mon, 25 Nov 2024 16:59:01 +0100 Subject: [PATCH 1/3] appsec: never start a new WAF Context when one is in context Signed-off-by: Eliott Bouhana --- internal/appsec/dyngo/operation.go | 44 +++++++++++----- internal/appsec/dyngo/operation_test.go | 51 +++++++++++++++++++ internal/appsec/emitter/graphqlsec/request.go | 21 +++++--- internal/appsec/emitter/grpcsec/grpc.go | 13 ++++- internal/appsec/emitter/httpsec/http.go | 16 ++++-- .../emitter/trace/service_entry_span.go | 2 + internal/appsec/emitter/waf/context.go | 2 + 7 files changed, 123 insertions(+), 26 deletions(-) diff --git a/internal/appsec/dyngo/operation.go b/internal/appsec/dyngo/operation.go index b5c0a16831..e66b85fa7b 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,35 @@ 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 + } + + const maxUnwrapping = 32 + for i := 0; i < maxUnwrapping && op != nil; i++ { + if o, ok := op.(O); ok { + return o, true + } + op = op.Parent() + } + + 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 +224,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 +294,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 = ` diff --git a/internal/appsec/emitter/trace/service_entry_span.go b/internal/appsec/emitter/trace/service_entry_span.go index 98e14b092f..e937017c37 100644 --- a/internal/appsec/emitter/trace/service_entry_span.go +++ b/internal/appsec/emitter/trace/service_entry_span.go @@ -126,6 +126,8 @@ func (op *ServiceEntrySpanOperation) OnSpanTagEvent(tag SpanTag) { op.SetTag(tag.Key, tag.Value) } +// StartServiceEntrySpanOperation starts a new ServiceEntrySpanOperation and returns it along with a new context IF one is not already started +// If a ServiceEntrySpanOperation is already started, it will return the existing one and the existing context func StartServiceEntrySpanOperation(ctx context.Context) (*ServiceEntrySpanOperation, context.Context) { parent, _ := dyngo.FromContext(ctx) op := &ServiceEntrySpanOperation{ diff --git a/internal/appsec/emitter/waf/context.go b/internal/appsec/emitter/waf/context.go index 698e721880..47678398e5 100644 --- a/internal/appsec/emitter/waf/context.go +++ b/internal/appsec/emitter/waf/context.go @@ -61,6 +61,8 @@ type ( func (ContextArgs) IsArgOf(*ContextOperation) {} func (ContextRes) IsResultOf(*ContextOperation) {} +// StartContextOperation starts a new WAF context operation and returns the operation and the new context. +// If a context operation already exist in the operation parent tree, it will return the existing operation. func StartContextOperation(ctx context.Context) (*ContextOperation, context.Context) { entrySpanOp, ctx := trace.StartServiceEntrySpanOperation(ctx) op := &ContextOperation{ From 12b69f266031e206f5cfc02e7f2af33ec924da45 Mon Sep 17 00:00:00 2001 From: Eliott Bouhana Date: Mon, 25 Nov 2024 17:01:55 +0100 Subject: [PATCH 2/3] reduce diff Signed-off-by: Eliott Bouhana --- internal/appsec/emitter/trace/service_entry_span.go | 2 -- internal/appsec/emitter/waf/context.go | 2 -- 2 files changed, 4 deletions(-) diff --git a/internal/appsec/emitter/trace/service_entry_span.go b/internal/appsec/emitter/trace/service_entry_span.go index e937017c37..98e14b092f 100644 --- a/internal/appsec/emitter/trace/service_entry_span.go +++ b/internal/appsec/emitter/trace/service_entry_span.go @@ -126,8 +126,6 @@ func (op *ServiceEntrySpanOperation) OnSpanTagEvent(tag SpanTag) { op.SetTag(tag.Key, tag.Value) } -// StartServiceEntrySpanOperation starts a new ServiceEntrySpanOperation and returns it along with a new context IF one is not already started -// If a ServiceEntrySpanOperation is already started, it will return the existing one and the existing context func StartServiceEntrySpanOperation(ctx context.Context) (*ServiceEntrySpanOperation, context.Context) { parent, _ := dyngo.FromContext(ctx) op := &ServiceEntrySpanOperation{ diff --git a/internal/appsec/emitter/waf/context.go b/internal/appsec/emitter/waf/context.go index 47678398e5..698e721880 100644 --- a/internal/appsec/emitter/waf/context.go +++ b/internal/appsec/emitter/waf/context.go @@ -61,8 +61,6 @@ type ( func (ContextArgs) IsArgOf(*ContextOperation) {} func (ContextRes) IsResultOf(*ContextOperation) {} -// StartContextOperation starts a new WAF context operation and returns the operation and the new context. -// If a context operation already exist in the operation parent tree, it will return the existing operation. func StartContextOperation(ctx context.Context) (*ContextOperation, context.Context) { entrySpanOp, ctx := trace.StartServiceEntrySpanOperation(ctx) op := &ContextOperation{ From 4e23ec9abce87377b9d7c5c20a432ff66540ef42 Mon Sep 17 00:00:00 2001 From: Eliott Bouhana Date: Tue, 26 Nov 2024 10:47:33 +0100 Subject: [PATCH 3/3] remove maxUnwrapping limit Signed-off-by: Eliott Bouhana --- internal/appsec/dyngo/operation.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/internal/appsec/dyngo/operation.go b/internal/appsec/dyngo/operation.go index e66b85fa7b..f3c58d4209 100644 --- a/internal/appsec/dyngo/operation.go +++ b/internal/appsec/dyngo/operation.go @@ -172,12 +172,10 @@ func FindOperation[T any, O interface { return nil, false } - const maxUnwrapping = 32 - for i := 0; i < maxUnwrapping && op != nil; i++ { - if o, ok := op.(O); ok { + for current := op; current != nil; current = current.unwrap().parent { + if o, ok := current.(O); ok { return o, true } - op = op.Parent() } return nil, false