wip

Signed-off-by: Eliott Bouhana <eliott.bouhana@datadoghq.com>

internal/stacktrace/event.go

@@ -11,7 +11,6 @@ import (
 	"fmt"
 	"github.com/google/uuid"
 	"github.com/tinylib/msgp/msgp"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 )
 
 var _ msgp.Marshaler = (*Event)(nil)
@@ -27,6 +26,7 @@ const (
 	ExploitEvent EventCategory = "exploit"
 )
 
+// Event is the toplevel structure to contain a stacktrace and the additional information needed to correlate it with other data
 type Event struct {
 	// Category is a well-known type of the event, not optional
 	Category EventCategory `msg:"-"`
@@ -42,31 +42,23 @@ type Event struct {
 	Frames StackTrace `msg:"frames"`
 }
 
-func tagName(eventCategory EventCategory) string {
-	return fmt.Sprintf("_dd.stack.%s", eventCategory)
+// TagName returns the tag name for the event
+func (e *Event) TagName() string {
+	return fmt.Sprintf("_dd.stack.%s", e.Category)
 }
 
-func newEvent(eventCat EventCategory, message string) *Event {
+// NewEvent creates a new stacktrace event with the given category, type and message
+func NewEvent(eventCat EventCategory, type_, message string) *Event {
 	return &Event{
 		Category: eventCat,
+		Type:     type_,
 		Language: "go",
 		Message:  message,
 		Frames:   TakeWithSkip(defaultCallerSkip + 1),
 	}
 }
 
-func NewException(message string) *Event {
-	return newEvent(ExceptionEvent, message)
-}
-
-func NewVulnerability(message string) *Event {
-	return newEvent(VulnerabilityEvent, message)
-}
-
-func NewExploit(message string) *Event {
-	return newEvent(ExploitEvent, message)
-}
-
+// IDLink returns a UUID to link the stacktrace event with other data
 func (e *Event) IDLink() string {
 	if e.ID != "" {
 		newUUID, err := uuid.NewUUID()
@@ -80,6 +72,7 @@ func (e *Event) IDLink() string {
 	return e.ID
 }
 
-func (e *Event) AddToSpan(span ddtrace.Span) {
-	span.SetTag(tagName(e.Category), e)
+// AddToSpan uses (*Event).TagName to add the event to a span using span.SetTag
+func (e *Event) AddToSpan(span interface{ SetTag(key string, value any) }) {
+	span.SetTag(e.TagName(), *e)
 }

internal/stacktrace/event_test.go

@@ -7,14 +7,31 @@ package stacktrace
 
 import (
 	"github.com/stretchr/testify/require"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 	"testing"
 )
 
 func TestNewEvent(t *testing.T) {
-	event := NewException("message")
+	event := NewEvent(ExceptionEvent, "", "message")
 	require.Equal(t, ExceptionEvent, event.Category)
 	require.Equal(t, "go", event.Language)
 	require.Equal(t, "message", event.Message)
 	require.GreaterOrEqual(t, len(event.Frames), 3)
 	require.Equal(t, "gopkg.in/DataDog/dd-trace-go.v1/internal/stacktrace.TestNewEvent", event.Frames[len(event.Frames)-1].Function)
 }
+
+func TestEventToSpan(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	span := ddtracer.StartSpan("op")
+	event := NewEvent(ExceptionEvent, "", "message")
+	event.AddToSpan(span)
+	span.Finish()
+
+	spans := mt.FinishedSpans()
+	require.Len(t, spans, 1)
+	require.Equal(t, "op", spans[0].OperationName())
+	require.Equal(t, *event, spans[0].Tag("_dd.stack.exception"))
+}

internal/stacktrace/stacktrace.go

@@ -8,12 +8,37 @@
 package stacktrace
 
 import (
+	"os"
 	"runtime"
+	"strconv"
 	"strings"
 )
 
-const defaultMaxDepth = 32
+var enabled = true
+var topFrames = 8
+var defaultMaxDepth = 32
+
 const defaultCallerSkip = 4
+const stackTraceDepthEnvVar = "DD_APPSEC_MAX_STACK_TRACE_DEPTH"
+const stackTraceDisabledEnvVar = "DD_APPSEC_STACK_TRACE_ENABLE"
+
+func init() {
+	if env := os.Getenv(stackTraceDepthEnvVar); env != "" {
+		if depth, err := strconv.ParseUint(env, 10, 64); err == nil {
+			defaultMaxDepth = int(depth)
+
+			if defaultMaxDepth < topFrames {
+				topFrames = 0
+			}
+		}
+	}
+
+	if env := os.Getenv(stackTraceDisabledEnvVar); env != "" {
+		if e, err := strconv.ParseBool(env); err == nil {
+			enabled = e
+		}
+	}
+}
 
 // StackTrace is intended to be sent over the span tag `_dd.stack`, the first frame is the top of the stack
 type StackTrace []StackFrame
@@ -74,13 +99,26 @@ func Take() StackTrace {
 
 // TakeWithSkip creates a new stack trace from the current call stack, skipping the first `skip` frames
 func TakeWithSkip(skip int) StackTrace {
-	frames, depth := callers(skip, defaultMaxDepth)
-	stack := make([]StackFrame, depth)
+	callers, realDepth := callers(skip, defaultMaxDepth)
 
 	// There can be way more frames than callers, so we need to check again that we don't store more frames that the depth specified
+	frames := runtime.CallersFrames(callers)
+	framesArray := make([]runtime.Frame, 0, defaultMaxDepth)
+	ok := true
 	var frame runtime.Frame
+	for ok {
+		frame, ok = frames.Next()
+		framesArray = append(framesArray, frame)
+	}
+
+	if len(framesArray) > defaultMaxDepth {
+		framesArray = append(framesArray[:defaultMaxDepth-topFrames], framesArray[len(framesArray)-topFrames:]...)
+	}
+
+	stack := make([]StackFrame, len(framesArray))
+
+	// We revert the order of
 	for i := depth - 1; i >= 0; i-- {
-		frame, _ = frames.Next()
 		stack[i] = StackFrame{
 			Index:     uint32(i),
 			Text:      "",
@@ -96,17 +134,35 @@ func TakeWithSkip(skip int) StackTrace {
 	return stack
 }
 
-func callers(skip, maxDepth int) (*runtime.Frames, int) {
-	pcs := make([]uintptr, maxDepth)
+// callers returns a maximum of `maxDepth` frames from the call stack, skipping the first `skip` frames
+// of the whole stack is bigger than `maxDepth`, the 8 first and the last `maxDepth-8` frames are returned
+// the depth
+// Lastly, keep in mind that pcs[0] is the current function, not the top of the stack
+// callers also returns the real depth of the stack so that the indexs of the frames can be calculated
+func callers(skip, maxDepth int) ([]uintptr, int) {
+	pcs := make([]uintptr, maxDepth+1)
 	n := runtime.Callers(skip, pcs[:])
 
-	depth := maxDepth
-	// Find the real depth of the stack (if the stack is smaller than the max depth)
-	for ; depth > 0; depth-- {
-		if pcs[depth-1] != 0 {
-			break
+	// The stack is smaller or equal to the max depth, return the whole stack
+	if n <= maxDepth+1 {
+		// Find the real depth of the stack (if the stack is smaller than the max depth)
+		for ; maxDepth >= 0; maxDepth-- {
+			if pcs[maxDepth-1] != 0 {
+				break
+			}
 		}
+
+		return pcs[:maxDepth], maxDepth
+	}
+
+	// The stack is bigger than the max depth, proceed to find the top 8 frames and stitch them to the ones we have
+	topPcs := make([]uintptr, topFrames)
+	i := 0
+	var topN int
+	for ; topN > 0 && runtime.FuncForPC(topPcs[topN-1]).Name() != "goexit"; i += topFrames {
+		topN = runtime.Callers(skip+maxDepth+i, topPcs)
 	}
 
-	return runtime.CallersFrames(pcs[:n]), depth
+	// stitch the top frames to the ones we have
+	return append(pcs[:maxDepth-topFrames], topPcs[:topN]...), maxDepth + i
 }

internal/stacktrace/stacktrace_test.go

@@ -7,6 +7,7 @@ package stacktrace
 
 import (
 	"github.com/stretchr/testify/require"
+	"runtime"
 	"testing"
 )
 
@@ -48,3 +49,9 @@ func TestStackMethodReceiver(t *testing.T) {
 	require.Equal(t, "gopkg.in/DataDog/dd-trace-go.v1/internal/stacktrace.(*Test).Method", frame.Function)
 	require.Contains(t, frame.File, "internal/stacktrace/stacktrace_test.go")
 }
+
+func BenchmarkTakeStackTrace(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		runtime.KeepAlive(Take())
+	}
+}