opentelemetry-go/propagation/baggage_test.go

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package propagation_test

import (
	"fmt"
	"net/http"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/assert"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/baggage"
	"go.opentelemetry.io/otel/propagation"
)

const maxBytesPerBaggageString = 8192

type property struct {
	Key, Value string
}

type member struct {
	Key, Value string

	Properties []property
}

func (m member) Member(t *testing.T) baggage.Member {
	props := make([]baggage.Property, 0, len(m.Properties))
	for _, p := range m.Properties {
		p, err := baggage.NewKeyValuePropertyRaw(p.Key, p.Value)
		if err != nil {
			t.Fatal(err)
		}
		props = append(props, p)
	}
	bMember, err := baggage.NewMemberRaw(m.Key, m.Value, props...)
	if err != nil {
		t.Fatal(err)
	}
	return bMember
}

type members []member

func (m members) Baggage(t *testing.T) baggage.Baggage {
	bMembers := make([]baggage.Member, 0, len(m))
	for _, mem := range m {
		bMembers = append(bMembers, mem.Member(t))
	}
	bag, err := baggage.New(bMembers...)
	if err != nil {
		t.Fatal(err)
	}
	return bag
}

func TestExtractValidBaggage(t *testing.T) {
	prop := propagation.TextMapPropagator(propagation.Baggage{})
	tests := []struct {
		name   string
		header string
		want   members
	}{
		{
			name:   "valid w3cHeader",
			header: "key1=val1,key2=val2",
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "valid w3cHeader with spaces",
			header: "key1 =   val1,  key2 =val2   ",
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "valid w3cHeader with properties",
			header: "key1=val1,key2=val2;prop=1",
			want: members{
				{Key: "key1", Value: "val1"},
				{
					Key:   "key2",
					Value: "val2",
					Properties: []property{
						{Key: "prop", Value: "1"},
					},
				},
			},
		},
		{
			name:   "valid header with an invalid header",
			header: "key1=val1,key2=val2,a,val3",
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "valid header with no value",
			header: "key1=,key2=val2",
			want: members{
				{Key: "key1", Value: ""},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "valid header with url encoded string",
			header: "key1=val%252",
			want: members{
				{Key: "key1", Value: "val%2"},
			},
		},
		{
			name:   "empty header",
			header: "",
			want:   members{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mapCarr := propagation.MapCarrier{}
			mapCarr["baggage"] = tt.header
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			req.Header.Set("baggage", tt.header)

			// test with http header carrier (which implements ValuesGetter)
			ctx := prop.Extract(t.Context(), propagation.HeaderCarrier(req.Header))
			expected := tt.want.Baggage(t)
			assert.Equal(t, expected, baggage.FromContext(ctx), "should extract baggage for HeaderCarrier")

			// test with map carrier (which does not implement ValuesGetter)
			ctx = prop.Extract(t.Context(), mapCarr)
			expected = tt.want.Baggage(t)
			assert.Equal(t, expected, baggage.FromContext(ctx), "should extract baggage for MapCarrier")
		})
	}
}

// generateBaggageHeader creates a baggage header string with n members.
func generateBaggageHeader(n int, prefix string) string {
	parts := make([]string, n)
	for i := range parts {
		parts[i] = fmt.Sprintf("%s%d=v%d", prefix, i, i)
	}
	return strings.Join(parts, ",")
}

// generateMembers creates n members with keys like "prefix0", "prefix1", etc.
func generateMembers(n int, prefix string) members {
	m := make(members, n)
	for i := range m {
		m[i] = member{Key: fmt.Sprintf("%s%d", prefix, i), Value: fmt.Sprintf("v%d", i)}
	}
	return m
}

// generateFixedWidthBaggageHeader creates a baggage header where every member
// has the same byte length, ensuring deterministic String() size regardless of
// which members baggage.New() keeps after map-order truncation.
// Each member is "prefixNN=vNN" with zero-padded indices.
func generateFixedWidthBaggageHeader(n int, prefix string) string {
	parts := make([]string, n)
	for i := range parts {
		parts[i] = fmt.Sprintf("%s%02d=v%02d", prefix, i, i)
	}
	return strings.Join(parts, ",")
}

func TestExtractValidMultipleBaggageHeaders(t *testing.T) {
	// W3C Baggage spec limits: https://www.w3.org/TR/baggage/#limits
	const maxMembers = 64

	prop := propagation.TextMapPropagator(propagation.Baggage{})
	tests := []struct {
		name      string
		headers   []string
		want      members
		wantCount int // Used when want is nil and we only care about count.
		wantBytes int // Used to check exact baggage size when want is nil.
	}{
		{
			name:    "non conflicting headers",
			headers: []string{"key1=val1", "key2=val2"},
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:    "conflicting keys, uses last val",
			headers: []string{"key1=val1", "key1=val2"},
			want: members{
				{Key: "key1", Value: "val2"},
			},
		},
		{
			name:    "single empty",
			headers: []string{"", "key1=val1"},
			want: members{
				{Key: "key1", Value: "val1"},
			},
		},
		{
			name:    "all empty",
			headers: []string{"", ""},
			want:    members{},
		},
		{
			name:    "none",
			headers: []string{},
			want:    members{},
		},
		{
			name:    "single header with one invalid skips invalid",
			headers: []string{"key1=val1,invalid-no-equals,key2=val2"},
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name: "multiple headers with one invalid skips invalid and continues",
			headers: []string{
				"key1=val1",
				"invalid-no-equals",
				"key2=val2",
			},
			want: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:    "single header at max members limit",
			headers: []string{generateBaggageHeader(maxMembers, "k")},
			want:    generateMembers(maxMembers, "k"),
		},
		{
			name:    "single header exceeds max members limit keeps 64",
			headers: []string{generateBaggageHeader(maxMembers+1, "k")},
			want:    generateMembers(maxMembers, "k"),
		},
		{
			name: "multiple headers exceeds total max members limit keeps 64",
			headers: []string{
				generateFixedWidthBaggageHeader(maxMembers/2, "a"),
				generateFixedWidthBaggageHeader(maxMembers/2, "b"),
				generateFixedWidthBaggageHeader(1, "c"),
			},
			want:      nil, // Non-deterministic truncation by baggage.New()
			wantCount: maxMembers,
			wantBytes: maxMembers*7 + maxMembers - 1, // 64 uniform "xNN=vNN" members + 63 commas
		},
		{
			name:    "single header at max bytes limit",
			headers: []string{"k=" + strings.Repeat("v", maxBytesPerBaggageString-2)},
			want: members{
				{Key: "k", Value: strings.Repeat("v", maxBytesPerBaggageString-2)},
			},
		},
		{
			name:    "single header exceeds max bytes limit will be dropped entirely",
			headers: []string{"k=" + strings.Repeat("v", maxBytesPerBaggageString-1)},
			want:    members{},
		},
		{
			name: "multiple headers exceed total max bytes drops everything",
			headers: []string{
				"k=" + strings.Repeat("v", maxBytesPerBaggageString-2),
				"y=" + strings.Repeat("v", maxBytesPerBaggageString-2),
			},
			want: members{},
		},
		{
			name: "multiple headers within total max bytes",
			headers: []string{
				"k=" + strings.Repeat("v", maxBytesPerBaggageString/2-2),
				// The comma as the separator of member would take 1 byte.
				"y=" + strings.Repeat("v", maxBytesPerBaggageString/2-2-1),
			},
			want: members{
				{Key: "k", Value: strings.Repeat("v", maxBytesPerBaggageString/2-2)},
				{Key: "y", Value: strings.Repeat("v", maxBytesPerBaggageString/2-2-1)},
			},
		},
		{
			name: "empty headers between non-empty do count comma separators",
			headers: []string{
				"k=" + strings.Repeat("v", maxBytesPerBaggageString/2-2),
				"",
				"",
				// The comma as the separator of member would take 1 byte.
				"y=" + strings.Repeat("v", maxBytesPerBaggageString/2-2-1),
			},
			want: members{},
		},
		{
			name: "empty headers before non-empty do count comma separators",
			headers: []string{
				// The comma as the separator of member would take 1 byte.
				"",
				"k=" + strings.Repeat("v", maxBytesPerBaggageString/2-2),
				// The comma as the separator of member would take 1 byte.
				"y=" + strings.Repeat("v", maxBytesPerBaggageString/2-2-1),
			},
			want: members{},
		},
		{
			name: "multiple headers exceed total max bytes due to comma separator",
			headers: []string{
				// Two equal halves: each is exactly maxBytesPerBaggageString/2 bytes.
				// Without the comma separator: 4096 + 4096 = 8192 (at limit).
				// With the comma separator: 4096 + 1 + 4096 = 8193 (over limit).
				"k=" + strings.Repeat("v", maxBytesPerBaggageString/2-2),
				"y=" + strings.Repeat("v", maxBytesPerBaggageString/2-2),
			},
			want: members{},
		},
		{
			name: "many headers total max members and max bytes",
			headers: func() []string {
				// 100 headers with 10 members each = 1000 total members.
				h := make([]string, 100)
				for i := range h {
					h[i] = generateFixedWidthBaggageHeader(10, fmt.Sprintf("h%02d_k", i))
				}
				return h
			}(),
			want: members{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			req.Header["Baggage"] = tt.headers

			ctx := t.Context()
			ctx = prop.Extract(ctx, propagation.HeaderCarrier(req.Header))
			got := baggage.FromContext(ctx)

			// If want is specified, check exact match
			if tt.want != nil {
				expected := tt.want.Baggage(t)
				assert.Equal(t, expected, got)
			} else {
				assert.Equal(t, tt.wantCount, got.Len(), "expected member count")
				assert.Len(t, got.String(), tt.wantBytes, "expected baggage size")
			}
		})
	}
}

func TestExtractInvalidDistributedContextFromHTTPReq(t *testing.T) {
	prop := propagation.TextMapPropagator(propagation.Baggage{})
	tests := []struct {
		name   string
		header string
		has    members
	}{
		{
			name:   "no key values",
			header: "header1",
		},
		{
			name:   "invalid header with existing context",
			header: "header2",
			has: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "empty header value",
			header: "",
			has: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name:   "with properties",
			header: "key1=val1,key2=val2;prop=1",
			has: members{
				{Key: "key1", Value: "val1"},
				{
					Key:   "key2",
					Value: "val2",
					Properties: []property{
						{Key: "prop", Value: "1"},
					},
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			req.Header.Set("baggage", tt.header)

			expected := tt.has.Baggage(t)
			ctx := baggage.ContextWithBaggage(t.Context(), expected)
			ctx = prop.Extract(ctx, propagation.HeaderCarrier(req.Header))
			assert.Equal(t, expected, baggage.FromContext(ctx))
		})
	}
}

func TestInjectBaggageToHTTPReq(t *testing.T) {
	propagator := propagation.Baggage{}
	tests := []struct {
		name         string
		mems         members
		wantInHeader []string
	}{
		{
			name: "two simple values",
			mems: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
			wantInHeader: []string{"key1=val1", "key2=val2"},
		},
		{
			name: "values with escaped chars",
			mems: members{
				{Key: "key2", Value: "val3,4"},
			},
			wantInHeader: []string{"key2=val3%2C4"},
		},
		{
			name: "with properties",
			mems: members{
				{Key: "key1", Value: "val1"},
				{
					Key:   "key2",
					Value: "val2",
					Properties: []property{
						{Key: "prop", Value: "1"},
					},
				},
			},
			wantInHeader: []string{
				"key1=val1",
				"key2=val2;prop=1",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			ctx := baggage.ContextWithBaggage(t.Context(), tt.mems.Baggage(t))
			propagator.Inject(ctx, propagation.HeaderCarrier(req.Header))

			got := strings.Split(req.Header.Get("baggage"), ",")
			assert.ElementsMatch(t, tt.wantInHeader, got)
		})
	}
}

func TestBaggageInjectExtractRoundtrip(t *testing.T) {
	propagator := propagation.Baggage{}
	tests := []struct {
		name string
		mems members
	}{
		{
			name: "two simple values",
			mems: members{
				{Key: "key1", Value: "val1"},
				{Key: "key2", Value: "val2"},
			},
		},
		{
			name: "values with escaped chars",
			mems: members{
				{Key: "key1", Value: "val3=4"},
				{Key: "key2", Value: "mess,me%up"},
			},
		},
		{
			name: "with properties",
			mems: members{
				{Key: "key1", Value: "val1"},
				{
					Key:   "key2",
					Value: "val2",
					Properties: []property{
						{Key: "prop", Value: "1"},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			b := tt.mems.Baggage(t)
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			ctx := baggage.ContextWithBaggage(t.Context(), b)
			propagator.Inject(ctx, propagation.HeaderCarrier(req.Header))

			ctx = propagator.Extract(t.Context(), propagation.HeaderCarrier(req.Header))
			got := baggage.FromContext(ctx)

			assert.Equal(t, b, got)
		})
	}
}

func TestBaggagePropagatorGetAllKeys(t *testing.T) {
	var propagator propagation.Baggage
	want := []string{"baggage"}
	got := propagator.Fields()
	if diff := cmp.Diff(got, want); diff != "" {
		t.Errorf("GetAllKeys: -got +want %s", diff)
	}
}

func TestExtractOversizedSingleBaggageHeader(t *testing.T) {
	prop := propagation.Baggage{}
	req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
	// Set a single baggage header exceeding 8192 bytes.
	req.Header.Set("baggage", "key="+strings.Repeat("v", maxBytesPerBaggageString))

	ctx := prop.Extract(t.Context(), propagation.HeaderCarrier(req.Header))
	got := baggage.FromContext(ctx)
	assert.Equal(t, 0, got.Len(), "oversized header should result in empty baggage")
}

type errHandler struct {
	err error
}

func (e *errHandler) Handle(err error) { e.err = err }

func TestExtractManyBaggageHeader(t *testing.T) {
	tests := []struct {
		name       string
		headers    func() []string
		want       func() members
		wantErrStr []string
	}{
		{
			name: "aggregate byte budget exceeded drops everything",
			headers: func() []string {
				// 100 headers, each ~195 bytes. Total: ~19.5KB, well over 8192.
				h := make([]string, 100)
				for i := range 100 {
					h[i] = fmt.Sprintf("k%d=%s", i, strings.Repeat("v", 190))
				}
				return h
			},
			want: func() members {
				return members{}
			},
			wantErrStr: []string{"aggregate header size 8374 exceeds 8192 byte limit"},
		},
		{
			name: "too many invalid headers triggers error cap",
			headers: func() []string {
				// 10 invalid headers (no '=' delimiter) followed by 1 valid.
				// Each invalid header produces 1 parse error from baggage.Parse.
				// maxParseErrors=5, so 5 errors are kept and 5 are summarized.
				h := make([]string, 11)
				for i := range 10 {
					h[i] = fmt.Sprintf("bad%d", i)
				}
				h[10] = "good=val"
				return h
			},
			want: func() members {
				return members{{Key: "good", Value: "val"}}
			},
			wantErrStr: []string{"invalid baggage list-member", "and 5 more error(s)"},
		},
		{
			name: "single header with many invalid members preserves inner error",
			headers: func() []string {
				// One header with 10 invalid members. baggage.Parse caps its own
				// joined error at maxParseErrors=5 and appends its own summary.
				// extractMultiBaggage must count this as a single failing header
				// so the inner error is still attached to the outer error chain.
				invalid := make([]string, 10)
				for i := range invalid {
					invalid[i] = fmt.Sprintf("bad%d", i)
				}
				return []string{strings.Join(invalid, ",") + ",good=val"}
			},
			want: func() members {
				return members{{Key: "good", Value: "val"}}
			},
			wantErrStr: []string{"invalid baggage list-member", "and 5 more invalid member(s)"},
		},
		{
			name: "member cap does not stop aggregate byte accounting",
			headers: func() []string {
				return []string{
					generateFixedWidthBaggageHeader(64, "a"),
					"oversized=" + strings.Repeat("v", maxBytesPerBaggageString),
				}
			},
			want: func() members {
				return members{}
			},
			wantErrStr: []string{"exceeds 8192 byte limit"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			originalErrorHandler := otel.GetErrorHandler()
			eh := &errHandler{}
			otel.SetErrorHandler(eh)
			t.Cleanup(func() { otel.SetErrorHandler(originalErrorHandler) })

			prop := propagation.Baggage{}
			req, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://example.com", http.NoBody)
			for _, h := range tt.headers() {
				req.Header.Add("baggage", h)
			}

			ctx := prop.Extract(t.Context(), propagation.HeaderCarrier(req.Header))
			got := baggage.FromContext(ctx)

			assert.Equal(t, tt.want().Baggage(t), got)
			assert.Error(t, eh.err)
			for _, s := range tt.wantErrStr {
				assert.Contains(t, eh.err.Error(), s)
			}
		})
	}
}