Add B3 trace IDs to cf cli commands

Author: Samze

api/cloudcontroller/wrapper/b3_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, span string) *CCTraceHeaderRequest {
+	return &CCTraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace, span),
+	}
+}
+
+// Add tracing headers
+func (t *CCTraceHeaderRequest) Make(request *cloudcontroller.Request, passedResponse *cloudcontroller.Response) error {
+	t.headers.SetHeaders(request.Request, passedResponse.HTTPResponse)
+	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/router/wrapper/b3_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, span string) *RoutingTraceHeaderRequest {
+	return &RoutingTraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace, span),
+	}
+}
+
+// Add tracing headers
+func (t *RoutingTraceHeaderRequest) Make(request *router.Request, passedResponse *router.Response) error {
+	t.headers.SetHeaders(request.Request, passedResponse.HTTPResponse)
+	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/shared/b3_trace_headers.go

@@ -0,0 +1,38 @@
+package shared
+
+import (
+	"net/http"
+)
+
+const (
+	B3TraceIDHeader = "X-B3-TraceId"
+	B3SpanIDHeader  = "X-B3-SpanId"
+)
+
+// TODO
+// 1. tests
+
+// TraceHeaders sets b3 trace headers to requests.
+type TraceHeaders struct {
+	b3trace string
+	b3span  string
+}
+
+// NewTraceHeaders returns a pointer to a TraceHeaderRequest.
+func NewTraceHeaders(trace, span string) *TraceHeaders {
+	return &TraceHeaders{
+		b3trace: trace,
+		b3span:  span,
+	}
+}
+
+// Add tracing headers if they are not already set.
+func (t *TraceHeaders) SetHeaders(request *http.Request, passedResponse *http.Response) {
+	// 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, t.b3span)
+	}
+}

api/uaa/wrapper/b3_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, span string) *UAATraceHeaderRequest {
+	return &UAATraceHeaderRequest{
+		headers: shared.NewTraceHeaders(trace, span),
+	}
+}
+
+// Add tracing headers
+func (t *UAATraceHeaderRequest) Make(request *http.Request, passedResponse *uaa.Response) error {
+	t.headers.SetHeaders(request, passedResponse.HTTPResponse)
+	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
+}

command/commandfakes/fake_config.go

@@ -15,6 +15,8 @@ type Config interface {
 	AddPluginRepository(name string, url string)
 	AuthorizationEndpoint() string
 	APIVersion() string
+	B3TraceID() string
+	B3SpanID() string
 	BinaryName() string
 	BinaryVersion() string
 	CFPassword() string

command/config.go

@@ -47,6 +47,7 @@ func NewWrappedCloudControllerClient(config command.Config, ui command.UI, extra
 	}
 
 	ccWrappers = append(ccWrappers, extraWrappers...)
+	ccWrappers = append(ccWrappers, ccWrapper.NewCCTraceHeaderRequest(config.B3TraceID(), config.B3SpanID()))
 	ccWrappers = append(ccWrappers, ccWrapper.NewRetryRequest(config.RequestRetryCount()))
 
 	return ccv3.NewClient(ccv3.Config{
@@ -85,6 +86,7 @@ func newWrappedUAAClient(config command.Config, ui command.UI) (*uaa.Client, err
 
 	uaaAuthWrapper := uaaWrapper.NewUAAAuthentication(uaaClient, config)
 	uaaClient.WrapConnection(uaaAuthWrapper)
+	uaaClient.WrapConnection(uaaWrapper.NewUAATraceHeaderRequest(config.B3TraceID(), config.B3SpanID()))
 	uaaClient.WrapConnection(uaaWrapper.NewRetryRequest(config.RequestRetryCount()))
 
 	err = uaaClient.SetupResources(config.UAAEndpoint(), config.AuthorizationEndpoint())

command/v7/shared/new_clients.go

@@ -5,6 +5,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"code.cloudfoundry.org/cli/util/random"
 )
 
 // EnvOverride represents all the environment variables read by the CF CLI
@@ -20,6 +22,8 @@ type EnvOverride struct {
 	CFStartupTimeout string
 	CFTrace          string
 	CFUsername       string
+	CFB3TraceID      string
+	CFB3SpanID       string
 	DockerPassword   string
 	CNBCredentials   string
 	Experimental     string
@@ -160,3 +164,17 @@ func (config *Config) StartupTimeout() time.Duration {
 
 	return DefaultStartupTimeout
 }
+
+func (config *Config) B3TraceID() string {
+	if config.ENV.CFB3TraceID == "" {
+		config.ENV.CFB3TraceID = random.GenerateHex(32)
+	}
+	return config.ENV.CFB3TraceID
+}
+
+func (config *Config) B3SpanID() string {
+	if config.ENV.CFB3SpanID == "" {
+		config.ENV.CFB3SpanID = random.GenerateHex(16)
+	}
+	return config.ENV.CFB3SpanID
+}

util/configv3/env.go

@@ -127,6 +127,8 @@ func LoadConfig(flags ...FlagOverride) (*Config, error) {
 		CFStartupTimeout: os.Getenv("CF_STARTUP_TIMEOUT"),
 		CFTrace:          os.Getenv("CF_TRACE"),
 		CFUsername:       os.Getenv("CF_USERNAME"),
+		CFB3TraceID:      os.Getenv("CF_B3_TRACE_ID"),
+		CFB3SpanID:       os.Getenv("CF_B3_SPAN_ID"),
 		DockerPassword:   os.Getenv("CF_DOCKER_PASSWORD"),
 		CNBCredentials:   os.Getenv("CNB_REGISTRY_CREDS"),
 		Experimental:     os.Getenv("CF_CLI_EXPERIMENTAL"),

util/configv3/load_config.go

@@ -0,0 +1,15 @@
+package random
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+)
+
+func GenerateHex(length int) string {
+	b := make([]byte, length/2)
+	if _, err := rand.Read(b); err != nil {
+		panic(err)
+	}
+
+	return hex.EncodeToString(b)
+}