[serverless] Add S3 span pointers

contrib/aws/aws-sdk-go-v2/aws/aws.go

@@ -8,6 +8,7 @@ package aws
 import (
 	"context"
 	"fmt"
+	spanpointers "gopkg.in/DataDog/dd-trace-go.v1/contrib/aws/internal/span_pointers"
 	"math"
 	"strings"
 	"time"
@@ -124,10 +125,6 @@ func (mw *traceMiddleware) startTraceMiddleware(stack *middleware.Stack) error {
 
 		// Handle initialize and continue through the middleware chain.
 		out, metadata, err = next.HandleInitialize(spanctx, in)
-		if err != nil && (mw.cfg.errCheck == nil || mw.cfg.errCheck(err)) {
-			span.SetTag(ext.Error, err)
-		}
-		span.Finish()
 
 		return out, metadata, err
 	}), middleware.After)
@@ -357,6 +354,15 @@ func (mw *traceMiddleware) deserializeTraceMiddleware(stack *middleware.Stack) e
 			span.SetTag(tags.AWSRequestID, requestID)
 		}
 
+		// Create span pointers
+		serviceID := awsmiddleware.GetServiceID(ctx)
+		spanpointers.AddSpanPointers(serviceID, in, out, span)
+
+		if err != nil && (mw.cfg.errCheck == nil || mw.cfg.errCheck(err)) {
+			span.SetTag(ext.Error, err)
+		}
+		span.Finish()
+
 		return out, metadata, err
 	}), middleware.Before)
 }

contrib/aws/internal/span_pointers/span_pointers.go

@@ -0,0 +1,97 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package spanpointers
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"strings"
+)
+
+const (
+	// SpanPointerHashLengthBytes 16 bytes = 32 chars.
+	// See https://github.com/DataDog/dd-span-pointer-rules/blob/main/README.md#general-hashing-rules
+	SpanPointerHashLengthBytes = 16
+	PointerDownDirection       = "d"
+	LinkKind                   = "span-pointer"
+	S3PointerKind              = "aws.s3.object"
+)
+
+func AddSpanPointers(serviceID string, in middleware.DeserializeInput, out middleware.DeserializeOutput, span tracer.Span) {
+	switch serviceID {
+	case "S3":
+		handleS3Operation(in, out, span)
+	}
+}
+
+func handleS3Operation(in middleware.DeserializeInput, out middleware.DeserializeOutput, span tracer.Span) {
+	spanWithLinks, ok := span.(tracer.SpanWithLinks)
+	if !ok {
+		return
+	}
+
+	req, ok := in.Request.(*smithyhttp.Request)
+	if !ok {
+		return
+	}
+	res, ok := out.RawResponse.(*smithyhttp.Response)
+	if !ok {
+		return
+	}
+
+	// URL format: https://BUCKETNAME.s3.REGION.amazonaws.com/KEYNAME?x-id=OPERATIONNAME
+	key := strings.TrimPrefix(req.URL.Path, "/")
+	bucket := strings.Split(req.URL.Host, ".")[0]
+	// the AWS SDK sometimes wraps the eTag in quotes
+	etag := strings.Trim(res.Header.Get("ETag"), "\"")
+	if key == "" || bucket == "" || etag == "" {
+		log.Debug("Unable to create S3 span pointer because required fields could not be found.")
+		return
+	}
+
+	// Hash calculation rules: https://github.com/DataDog/dd-span-pointer-rules/blob/main/AWS/S3/Object/README.md
+	components := []string{bucket, key, etag}
+	hash := generatePointerHash(components)
+
+	link := ddtrace.SpanLink{
+		// We leave trace_id, span_id, trade_id_high, tracestate, and flags as 0 or empty.
+		// The Datadog frontend will use `ptr.hash` to find the linked span.
+		TraceID:     0,
+		SpanID:      0,
+		TraceIDHigh: 0,
+		Flags:       0,
+		Tracestate:  "",
+		Attributes: map[string]string{
+			"ptr.kind":  S3PointerKind,
+			"ptr.dir":   PointerDownDirection,
+			"ptr.hash":  hash,
+			"link.kind": LinkKind,
+		},
+	}
+
+	spanWithLinks.AddSpanLink(link)
+}
+
+// generatePointerHash generates a unique hash from an array of strings by joining them with | before hashing.
+// Used to uniquely identify AWS requests for span pointers.
+// Returns a 32-character hash uniquely identifying the components.
+func generatePointerHash(components []string) string {
+	h := sha256.New()
+	for i, component := range components {
+		if i > 0 {
+			h.Write([]byte("|"))
+		}
+		h.Write([]byte(component))
+	}
+
+	fullHash := h.Sum(nil)
+	return hex.EncodeToString(fullHash[:SpanPointerHashLengthBytes])
+}

contrib/aws/internal/span_pointers/span_pointers_test.go

