Fix traced mTLS identity client transport

Author: cursoragent

identityclient/client.go

@@ -18,7 +18,6 @@ import (
 
 	identityv1 "github.com/evalops/proto/gen/go/identity/v1"
 	"github.com/evalops/service-runtime/mtls"
-	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -159,7 +158,7 @@ func New(config Config) *Client {
 		httpClient = http.DefaultClient
 	}
 	usesMTLSCert := httpClientUsesMTLSCertificate(httpClient)
-	httpClient = tracedHTTPClient(httpClient)
+	httpClient = mtls.TraceHTTPClient(httpClient)
 	maxSize := config.MaxCacheSize
 	if maxSize <= 0 {
 		maxSize = defaultMaxCacheSize
@@ -177,19 +176,6 @@ func New(config Config) *Client {
 	}
 }
 
-func tracedHTTPClient(client *http.Client) *http.Client {
-	if client == nil {
-		client = http.DefaultClient
-	}
-	cloned := *client
-	baseTransport := cloned.Transport
-	if baseTransport == nil {
-		baseTransport = http.DefaultTransport
-	}
-	cloned.Transport = otelhttp.NewTransport(baseTransport)
-	return &cloned
-}
-
 // NewClient creates a Client that introspects tokens at the given URL.
 func NewClient(introspectURL string, requestTimeout time.Duration, httpClient *http.Client) *Client {
 	return New(Config{
@@ -422,11 +408,22 @@ func httpClientUsesMTLSCertificate(client *http.Client) bool {
 	if client == nil {
 		return false
 	}
-	transport, ok := client.Transport.(*http.Transport)
-	if !ok || transport == nil {
-		return false
+	transport := client.Transport
+	if transport == nil {
+		transport = http.DefaultTransport
+	}
+	for transport != nil {
+		httpTransport, ok := transport.(*http.Transport)
+		if ok {
+			return tlsConfigHasClientCertificate(httpTransport.TLSClientConfig)
+		}
+		wrapped, ok := transport.(interface{ Unwrap() http.RoundTripper })
+		if !ok {
+			return false
+		}
+		transport = wrapped.Unwrap()
 	}
-	return tlsConfigHasClientCertificate(transport.TLSClientConfig)
+	return false
 }
 
 func tlsConfigHasClientCertificate(cfg *tls.Config) bool {

identityclient/client_test.go

@@ -18,6 +18,7 @@ import (
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 )
 
 type roundTripFunc func(*http.Request) (*http.Response, error)
@@ -67,6 +68,65 @@ func TestConfigured(t *testing.T) {
 	}
 }
 
+func TestConfiguredDetectsMTLSCertificatesThroughTracedTransport(t *testing.T) {
+	if !New(Config{
+		ServiceTokensURL: "https://identity.internal/v1/service-tokens",
+		HTTPClient: mtls.TraceHTTPClient(&http.Client{
+			Transport: &http.Transport{
+				TLSClientConfig: &tls.Config{Certificates: []tls.Certificate{{}}},
+			},
+		}),
+	}).ServiceTokensConfigured() {
+		t.Fatal("expected traced mtls-authenticated service tokens to be configured")
+	}
+}
+
+func TestNewDoesNotDoubleWrapTracedHTTPClient(t *testing.T) {
+	originalProvider := otel.GetTracerProvider()
+	originalPropagator := otel.GetTextMapPropagator()
+	recorder := tracetest.NewSpanRecorder()
+	tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	otel.SetTracerProvider(tracerProvider)
+	otel.SetTextMapPropagator(propagation.TraceContext{})
+	t.Cleanup(func() {
+		otel.SetTracerProvider(originalProvider)
+		otel.SetTextMapPropagator(originalPropagator)
+		_ = tracerProvider.Shutdown(context.Background())
+	})
+
+	httpClient := mtls.TraceHTTPClient(&http.Client{
+		Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       io.NopCloser(strings.NewReader(`{"active":true,"organization_id":"org-123"}`)),
+				Header:     make(http.Header),
+			}, nil
+		}),
+	})
+	client := New(Config{
+		IntrospectURL:  "https://identity.test/v1/tokens/introspect",
+		RequestTimeout: time.Second,
+		HTTPClient:     httpClient,
+	})
+
+	ctx, span := tracerProvider.Tracer("identityclient-test").Start(context.Background(), "root")
+	defer span.End()
+
+	if _, err := client.Introspect(ctx, "write-token"); err != nil {
+		t.Fatalf("introspect: %v", err)
+	}
+
+	httpSpanCount := 0
+	for _, ended := range recorder.Ended() {
+		if ended.Name() == "HTTP POST" {
+			httpSpanCount++
+		}
+	}
+	if httpSpanCount != 1 {
+		t.Fatalf("expected one HTTP client span, got %d", httpSpanCount)
+	}
+}
+
 func TestIntrospectSuccess(t *testing.T) {
 	server := testutil.NewTestServer(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		if got := request.Header.Get("Authorization"); got != "Bearer write-token" {

mtls/mtls.go

@@ -69,7 +69,7 @@ func BuildHTTPClient(cfg ClientConfig) (*http.Client, error) {
 		return nil, err
 	}
 	if tlsConfig == nil {
-		return traceHTTPClient(http.DefaultClient), nil
+		return TraceHTTPClient(http.DefaultClient), nil
 	}
 
 	transport, ok := http.DefaultTransport.(*http.Transport)
@@ -78,10 +78,11 @@ func BuildHTTPClient(cfg ClientConfig) (*http.Client, error) {
 	}
 	clone := transport.Clone()
 	clone.TLSClientConfig = tlsConfig
-	return traceHTTPClient(&http.Client{Transport: clone}), nil
+	return TraceHTTPClient(&http.Client{Transport: clone}), nil
 }
 
-func traceHTTPClient(client *http.Client) *http.Client {
+// TraceHTTPClient clones client and wraps its transport with OTel propagation.
+func TraceHTTPClient(client *http.Client) *http.Client {
 	if client == nil {
 		client = http.DefaultClient
 	}
@@ -90,10 +91,39 @@ func traceHTTPClient(client *http.Client) *http.Client {
 	if transport == nil {
 		transport = http.DefaultTransport
 	}
-	cloned.Transport = otelhttp.NewTransport(transport)
+	cloned.Transport = traceRoundTripper(transport)
 	return &cloned
 }
 
+func traceRoundTripper(transport http.RoundTripper) http.RoundTripper {
+	if transport == nil {
+		transport = http.DefaultTransport
+	}
+	if _, ok := transport.(*tracedTransport); ok {
+		return transport
+	}
+	if _, ok := transport.(*otelhttp.Transport); ok {
+		return transport
+	}
+	return &tracedTransport{
+		base:   transport,
+		traced: otelhttp.NewTransport(transport),
+	}
+}
+
+type tracedTransport struct {
+	base   http.RoundTripper
+	traced http.RoundTripper
+}
+
+func (t *tracedTransport) RoundTrip(request *http.Request) (*http.Response, error) {
+	return t.traced.RoundTrip(request)
+}
+
+func (t *tracedTransport) Unwrap() http.RoundTripper {
+	return t.base
+}
+
 // BuildClientTLSConfig returns a *tls.Config for an outbound mTLS client, or nil when all fields are empty.
 func BuildClientTLSConfig(cfg ClientConfig) (*tls.Config, error) {
 	if cfg.CAFile == "" && cfg.CertFile == "" && cfg.KeyFile == "" && cfg.ServerName == "" {