Tracestate propagate (#901)

* Propagate tracestate

* Fixed review comments.

* Fixed leading/trailing OWS removal.

* Refactor to create separate package to inject/extract tracestate for http.

* Fix review comments.

* removed unnecessary check and import comment.

* Revert "removed unnecessary check and import comment."

This reverts commit 5349341.

* Revert "Fix review comments."

This reverts commit 40c2858.

* Revert "Refactor to create separate package to inject/extract tracestate for http."

This reverts commit 5574ce8.

* shorten variable name.

plugin/ochttp/propagation/tracecontext/propagation.go

@@ -24,22 +24,29 @@ import (
 
 	"go.opencensus.io/trace"
 	"go.opencensus.io/trace/propagation"
+	"go.opencensus.io/trace/tracestate"
+	"regexp"
 )
 
 const (
-	supportedVersion = 0
-	maxVersion       = 254
-	header           = "traceparent"
+	supportedVersion  = 0
+	maxVersion        = 254
+	maxTracestateLen  = 512
+	traceparentHeader = "traceparent"
+	tracestateHeader  = "tracestate"
+	trimOWSRegexFmt   = `^[\x09\x20]*(.*[^\x20\x09])[\x09\x20]*$`
 )
 
+var trimOWSRegExp = regexp.MustCompile(trimOWSRegexFmt)
+
 var _ propagation.HTTPFormat = (*HTTPFormat)(nil)
 
 // HTTPFormat implements the TraceContext trace propagation format.
 type HTTPFormat struct{}
 
 // SpanContextFromRequest extracts a span context from incoming requests.
 func (f *HTTPFormat) SpanContextFromRequest(req *http.Request) (sc trace.SpanContext, ok bool) {
-	h := req.Header.Get(header)
+	h := req.Header.Get(traceparentHeader)
 	if h == "" {
 		return trace.SpanContext{}, false
 	}
@@ -87,15 +94,69 @@ func (f *HTTPFormat) SpanContextFromRequest(req *http.Request) (sc trace.SpanCon
 		return trace.SpanContext{}, false
 	}
 
+	sc.Tracestate = tracestateFromRequest(req)
 	return sc, true
 }
 
-// SpanContextToRequest modifies the given request to include a header.
+// TODO(rghetia): return an empty Tracestate when parsing tracestate header encounters an error.
+// Revisit to return additional boolean value to indicate parsing error when following issues
+// are resolved.
+// https://github.com/w3c/distributed-tracing/issues/172
+// https://github.com/w3c/distributed-tracing/issues/175
+func tracestateFromRequest(req *http.Request) *tracestate.Tracestate {
+	h := req.Header.Get(tracestateHeader)
+	if h == "" {
+		return nil
+	}
+
+	var entries []tracestate.Entry
+	pairs := strings.Split(h, ",")
+	hdrLenWithoutOWS := len(pairs) - 1 // Number of commas
+	for _, pair := range pairs {
+		matches := trimOWSRegExp.FindStringSubmatch(pair)
+		if matches == nil {
+			return nil
+		}
+		pair = matches[1]
+		hdrLenWithoutOWS += len(pair)
+		if hdrLenWithoutOWS > maxTracestateLen {
+			return nil
+		}
+		kv := strings.Split(pair, "=")
+		if len(kv) != 2 {
+			return nil
+		}
+		entries = append(entries, tracestate.Entry{Key: kv[0], Value: kv[1]})
+	}
+	ts, err := tracestate.New(nil, entries...)
+	if err != nil {
+		return nil
+	}
+
+	return ts
+}
+
+func tracestateToRequest(sc trace.SpanContext, req *http.Request) {
+	var pairs = make([]string, 0, len(sc.Tracestate.Entries()))
+	if sc.Tracestate != nil {
+		for _, entry := range sc.Tracestate.Entries() {
+			pairs = append(pairs, strings.Join([]string{entry.Key, entry.Value}, "="))
+		}
+		h := strings.Join(pairs, ",")
+
+		if h != "" && len(h) <= maxTracestateLen {
+			req.Header.Set(tracestateHeader, h)
+		}
+	}
+}
+
+// SpanContextToRequest modifies the given request to include traceparent and tracestate headers.
 func (f *HTTPFormat) SpanContextToRequest(sc trace.SpanContext, req *http.Request) {
 	h := fmt.Sprintf("%x-%x-%x-%x",
 		[]byte{supportedVersion},
 		sc.TraceID[:],
 		sc.SpanID[:],
 		[]byte{byte(sc.TraceOptions)})
-	req.Header.Set(header, h)
+	req.Header.Set(traceparentHeader, h)
+	tracestateToRequest(sc, req)
 }

plugin/ochttp/propagation/tracecontext/propagation_test.go

@@ -15,11 +15,29 @@
 package tracecontext
 
 import (
+	"fmt"
 	"net/http"
 	"reflect"
 	"testing"
 
 	"go.opencensus.io/trace"
+	"go.opencensus.io/trace/tracestate"
+	"strings"
+)
+
+var (
+	tpHeader        = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+	traceID         = trace.TraceID{75, 249, 47, 53, 119, 179, 77, 166, 163, 206, 146, 157, 14, 14, 71, 54}
+	spanID          = trace.SpanID{0, 240, 103, 170, 11, 169, 2, 183}
+	traceOpt        = trace.TraceOptions(1)
+	oversizeValue   = strings.Repeat("a", maxTracestateLen/2)
+	oversizeEntry1  = tracestate.Entry{Key: "foo", Value: oversizeValue}
+	oversizeEntry2  = tracestate.Entry{Key: "hello", Value: oversizeValue}
+	entry1          = tracestate.Entry{Key: "foo", Value: "bar"}
+	entry2          = tracestate.Entry{Key: "hello", Value: "world   example"}
+	oversizeTs, _   = tracestate.New(nil, oversizeEntry1, oversizeEntry2)
+	defaultTs, _    = tracestate.New(nil, nil...)
+	nonDefaultTs, _ = tracestate.New(nil, entry1, entry2)
 )
 
 func TestHTTPFormat_FromRequest(t *testing.T) {
@@ -113,3 +131,137 @@ func TestHTTPFormat_ToRequest(t *testing.T) {
 		})
 	}
 }
