feat: remove timeout metrics and rework timer usages (#122)

This PR reworks go-libddwaf context and run API to better work with
timers. It also removes timeout summing that is now done on dd-trace-go
side.

- [x] Change `NewContext()` so it takes a list of timer options instead
of simply the time budget for this context.
- [x] Add a `TimerKey` to the parameters passed to `Context.Run()` to
select the component where this run duration will be summed in (ex: rasp
or waf)
- [x] Create 3 public constants `EncodeTimeKey`, `DecodeTimeKey` and
`DurationTimeKey` that all represent a part of a single WAF run that are
supposed to be used as key to the map `Result.TimerStats` that was added
as return value to `Context.Run()`.
- [x] Rework tests to they work with the changes
- [x] Adapt the README.md file to surface the main changes to the API
- [x] Remove metrics.go

---------

Signed-off-by: Eliott Bouhana <[email protected]>

Author: eliottness

README.md

@@ -33,13 +33,14 @@ func main() {
     }
     defer wafHandle.Close()
 
-    wafCtx := wafHandle.NewContext()
+    wafCtx := wafHandle.NewContext(timer.WithUnlimitedBudget(), timer.WithComponent("waf", "rasp"))
     defer wafCtx.Close()
 
     matches, actions := wafCtx.Run(RunAddressData{
         Persistent: map[string]any{
             "server.request.path_params": "/rfiinc.txt",
         },
+		TimerKey: "waf",
     })
 }
 ```

builder_test.go

@@ -71,7 +71,7 @@ func TestBuilder(t *testing.T) {
 		require.NotNil(t, handle)
 		defer handle.Close()
 
-		ctx, err := handle.NewContext(timer.UnlimitedBudget)
+		ctx, err := handle.NewContext(timer.WithBudget(timer.UnlimitedBudget))
 		require.NoError(t, err)
 		require.NotNil(t, ctx)
 		defer ctx.Close()
@@ -195,7 +195,7 @@ func TestBuilder(t *testing.T) {
 		waf := builder.Build()
 		require.NotNil(t, waf)
 		defer waf.Close()
-		ctx, err := waf.NewContext(timer.UnlimitedBudget)
+		ctx, err := waf.NewContext(timer.WithBudget(timer.UnlimitedBudget))
 		require.NoError(t, err)
 		require.NotNil(t, ctx)
 		defer ctx.Close()
@@ -236,7 +236,7 @@ func TestBuilder(t *testing.T) {
 		waf = builder.Build()
 		require.NotNil(t, waf)
 		defer waf.Close()
-		ctx, err = waf.NewContext(timer.UnlimitedBudget)
+		ctx, err = waf.NewContext(timer.WithBudget(timer.UnlimitedBudget))
 		require.NoError(t, err)
 		require.NotNil(t, ctx)
 		defer ctx.Close()
@@ -369,7 +369,7 @@ func TestBuilder(t *testing.T) {
 		require.NotNil(t, waf)
 		defer waf.Close()
 
-		ctx, err := waf.NewContext(time.Hour)
+		ctx, err := waf.NewContext(timer.WithBudget(time.Hour))
 		require.NoError(t, err)
 		defer ctx.Close()
 

context.go

@@ -6,9 +6,9 @@
 package libddwaf
 
 import (
+	"maps"
 	"runtime"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/DataDog/go-libddwaf/v4/internal/bindings"
@@ -22,25 +22,19 @@ import (
 // its own [Context]. New [Context] instances can be created by calling
 // [Handle.NewContext].
 type Context struct {
+	// Timer registers the time spent in the WAF and go-libddwaf. It is created alongside the Context using the options
+	// passed in to NewContext. Once its time budget is exhausted, each new call to Context.Run will return a timeout error.
+	Timer timer.NodeTimer
+
 	handle *Handle // Instance of the WAF
 
 	cContext bindings.WAFContext // The C ddwaf_context pointer
 
-	// timeoutCount count all calls which have timeout'ed by scope. Keys are fixed at creation time.
-	timeoutCount map[Scope]*atomic.Uint64
-
-	// mutex protecting the use of cContext which is not thread-safe and cgoRefs.
+	// mutex protecting the use of cContext which is not thread-safe and truncations
 	mutex sync.Mutex
 
-	// timer registers the time spent in the WAF and go-libddwaf
-	timer timer.NodeTimer
-
-	// metrics stores the cumulative time spent in various parts of the WAF
-	metrics metricsStore
-
-	// truncations provides details about truncations that occurred while encoding address data for
-	// WAF execution.
-	truncations map[Scope]map[TruncationReason][]int
+	// truncations provides details about truncations that occurred while encoding address data for the WAF execution.
+	truncations map[TruncationReason][]int
 
 	// pinner is used to retain Go data that is being passed to the WAF as part of
 	// [RunAddressData.Persistent] until the [Context.Close] method results in the context being
@@ -49,109 +43,111 @@ type Context struct {
 }
 
 // RunAddressData provides address data to the [Context.Run] method. If a given key is present in
-// both [RunAddressData.Persistent] and [RunAddressData.Ephemeral], the value from
-// [RunAddressData.Persistent] will take precedence.
+// both `Persistent` and `Ephemeral`, the value from `Persistent` will take precedence.
 // When encoding Go structs to the WAF-compatible format, fields with the `ddwaf:"ignore"` tag are
 // ignored and will not be visible to the WAF.
 type RunAddressData struct {
 	// Persistent address data is scoped to the lifetime of a given Context, and subsquent calls to
-	// [Context.Run] with the same address name will be silently ignored.
+	// Context.Run with the same address name will be silently ignored.
 	Persistent map[string]any
-	// Ephemeral address data is scoped to a given [Context.Run] call and is not persisted across
+	// Ephemeral address data is scoped to a given Context.Run call and is not persisted across
 	// calls. This is used for protocols such as gRPC client/server streaming or GraphQL, where a
 	// single request can incur multiple subrequests.
 	Ephemeral map[string]any
-	// Scope is the way to classify the different runs in the same context in order to have different
-	// metrics.
-	Scope Scope
+
+	// TimerKey is the key used to track the time spent in the WAF for this run.
+	// If left empty, a new timer with unlimited budget is started.
+	TimerKey timer.Key
 }
 
 func (d RunAddressData) isEmpty() bool {
 	return len(d.Persistent) == 0 && len(d.Ephemeral) == 0
 }
 
+// newTimer creates a new timer for this run. If the TimerKey is empty, a new timer without taking the parent into account is created.
+func (d RunAddressData) newTimer(parent timer.NodeTimer) (timer.NodeTimer, error) {
+	if d.TimerKey == "" {
+		return timer.NewTreeTimer(
+			timer.WithComponents(
+				EncodeTimeKey,
+				DurationTimeKey,
+				DecodeTimeKey,
+			),
+			timer.WithBudget(parent.SumRemaining()),
+		)
+	}
+
+	return parent.NewNode(d.TimerKey,
+		timer.WithComponents(
+			EncodeTimeKey,
+			DurationTimeKey,
+			DecodeTimeKey,
+		),
+		timer.WithInheritedSumBudget(),
+	)
+}
+
 // Run encodes the given [RunAddressData] values and runs them against the WAF rules.
 // Callers must check the returned [Result] object even when an error is returned, as the WAF might
 // have been able to match some rules and generate events or actions before the error was reached;
 // especially when the error is [waferrors.ErrTimeout].
 func (context *Context) Run(addressData RunAddressData) (res Result, err error) {
 	if addressData.isEmpty() {
-		return
-	}
-
-	if addressData.Scope == "" {
-		addressData.Scope = DefaultScope
+		return Result{}, nil
 	}
 
-	defer func() {
-		if err == waferrors.ErrTimeout {
-			context.timeoutCount[addressData.Scope].Add(1)
-		}
-	}()
-
 	// If the context has already timed out, we don't need to run the WAF again
-	if context.timer.SumExhausted() {
+	if context.Timer.SumExhausted() {
 		return Result{}, waferrors.ErrTimeout
 	}
 
-	runTimer, err := context.timer.NewNode(wafRunTag,
-		timer.WithComponents(
-			wafEncodeTag,
-			wafDecodeTag,
-			wafDurationTag,
-		),
-	)
+	runTimer, err := addressData.newTimer(context.Timer)
 	if err != nil {
 		return Result{}, err
 	}
 
-	runTimer.Start()
 	defer func() {
-		context.metrics.add(addressData.Scope, wafRunTag, runTimer.Stop())
-		context.metrics.merge(addressData.Scope, runTimer.Stats())
+		res.TimerStats = runTimer.Stats()
 	}()
 
-	wafEncodeTimer := runTimer.MustLeaf(wafEncodeTag)
+	runTimer.Start()
+	defer runTimer.Stop()
+
+	wafEncodeTimer := runTimer.MustLeaf(EncodeTimeKey)
 	wafEncodeTimer.Start()
-	persistentData, err := context.encodeOneAddressType(&context.pinner, addressData.Scope, addressData.Persistent, wafEncodeTimer)
+	defer wafEncodeTimer.Stop()
+
+	persistentData, err := context.encodeOneAddressType(&context.pinner, addressData.Persistent, wafEncodeTimer)
 	if err != nil {
-		wafEncodeTimer.Stop()
-		return res, err
+		return Result{}, err
 	}
 
 	// The WAF releases ephemeral address data at the max of each run call, so we need not keep the Go
 	// values live beyond that in the same way we need for persistent data. We hence use a separate
 	// encoder.
 	var ephemeralPinner runtime.Pinner
 	defer ephemeralPinner.Unpin()
-	ephemeralData, err := context.encodeOneAddressType(&ephemeralPinner, addressData.Scope, addressData.Ephemeral, wafEncodeTimer)
+	ephemeralData, err := context.encodeOneAddressType(&ephemeralPinner, addressData.Ephemeral, wafEncodeTimer)
 	if err != nil {
-		wafEncodeTimer.Stop()
-		return res, err
+		return Result{}, err
 	}
 
 	wafEncodeTimer.Stop()
 
-	// ddwaf_run cannot run concurrently and we are going to mutate the context.cgoRefs, so we need to
-	// lock the context
+	// ddwaf_run cannot run concurrently, so we need to lock the context
 	context.mutex.Lock()
 	defer context.mutex.Unlock()
 
 	if context.cContext == 0 {
 		// Context has been closed, returning an empty result...
-		return res, waferrors.ErrContextClosed
+		return Result{}, waferrors.ErrContextClosed
 	}
 
 	if runTimer.SumExhausted() {
-		return res, waferrors.ErrTimeout
+		return Result{}, waferrors.ErrTimeout
 	}
 
-	wafDecodeTimer := runTimer.MustLeaf(wafDecodeTag)
-	res, err = context.run(persistentData, ephemeralData, wafDecodeTimer, runTimer.SumRemaining())
-
-	runTimer.AddTime(wafDurationTag, res.TimeSpent)
-
-	return
+	return context.run(persistentData, ephemeralData, runTimer)
 }
 
 // merge merges two maps of slices into a single map of slices. The resulting map will contain all
@@ -195,7 +191,7 @@ func merge[K comparable, V any](a, b map[K][]V) (merged map[K][]V) {
 // top level object is a nil map, but this  behaviour is expected since either persistent or
 // ephemeral addresses are allowed to be null one at a time. In this case, Encode will return nil,
 // which is what we need to send to ddwaf_run to signal that the address data is empty.
-func (context *Context) encodeOneAddressType(pinner pin.Pinner, scope Scope, addressData map[string]any, timer timer.Timer) (*bindings.WAFObject, error) {
+func (context *Context) encodeOneAddressType(pinner pin.Pinner, addressData map[string]any, timer timer.Timer) (*bindings.WAFObject, error) {
 	encoder := newLimitedEncoder(pinner, timer)
 	if addressData == nil {
 		return nil, nil
@@ -206,7 +202,7 @@ func (context *Context) encodeOneAddressType(pinner pin.Pinner, scope Scope, add
 		context.mutex.Lock()
 		defer context.mutex.Unlock()
 
-		context.truncations[scope] = merge(context.truncations[scope], encoder.truncations)
+		context.truncations = merge(context.truncations, encoder.truncations)
 	}
 
 	if timer.Exhausted() {
@@ -218,22 +214,25 @@ func (context *Context) encodeOneAddressType(pinner pin.Pinner, scope Scope, add
 
 // run executes the ddwaf_run call with the provided data on this context. The caller is responsible for locking the
 // context appropriately around this call.
-func (context *Context) run(persistentData, ephemeralData *bindings.WAFObject, wafDecodeTimer timer.Timer, timeBudget time.Duration) (Result, error) {
+func (context *Context) run(persistentData, ephemeralData *bindings.WAFObject, runTimer timer.NodeTimer) (Result, error) {
 	result := new(bindings.WAFResult)
 	defer wafLib.ResultFree(result)
 
 	// The value of the timeout cannot exceed 2^55
 	// cf. https://en.cppreference.com/w/cpp/chrono/duration
-	timeout := uint64(timeBudget.Microseconds()) & 0x008FFFFFFFFFFFFF
+	timeout := uint64(runTimer.SumRemaining().Microseconds()) & 0x008FFFFFFFFFFFFF
 	ret := wafLib.Run(context.cContext, persistentData, ephemeralData, result, timeout)
 
-	wafDecodeTimer.Start()
-	defer wafDecodeTimer.Stop()
+	decodeTimer := runTimer.MustLeaf(DecodeTimeKey)
+	decodeTimer.Start()
+	defer decodeTimer.Stop()
 
-	return unwrapWafResult(ret, result)
+	res, duration, err := unwrapWafResult(ret, result)
+	runTimer.AddTime(DurationTimeKey, duration)
+	return res, err
 }
 
-func unwrapWafResult(ret bindings.WAFReturnCode, result *bindings.WAFResult) (res Result, err error) {
+func unwrapWafResult(ret bindings.WAFReturnCode, result *bindings.WAFResult) (res Result, duration time.Duration, err error) {
 	if result.Timeout > 0 {
 		err = waferrors.ErrTimeout
 	} else {
@@ -242,28 +241,28 @@ func unwrapWafResult(ret bindings.WAFReturnCode, result *bindings.WAFResult) (re
 		res.Derivatives, err = decodeMap(&result.Derivatives)
 	}
 
-	res.TimeSpent = time.Duration(result.TotalRuntime) * time.Nanosecond
+	duration = time.Duration(result.TotalRuntime) * time.Nanosecond
 
 	if ret == bindings.WAFOK {
-		return res, err
+		return res, duration, err
 	}
 
 	if ret != bindings.WAFMatch {
-		return res, goRunError(ret)
+		return res, duration, goRunError(ret)
 	}
 
 	res.Events, err = decodeArray(&result.Events)
 	if err != nil {
-		return res, err
+		return res, duration, err
 	}
 	if size := result.Actions.NbEntries; size > 0 {
 		res.Actions, err = decodeMap(&result.Actions)
 		if err != nil {
-			return res, err
+			return res, duration, err
 		}
 	}
 
-	return res, err
+	return res, duration, err
 }
 
 // Close disposes of the underlying `ddwaf_context` and releases the associated
@@ -281,43 +280,13 @@ func (context *Context) Close() {
 	context.pinner.Unpin() // The pinned data is no longer needed, explicitly release
 }
 
-// Stats returns the cumulative time spent in various parts of the WAF, at
-// nanosecond resolution, as well as the timeout value used, and other
-// information about this [Context]'s usage.
-func (context *Context) Stats() Stats {
+// Truncations returns the truncations that occurred while encoding address data for WAF execution.
+// The key is the truncation reason: either because the object was too deep, the arrays where to large or the strings were too long.
+// The value is a slice of integers, each integer being the original size of the object that was truncated.
+// In case of the [ObjectTooDeep] reason, the original size can only be approximated because of recursive objects.
+func (context *Context) Truncations() map[TruncationReason][]int {
 	context.mutex.Lock()
 	defer context.mutex.Unlock()
 
-	truncations := make(map[TruncationReason][]int, len(context.truncations[DefaultScope]))
-	for reason, counts := range context.truncations[DefaultScope] {
-		truncations[reason] = make([]int, len(counts))
-		copy(truncations[reason], counts)
-	}
-
-	raspTruncations := make(map[TruncationReason][]int, len(context.truncations[RASPScope]))
-	for reason, counts := range context.truncations[RASPScope] {
-		raspTruncations[reason] = make([]int, len(counts))
-		copy(raspTruncations[reason], counts)
-	}
-
-	var (
-		timeoutDefault uint64
-		timeoutRASP    uint64
-	)
-
-	if atomic, ok := context.timeoutCount[DefaultScope]; ok {
-		timeoutDefault = atomic.Load()
-	}
-
-	if atomic, ok := context.timeoutCount[RASPScope]; ok {
-		timeoutRASP = atomic.Load()
-	}
-
-	return Stats{
-		Timers:           context.metrics.timers(),
-		TimeoutCount:     timeoutDefault,
-		TimeoutRASPCount: timeoutRASP,
-		Truncations:      truncations,
-		TruncationsRASP:  raspTruncations,
-	}
+	return maps.Clone(context.truncations)
 }