[v8] Add B3 Trace headers (#3472)

* Add B3 trace IDs to cf cli commands

* Auto set span headers for each request

Author: Samze

api/cloudcontroller/wrapper/trace_request.go

@@ -0,0 +1,31 @@
+package wrapper
+
+import (
+	"code.cloudfoundry.org/cli/api/cloudcontroller"
+	"code.cloudfoundry.org/cli/api/shared"
+)
+
+// CCTraceHeaderRequest is a wrapper that adds b3 trace headers to requests.
+type CCTraceHeaderRequest struct {
+	headers    *shared.TraceHeaders
+	connection cloudcontroller.Connection
+}
+
+// NewCCTraceHeaderRequest returns a pointer to a CCTraceHeaderRequest wrapper.
+func NewCCTraceHeaderRequest(trace string) *CCTraceHeaderRequest {
+	return &CCTraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace),
+	}
+}
+
+// Add tracing headers
+func (t *CCTraceHeaderRequest) Make(request *cloudcontroller.Request, passedResponse *cloudcontroller.Response) error {
+	t.headers.SetHeaders(request.Request)
+	return t.connection.Make(request, passedResponse)
+}
+
+// Wrap sets the connection in the CCTraceHeaderRequest and returns itself.
+func (t *CCTraceHeaderRequest) Wrap(innerconnection cloudcontroller.Connection) cloudcontroller.Connection {
+	t.connection = innerconnection
+	return t
+}

api/cloudcontroller/wrapper/trace_request_test.go

@@ -0,0 +1,65 @@
+package wrapper_test
+
+import (
+	"bytes"
+	"net/http"
+
+	"code.cloudfoundry.org/cli/api/cloudcontroller"
+	"code.cloudfoundry.org/cli/api/cloudcontroller/cloudcontrollerfakes"
+	. "code.cloudfoundry.org/cli/api/cloudcontroller/wrapper"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("CCTraceHeaderRequest", func() {
+	var (
+		fakeConnection *cloudcontrollerfakes.FakeConnection
+
+		wrapper cloudcontroller.Connection
+
+		request  *cloudcontroller.Request
+		response *cloudcontroller.Response
+		makeErr  error
+
+		traceHeader string
+	)
+
+	BeforeEach(func() {
+		fakeConnection = new(cloudcontrollerfakes.FakeConnection)
+
+		traceHeader = "trace-id"
+
+		wrapper = NewCCTraceHeaderRequest(traceHeader).Wrap(fakeConnection)
+
+		body := bytes.NewReader([]byte("foo"))
+
+		req, err := http.NewRequest(http.MethodGet, "https://foo.bar.com/banana", body)
+		Expect(err).NotTo(HaveOccurred())
+
+		response = &cloudcontroller.Response{
+			RawResponse:  []byte("some-response-body"),
+			HTTPResponse: &http.Response{},
+		}
+		request = cloudcontroller.NewRequest(req, body)
+	})
+
+	JustBeforeEach(func() {
+		makeErr = wrapper.Make(request, response)
+	})
+
+	Describe("Make", func() {
+		It("Adds the request headers", func() {
+			Expect(makeErr).NotTo(HaveOccurred())
+			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
+		})
+
+		It("Calls the inner connection", func() {
+			Expect(fakeConnection.MakeCallCount()).To(Equal(1))
+			req, resp := fakeConnection.MakeArgsForCall(0)
+			Expect(req).To(Equal(request))
+			Expect(resp).To(Equal(response))
+		})
+	})
+})

api/router/wrapper/trace_request.go

@@ -0,0 +1,31 @@
+package wrapper
+
+import (
+	"code.cloudfoundry.org/cli/api/router"
+	"code.cloudfoundry.org/cli/api/shared"
+)
+
+// RoutingTraceHeaderRequest is a wrapper that adds b3 trace headers to requests.
+type RoutingTraceHeaderRequest struct {
+	headers    *shared.TraceHeaders
+	connection router.Connection
+}
+
+// NewRoutingTraceHeaderRequest returns a pointer to a RoutingTraceHeaderRequest wrapper.
+func NewRoutingTraceHeaderRequest(trace string) *RoutingTraceHeaderRequest {
+	return &RoutingTraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace),
+	}
+}
+
+// Add tracing headers
+func (t *RoutingTraceHeaderRequest) Make(request *router.Request, passedResponse *router.Response) error {
+	t.headers.SetHeaders(request.Request)
+	return t.connection.Make(request, passedResponse)
+}
+
+// Wrap sets the connection in the RoutingTraceHeaderRequest and returns itself.
+func (t *RoutingTraceHeaderRequest) Wrap(innerconnection router.Connection) router.Connection {
+	t.connection = innerconnection
+	return t
+}