+
+func TestHTTPFormatTracestate_FromRequest(t *testing.T) {
+	scWithNonDefaultTracestate := trace.SpanContext{
+		TraceID:      traceID,
+		SpanID:       spanID,
+		TraceOptions: traceOpt,
+		Tracestate:   nonDefaultTs,
+	}
+
+	scWithDefaultTracestate := trace.SpanContext{
+		TraceID:      traceID,
+		SpanID:       spanID,
+		TraceOptions: traceOpt,
+		Tracestate:   defaultTs,
+	}
+
+	tests := []struct {
+		name     string
+		tpHeader string
+		tsHeader string
+		wantSc   trace.SpanContext
+		wantOk   bool
+	}{
+		{
+			name:     "tracestate invalid entries delimiter",
+			tpHeader: tpHeader,
+			tsHeader: "foo=bar;hello=world",
+			wantSc:   scWithDefaultTracestate,
+			wantOk:   true,
+		},
+		{
+			name:     "tracestate invalid key-value delimiter",
+			tpHeader: tpHeader,
+			tsHeader: "foo=bar,hello-world",
+			wantSc:   scWithDefaultTracestate,
+			wantOk:   true,
+		},
+		{
+			name:     "tracestate invalid value character",
+			tpHeader: tpHeader,
+			tsHeader: "foo=bar,hello=world   example   \u00a0  ",
+			wantSc:   scWithDefaultTracestate,
+			wantOk:   true,
+		},
+		{
+			name:     "tracestate blank key-value",
+			tpHeader: tpHeader,
+			tsHeader: "foo=bar,    ",
+			wantSc:   scWithDefaultTracestate,
+			wantOk:   true,
+		},
+		{
+			name:     "tracestate oversize header",
+			tpHeader: tpHeader,
+			tsHeader: fmt.Sprintf("foo=%s,hello=%s", oversizeValue, oversizeValue),
+			wantSc:   scWithDefaultTracestate,
+			wantOk:   true,
+		},
+		{
+			name:     "tracestate valid",
+			tpHeader: tpHeader,
+			tsHeader: "foo=bar   ,   hello=world   example",
+			wantSc:   scWithNonDefaultTracestate,
+			wantOk:   true,
+		},
+	}
+
+	f := &HTTPFormat{}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req, _ := http.NewRequest("GET", "http://example.com", nil)
+			req.Header.Set("traceparent", tt.tpHeader)
+			req.Header.Set("tracestate", tt.tsHeader)
+
+			gotSc, gotOk := f.SpanContextFromRequest(req)
+			if !reflect.DeepEqual(gotSc, tt.wantSc) {
+				t.Errorf("HTTPFormat.FromRequest() gotTs = %v, want %v", gotSc.Tracestate, tt.wantSc.Tracestate)
+			}
+			if gotOk != tt.wantOk {
+				t.Errorf("HTTPFormat.FromRequest() gotOk = %v, want %v", gotOk, tt.wantOk)
+			}
+		})
+	}
+}
+
+func TestHTTPFormatTracestate_ToRequest(t *testing.T) {
+	tests := []struct {
+		name       string
+		sc         trace.SpanContext
+		wantHeader string
+	}{
+		{
+			name: "valid span context with default tracestate",
+			sc: trace.SpanContext{
+				TraceID:      traceID,
+				SpanID:       spanID,
+				TraceOptions: traceOpt,
+			},
+			wantHeader: "",
+		},
+		{
+			name: "valid span context with non default tracestate",
+			sc: trace.SpanContext{
+				TraceID:      traceID,
+				SpanID:       spanID,
+				TraceOptions: traceOpt,
+				Tracestate:   nonDefaultTs,
+			},
+			wantHeader: "foo=bar,hello=world   example",
+		},
+		{
+			name: "valid span context with oversize tracestate",
+			sc: trace.SpanContext{
+				TraceID:      traceID,
+				SpanID:       spanID,
+				TraceOptions: traceOpt,
+				Tracestate:   oversizeTs,
+			},
+			wantHeader: "",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			f := &HTTPFormat{}
+			req, _ := http.NewRequest("GET", "http://example.com", nil)
+			f.SpanContextToRequest(tt.sc, req)
+
+			h := req.Header.Get("tracestate")
+			if got, want := h, tt.wantHeader; got != want {
+				t.Errorf("HTTPFormat.ToRequest() tracestate header = %v, want %v", got, want)
+			}
+		})
+	}
+}