From 7f7e4d352c2fef3a70b324191ae3c6ba6ae9e1cb Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Wed, 25 Jul 2018 18:35:58 -0400
Subject: [PATCH] logging: populate entry trace from HTTP request header

If we see the X-Cloud-Trace-Context header and the user hasn't already
set the trace field, populate the trace field with the header.

See #448.

Change-Id: I4870616d6ff5e94ea3af94ff2fb5a8b346335766
Reviewed-on: https://code-review.googlesource.com/30890
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jean de Klerk <deklerk@google.com>
---
 logging/logging.go                 |  3 ++
 logging/logging_unexported_test.go | 52 ++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/logging/logging.go b/logging/logging.go
index 52e113495ec..c64d8e33be0 100644
--- a/logging/logging.go
+++ b/logging/logging.go
@@ -794,6 +794,9 @@ func toLogEntry(e Entry) (*logpb.LogEntry, error) {
 	if err != nil {
 		return nil, err
 	}
+	if e.Trace == "" && e.HTTPRequest != nil && e.HTTPRequest.Request != nil {
+		e.Trace = e.HTTPRequest.Request.Header.Get("X-Cloud-Trace-Context")
+	}
 	ent := &logpb.LogEntry{
 		Timestamp:      ts,
 		Severity:       logtypepb.LogSeverity(e.Severity),
diff --git a/logging/logging_unexported_test.go b/logging/logging_unexported_test.go
index 4c8cc651197..64c12c3b618 100644
--- a/logging/logging_unexported_test.go
+++ b/logging/logging_unexported_test.go
@@ -227,6 +227,58 @@ func TestToLogEntryPayload(t *testing.T) {
 	}
 }
 
+func TestToLogEntryTrace(t *testing.T) {
+	// Verify that we get the trace from the HTTP request if it isn't
+	// provided by the caller.
+	u := &url.URL{Scheme: "http"}
+	for _, test := range []struct {
+		in   Entry
+		want string
+	}{
+		{Entry{}, ""},
+		{Entry{Trace: "t1"}, "t1"},
+		{
+			Entry{
+				HTTPRequest: &HTTPRequest{
+					Request: &http.Request{URL: u, Header: http.Header{"foo": {"bar"}}},
+				},
+			},
+			"",
+		},
+		{
+			Entry{
+				HTTPRequest: &HTTPRequest{
+					Request: &http.Request{
+						URL:    u,
+						Header: http.Header{"X-Cloud-Trace-Context": {"t2"}},
+					},
+				},
+			},
+			"t2",
+		},
+		{
+			Entry{
+				HTTPRequest: &HTTPRequest{
+					Request: &http.Request{
+						URL:    u,
+						Header: http.Header{"X-Cloud-Trace-Context": {"t3"}},
+					},
+				},
+				Trace: "t4",
+			},
+			"t4",
+		},
+	} {
+		e, err := toLogEntry(test.in)
+		if err != nil {
+			t.Fatalf("%+v: %v", test.in, err)
+		}
+		if got := e.Trace; got != test.want {
+			t.Errorf("got %q, want %q", got, test.want)
+		}
+	}
+}
+
 func TestFromHTTPRequest(t *testing.T) {
 	const testURL = "http:://example.com/path?q=1"
 	u, err := url.Parse(testURL)