[Serverless] Type for extracting trace context from various event types [SVLS-3934].

LICENSE-3rdparty.csv

@@ -2075,17 +2075,31 @@ core,google.golang.org/protobuf/types/known/structpb,BSD-3-Clause,Copyright (c)
 core,google.golang.org/protobuf/types/known/timestamppb,BSD-3-Clause,Copyright (c) 2018 The Go Authors. All rights reserved
 core,google.golang.org/protobuf/types/known/wrapperspb,BSD-3-Clause,Copyright (c) 2018 The Go Authors. All rights reserved
 core,google.golang.org/protobuf/types/pluginpb,BSD-3-Clause,Copyright (c) 2018 The Go Authors. All rights reserved
+core,gopkg.in/DataDog/dd-trace-go.v1/datastreams/options,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/ddtrace,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation/grpcsec,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation/httpsec,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo/instrumentation/sharedsec,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/datastreams,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/azure,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/cachedfetch,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/ec2,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/ecs,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/gce,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/httputils,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/hostname/validate,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/log,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/namingschema,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
+core,gopkg.in/DataDog/dd-trace-go.v1/internal/normalizer,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/osinfo,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/remoteconfig,Apache-2.0,"Copyright 2016-Present Datadog, Inc."
 core,gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames,Apache-2.0,"Copyright 2016-Present Datadog, Inc."

go.mod

