From c28ef32905826a86f3322d726b9f2575f7bdfbc4 Mon Sep 17 00:00:00 2001
From: "Spindler, Justin" <Justin_Spindler@comcast.com>
Date: Mon, 26 Oct 2020 11:09:18 -0400
Subject: [PATCH] Add support for additional key/value pairs in X-MoneyTrace
 header

---
 trace.go      |  9 +++++----
 trace_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/trace.go b/trace.go
index d478d71..230f7b1 100644
--- a/trace.go
+++ b/trace.go
@@ -34,7 +34,7 @@ func decodeTraceContext(raw string) (tc *TraceContext, err error) {
 
 	pairs := strings.Split(raw, ";")
 
-	if len(pairs) != 3 {
+	if len(pairs) < 3 {
 		return nil, errPairsCount
 	}
 
@@ -66,12 +66,13 @@ func decodeTraceContext(raw string) (tc *TraceContext, err error) {
 				return nil, err
 			}
 			tc.PID, seen[k] = pv, true
-
-		default:
-			return nil, errBadTrace
 		}
 	}
 
+	if (!seen[tIDKey] || !seen[sIDKey] || !seen[pIDKey]) {
+		return nil, errBadTrace
+	}
+
 	return
 }
 
diff --git a/trace_test.go b/trace_test.go
index f552d75..abfdc1d 100644
--- a/trace_test.go
+++ b/trace_test.go
@@ -27,6 +27,16 @@ func TestDecodeTraceContext(t *testing.T) {
 			},
 			e: nil,
 		},
+		{
+			name: "extra",
+			i:    "trace-id=de305d54-75b4-431b-adb2-eb6b9e546013;parent-id=3285573610483682037;span-id=3285573610483682037;sampled=1",
+			o: &TraceContext{
+				PID: 3285573610483682037,
+				SID: 3285573610483682037,
+				TID: "de305d54-75b4-431b-adb2-eb6b9e546013",
+			},
+			e: nil,
+		},
 
 		{
 			name: "duplicateEntries",
@@ -47,6 +57,27 @@ func TestDecodeTraceContext(t *testing.T) {
 			o:    nil,
 			e:    errBadTrace,
 		},
+
+		{
+			name: "missingTraceId",
+			i:    "foo=bar;parent-id=3285573610483682037;span-id=3285573610483682037",
+			o:    nil,
+			e:    errBadTrace,
+		},
+
+		{
+			name: "missingParentId",
+			i:    "trace-id=de305d54-75b4-431b-adb2-eb6b9e546013;foo=bar;span-id=3285573610483682037",
+			o:    nil,
+			e:    errBadTrace,
+		},
+
+		{
+			name: "missingSpanId",
+			i:    "trace-id=de305d54-75b4-431b-adb2-eb6b9e546013;parent-id=3285573610483682037;foo=bar",
+			o:    nil,
+			e:    errBadTrace,
+		},
 	}
 
 	for _, test := range tests {