@@ -0,0 +1,190 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package spanpointers
+
+import (
+	"context"
+	"encoding/json"
+	"github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"net/http"
+	"net/url"
+	"testing"
+)
+
+func TestGeneratePointerHash(t *testing.T) {
+	tests := []struct {
+		name         string
+		components   []string
+		expectedHash string
+	}{
+		{
+			name: "basic values",
+			components: []string{
+				"some-bucket",
+				"some-key.data",
+				"ab12ef34",
+			},
+			expectedHash: "e721375466d4116ab551213fdea08413",
+		},
+		{
+			name: "non-ascii key",
+			components: []string{
+				"some-bucket",
+				"some-key.你好",
+				"ab12ef34",
+			},
+			expectedHash: "d1333a04b9928ab462b5c6cadfa401f4",
+		},
+		{
+			name: "multipart-upload",
+			components: []string{
+				"some-bucket",
+				"some-key.data",
+				"ab12ef34-5",
+			},
+			expectedHash: "2b90dffc37ebc7bc610152c3dc72af9f",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := generatePointerHash(tt.components)
+			if got != tt.expectedHash {
+				t.Errorf("GeneratePointerHash() = %v, want %v", got, tt.expectedHash)
+			}
+		})
+	}
+}
+
+func TestHandleS3Operation(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	tests := []struct {
+		name          string
+		bucket        string
+		key           string
+		etag          string
+		expectedHash  string
+		expectSuccess bool
+	}{
+		{
+			name:          "basic operation",
+			bucket:        "some-bucket",
+			key:           "some-key.data",
+			etag:          "ab12ef34",
+			expectedHash:  "e721375466d4116ab551213fdea08413",
+			expectSuccess: true,
+		},
+		{
+			name:          "quoted etag",
+			bucket:        "some-bucket",
+			key:           "some-key.data",
+			etag:          "\"ab12ef34\"",
+			expectedHash:  "e721375466d4116ab551213fdea08413",
+			expectSuccess: true,
+		},
+		{
+			name:          "non-ascii key",
+			bucket:        "some-bucket",
+			key:           "some-key.你好",
+			etag:          "ab12ef34",
+			expectedHash:  "d1333a04b9928ab462b5c6cadfa401f4",
+			expectSuccess: true,
+		},
+		{
+			name:          "empty bucket",
+			bucket:        "",
+			key:           "some_key",
+			etag:          "some_etag",
+			expectSuccess: false,
+		},
+		{
+			name:          "empty key",
+			bucket:        "some_bucket",
+			key:           "",
+			etag:          "some_etag",
+			expectSuccess: false,
+		},
+		{
+			name:          "empty etag",
+			bucket:        "some_bucket",
+			key:           "some_key",
+			etag:          "",
+			expectSuccess: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			span, ctx := tracer.StartSpanFromContext(ctx, "test.s3.operation")
+
+			// Create request
+			reqURL, _ := url.Parse("https://" + tt.bucket + ".s3.region.amazonaws.com/" + tt.key)
+			req := &smithyhttp.Request{
+				Request: &http.Request{
+					URL: reqURL,
+				},
+			}
+
+			// Create response
+			header := http.Header{}
+			header.Set("ETag", tt.etag)
+			res := &smithyhttp.Response{
+				Response: &http.Response{
+					Header: header,
+				},
+			}
+
+			// Create input/output
+			in := middleware.DeserializeInput{
+				Request: req,
+			}
+			out := middleware.DeserializeOutput{
+				RawResponse: res,
+			}
+
+			AddSpanPointers("S3", in, out, span)
+			span.Finish()
+			spans := mt.FinishedSpans()
+			if tt.expectSuccess {
+				require.Len(t, spans, 1)
+				meta := spans[0].Tags()
+
+				spanLinks, exists := meta["_dd.span_links"]
+				assert.True(t, exists, "Expected span links to be set")
+				assert.NotEmpty(t, spanLinks, "Expected span links to not be empty")
+
+				spanLinksStr, ok := spanLinks.(string)
+				assert.True(t, ok, "Expected span links to be a string")
+
+				var links []ddtrace.SpanLink
+				err := json.Unmarshal([]byte(spanLinksStr), &links)
+				require.NoError(t, err)
+				require.Len(t, links, 1)
+
+				attributes := links[0].Attributes
+				assert.Equal(t, S3PointerKind, attributes["ptr.kind"])
+				assert.Equal(t, PointerDownDirection, attributes["ptr.dir"])
+				assert.Equal(t, LinkKind, attributes["link.kind"])
+				assert.Equal(t, tt.expectedHash, attributes["ptr.hash"])
+			} else {
+				require.Len(t, spans, 1)
+				tags := spans[0].Tags()
+				_, exists := tags["_dd.span_links"]
+				assert.False(t, exists, "Expected no span links to be set")
+			}
+			mt.Reset()
+		})
+	}
+}

ddtrace/mocktracer/mockspan.go

@@ -6,6 +6,7 @@
 package mocktracer // import "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
 
 import (
+	"encoding/json"
 	"fmt"
 	"sync"
 	"time"
@@ -206,6 +207,7 @@ func (s *mockspan) Finish(opts ...ddtrace.FinishOption) {
 	if cfg.NoDebugStack {
 		s.SetTag(ext.ErrorStack, "<debug stack disabled>")
 	}
+	s.serializeSpanLinksInMeta()
 	s.Lock()
 	defer s.Unlock()
 	if s.finished {
@@ -284,3 +286,18 @@ func (s *mockspan) Root() tracer.Span {
 	root, _ := current.(*mockspan)
 	return root
 }
+
+func (s *mockspan) AddSpanLink(link ddtrace.SpanLink) {
+	s.links = append(s.links, link)
+}
+
+func (s *mockspan) serializeSpanLinksInMeta() {
+	if len(s.links) == 0 {
+		return
+	}
+	spanLinkBytes, err := json.Marshal(s.links)
+	if err != nil {
+		return
+	}
+	s.tags["_dd.span_links"] = string(spanLinkBytes)
+}