@@ -301,7 +301,7 @@ require (
 	github.com/acomagu/bufpipe v1.0.4 // indirect
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/agnivade/levenshtein v1.1.1 // indirect
-	github.com/andybalholm/brotli v1.0.5 // indirect
+	github.com/andybalholm/brotli v1.0.6 // indirect
 	github.com/aquasecurity/go-dep-parser v0.0.0-20230115135733-3be7cb085121 // indirect
 	github.com/aquasecurity/go-gem-version v0.0.0-20201115065557-8eed6fe000ce // indirect
 	github.com/aquasecurity/go-npm-version v0.0.0-20201110091526-0b796d180798 // indirect
@@ -413,7 +413,7 @@ require (
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
 	github.com/kjk/lzma v0.0.0-20161016003348-3fd93898850d // indirect
-	github.com/klauspost/compress v1.17.0 // indirect
+	github.com/klauspost/compress v1.17.1 // indirect
 	github.com/klauspost/pgzip v1.2.5 // indirect
 	github.com/knadh/koanf v1.5.0 // indirect
 	github.com/knqyf263/go-apk-version v0.0.0-20200609155635-041fdbb8563f // indirect
@@ -433,7 +433,7 @@ require (
 	github.com/masahiro331/go-vmdk-parser v0.0.0-20221225061455-612096e4bbbd // indirect
 	github.com/masahiro331/go-xfs-filesystem v0.0.0-20221225060805-c02764233454 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
-	github.com/mattn/go-isatty v0.0.19 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/mdlayher/socket v0.2.3 // indirect
@@ -706,3 +706,8 @@ exclude (
 	github.com/knadh/koanf/maps v0.1.1
 	github.com/knadh/koanf/providers/confmap v0.1.0
 )
+
+// Once a version of dd-trace-go with commit 1e7a3e0de599 is released, this can
+// be removed and the version of listed above dd-trace-go.v1 bumped to the
+// newer version.
+replace gopkg.in/DataDog/dd-trace-go.v1 => github.com/DataDog/dd-trace-go v0.0.0-20231030162158-1e7a3e0de599

pkg/serverless/trace/propagation/carriers.go

@@ -0,0 +1,167 @@
+// 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 2022-present Datadog, Inc.
+
+// Package propagation manages propagation of trace context headers.
+package propagation
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/aws/aws-lambda-go/events"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+const (
+	awsTraceHeader   = "AWSTraceHeader"
+	datadogSQSHeader = "_datadog"
+
+	rootPrefix    = "Root="
+	parentPrefix  = "Parent="
+	rootPadding   = len(rootPrefix + "1-00000000-00000000")
+	parentPadding = len(parentPrefix)
+)
+
+var rootRegex = regexp.MustCompile("Root=1-[0-9a-fA-F]{8}-00000000[0-9a-fA-F]{16}")
+
+// extractTraceContextfromAWSTraceHeader extracts trace context from the
+// AWSTraceHeader directly. Unlike the other carriers in this file, it should
+// not be passed to the tracer.Propagator, instead extracting context directly.
+func extractTraceContextfromAWSTraceHeader(value string) (*TraceContext, error) {
+	if !rootRegex.MatchString(value) {
+		return nil, errors.New("AWSTraceHeader does not match expected regex")
+	}
+	var (
+		startPart int
+		traceID   string
+		parentID  string
+		err       error
+	)
+	length := len(value)
+	for startPart < length {
+		endPart := strings.IndexRune(value[startPart:], ';') + startPart
+		if endPart < startPart {
+			endPart = length
+		}
+		part := value[startPart:endPart]
+		if strings.HasPrefix(part, rootPrefix) {
+			if traceID == "" {
+				traceID = part[rootPadding:]
+			}
+		} else if strings.HasPrefix(part, parentPrefix) {
+			if parentID == "" {
+				parentID = part[parentPadding:]
+			}
+		}
+		if traceID != "" && parentID != "" {
+			break
+		}
+		startPart = endPart + 1
+	}
+	tc := new(TraceContext)
+	tc.TraceID, err = strconv.ParseUint(traceID, 16, 64)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to parse trace ID from AWSTraceHeader: %w", err)
+	}
+	tc.ParentID, err = strconv.ParseUint(parentID, 16, 64)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to parse parent ID from AWSTraceHeader: %w", err)
+	}
+	if tc.TraceID == 0 || tc.ParentID == 0 {
+		return nil, errors.New("AWSTraceHeader does not contain trace ID and parent ID")
+	}
+	return tc, nil
+}
+
+// sqsMessageCarrier returns the tracer.TextMapReader used to extract trace
+// context from the events.SQSMessage type.
+func sqsMessageCarrier(event events.SQSMessage) (tracer.TextMapReader, error) {
+	if attr, ok := event.MessageAttributes[datadogSQSHeader]; ok {
+		return sqsMessageAttrCarrier(attr)
+	}
+	return snsSqsMessageCarrier(event)
+}
+
+// sqsMessageAttrCarrier returns the tracer.TextMapReader used to extract trace
+// context from the events.SQSMessageAttribute field on an events.SQSMessage
+// type.
+func sqsMessageAttrCarrier(attr events.SQSMessageAttribute) (tracer.TextMapReader, error) {
+	var bytes []byte
+	switch attr.DataType {
+	case "String":
+		if attr.StringValue == nil {
+			return nil, errors.New("String value not found in _datadog payload")
+		}
+		bytes = []byte(*attr.StringValue)
+	case "Binary":
+		// SNS => SQS => Lambda with SQS's subscription to SNS has enabled RAW
+		// MESSAGE DELIVERY option
+		bytes = attr.BinaryValue // No need to decode base64 because already decoded
+	default:
+		return nil, errors.New("Unsupported DataType in _datadog payload")
+	}
+
+	var carrier tracer.TextMapCarrier
+	if err := json.Unmarshal(bytes, &carrier); err != nil {
+		return nil, fmt.Errorf("Error unmarshaling payload value: %w", err)
+	}
+	return carrier, nil
+}
+
+// snsSqsMessageCarrier returns the tracer.TextMapReader used to extract trace
+// context from the body of an events.SQSMessage type.
+func snsSqsMessageCarrier(event events.SQSMessage) (tracer.TextMapReader, error) {
+	var body struct {
+		MessageAttributes map[string]struct {
+			Type  string
+			Value string
+		}
+	}
+	err := json.Unmarshal([]byte(event.Body), &body)
+	if err != nil {
+		return nil, fmt.Errorf("Error unmarshaling message body: %w", err)
+	}
+	msgAttrs, ok := body.MessageAttributes[datadogSQSHeader]
+	if !ok {
+		return nil, errors.New("No Datadog trace context found")
+	}
+	if msgAttrs.Type != "Binary" {
+		return nil, errors.New("Unsupported DataType in _datadog payload")
+	}
+	attr, err := base64.StdEncoding.DecodeString(string(msgAttrs.Value))
+	if err != nil {
+		return nil, fmt.Errorf("Error decoding binary: %w", err)
+	}
+	var carrier tracer.TextMapCarrier
+	if err = json.Unmarshal(attr, &carrier); err != nil {
+		return nil, fmt.Errorf("Error unmarshaling the decoded binary: %w", err)
+	}
+	return carrier, nil
+}
+
+type invocationPayload struct {
+	Headers tracer.TextMapCarrier `json:"headers"`
+}
+
+// rawPayloadCarrier returns the tracer.TextMapReader used to extract trace
+// context from the raw json event payload.
+func rawPayloadCarrier(rawPayload []byte) (tracer.TextMapReader, error) {
+	var payload invocationPayload
+	if err := json.Unmarshal(rawPayload, &payload); err != nil {
+		return nil, errors.New("Could not unmarshal the invocation event payload")
+	}
+	return payload.Headers, nil
+}
+
+// headersCarrier returns the tracer.TextMapReader used to extract trace
+// context from a Headers field of form map[string]string.
+func headersCarrier(hdrs map[string]string) (tracer.TextMapReader, error) {
+	return tracer.TextMapCarrier(hdrs), nil
+}