api/router/wrapper/trace_request_test.go

@@ -0,0 +1,64 @@
+package wrapper_test
+
+import (
+	"bytes"
+	"net/http"
+
+	"code.cloudfoundry.org/cli/api/router"
+	"code.cloudfoundry.org/cli/api/router/routerfakes"
+	. "code.cloudfoundry.org/cli/api/router/wrapper"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("CCTraceHeaderRequest", func() {
+	var (
+		fakeConnection *routerfakes.FakeConnection
+
+		wrapper router.Connection
+
+		request  *router.Request
+		response *router.Response
+		makeErr  error
+
+		traceHeader string
+	)
+
+	BeforeEach(func() {
+		fakeConnection = new(routerfakes.FakeConnection)
+
+		traceHeader = "trace-id"
+		wrapper = NewRoutingTraceHeaderRequest(traceHeader).Wrap(fakeConnection)
+
+		body := bytes.NewReader([]byte("foo"))
+
+		req, err := http.NewRequest(http.MethodGet, "https://foo.bar.com/banana", body)
+		Expect(err).NotTo(HaveOccurred())
+
+		response = &router.Response{
+			RawResponse:  []byte("some-response-body"),
+			HTTPResponse: &http.Response{},
+		}
+		request = router.NewRequest(req, body)
+	})
+
+	JustBeforeEach(func() {
+		makeErr = wrapper.Make(request, response)
+	})
+
+	Describe("Make", func() {
+		It("Adds the request headers", func() {
+			Expect(makeErr).NotTo(HaveOccurred())
+			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
+		})
+
+		It("Calls the inner connection", func() {
+			Expect(fakeConnection.MakeCallCount()).To(Equal(1))
+			req, resp := fakeConnection.MakeArgsForCall(0)
+			Expect(req).To(Equal(request))
+			Expect(resp).To(Equal(response))
+		})
+	})
+})

api/shared/trace_headers.go

@@ -0,0 +1,37 @@
+package shared
+
+import (
+	"net/http"
+
+	"code.cloudfoundry.org/cli/util/trace"
+)
+
+const (
+	B3TraceIDHeader = "X-B3-TraceId"
+	B3SpanIDHeader  = "X-B3-SpanId"
+)
+
+// TraceHeaders sets b3 trace headers to requests.
+type TraceHeaders struct {
+	b3trace string
+}
+
+// NewTraceHeaders returns a pointer to a TraceHeaderRequest.
+func NewTraceHeaders(trace string) *TraceHeaders {
+	return &TraceHeaders{
+		b3trace: trace,
+	}
+}
+
+// Add tracing headers if they are not already set.
+func (t *TraceHeaders) SetHeaders(request *http.Request) {
+	// only override the trace headers if they are not already set (e.g. already explicitly set by cf curl)
+	if request.Header.Get(B3TraceIDHeader) == "" {
+		request.Header.Add(B3TraceIDHeader, t.b3trace)
+	}
+	if request.Header.Get(B3SpanIDHeader) == "" {
+		request.Header.Add(B3SpanIDHeader, trace.GenerateRandomTraceID(16))
+	}
+
+	// request.Header.Add(("B3", request.Header.Get(B3TraceIDHeader)+request.Header.Get(B3SpanIDHeader)))
+}

api/shared/trace_headers_test.go

@@ -0,0 +1,42 @@
+package shared_test
+
+import (
+	"net/http"
+
+	. "code.cloudfoundry.org/cli/api/shared"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("B3 Trace Headers", func() {
+	Describe("SetHeaders", func() {
+		Context("when there are already headers set", func() {
+			It("does not add the headers", func() {
+				traceHeaders := NewTraceHeaders("new_trace_id")
+				request := &http.Request{
+					Header: http.Header{},
+				}
+				request.Header.Set("X-B3-TraceId", "old_trace_id")
+				request.Header.Set("X-B3-SpanId", "old_span_id")
+				traceHeaders.SetHeaders(request)
+
+				Expect(request.Header.Get("X-B3-TraceId")).To(Equal("old_trace_id"))
+				Expect(request.Header.Get("X-B3-SpanId")).To(Equal("old_span_id"))
+			})
+		})
+
+		Context("when there are no headers set", func() {
+			It("adds the headers", func() {
+				traceHeaders := NewTraceHeaders("new_trace_id")
+				request := &http.Request{
+					Header: http.Header{},
+				}
+				traceHeaders.SetHeaders(request)
+
+				Expect(request.Header.Get("X-B3-TraceId")).To(Equal("new_trace_id"))
+				Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
+			})
+		})
+	})
+})

api/uaa/wrapper/trace_request.go

@@ -0,0 +1,33 @@
+package wrapper
+
+import (
+	"net/http"
+
+	"code.cloudfoundry.org/cli/api/shared"
+	"code.cloudfoundry.org/cli/api/uaa"
+)
+
+// UAATraceHeaderRequest is a wrapper that adds b3 trace headers to requests.
+type UAATraceHeaderRequest struct {
+	headers    *shared.TraceHeaders
+	connection uaa.Connection
+}
+
+// NewUAATraceHeaderRequest returns a pointer to a UAATraceHeaderRequest wrapper.
+func NewUAATraceHeaderRequest(trace string) *UAATraceHeaderRequest {
+	return &UAATraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace),
+	}
+}
+
+// Add tracing headers
+func (t *UAATraceHeaderRequest) Make(request *http.Request, passedResponse *uaa.Response) error {
+	t.headers.SetHeaders(request)
+	return t.connection.Make(request, passedResponse)
+}
+
+// Wrap sets the connection in the UAATraceHeaderRequest and returns itself.
+func (t *UAATraceHeaderRequest) Wrap(innerconnection uaa.Connection) uaa.Connection {
+	t.connection = innerconnection
+	return t
+}

api/uaa/wrapper/trace_request_test.go

@@ -0,0 +1,65 @@
+package wrapper_test
+
+import (
+	"bytes"
+	"net/http"
+
+	"code.cloudfoundry.org/cli/api/uaa"
+	"code.cloudfoundry.org/cli/api/uaa/uaafakes"
+	. "code.cloudfoundry.org/cli/api/uaa/wrapper"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("CCTraceHeaderRequest", func() {
+	var (
+		fakeConnection *uaafakes.FakeConnection
+
+		wrapper uaa.Connection
+
+		request  *http.Request
+		response *uaa.Response
+		makeErr  error
+
+		traceHeader string
+	)
+
+	BeforeEach(func() {
+		fakeConnection = new(uaafakes.FakeConnection)
+
+		traceHeader = "trace-id"
+
+		wrapper = NewUAATraceHeaderRequest(traceHeader).Wrap(fakeConnection)
+
+		body := bytes.NewReader([]byte("foo"))
+
+		var err error
+		request, err = http.NewRequest(http.MethodGet, "https://foo.bar.com/banana", body)
+		Expect(err).NotTo(HaveOccurred())
+
+		response = &uaa.Response{
+			RawResponse:  []byte("some-response-body"),
+			HTTPResponse: &http.Response{},
+		}
+	})
+
+	JustBeforeEach(func() {
+		makeErr = wrapper.Make(request, response)
+	})
+
+	Describe("Make", func() {
+		It("Adds the request headers", func() {
+			Expect(makeErr).NotTo(HaveOccurred())
+			Expect(request.Header.Get("X-B3-TraceId")).To(Equal(traceHeader))
+			Expect(request.Header.Get("X-B3-SpanId")).ToNot(BeEmpty())
+		})
+
+		It("Calls the inner connection", func() {
+			Expect(fakeConnection.MakeCallCount()).To(Equal(1))
+			req, resp := fakeConnection.MakeArgsForCall(0)
+			Expect(req).To(Equal(request))
+			Expect(resp).To(Equal(response))
+		})
+	})
+})