From 52f7d778c2eea02e3e5bcf638c590c533127baa8 Mon Sep 17 00:00:00 2001
From: Arne Luenser <arne.luenser@ory.sh>
Date: Tue, 6 Aug 2024 15:15:14 +0200
Subject: [PATCH] fix: reuse OpenTelemetry-wrapped http.Transport (#802)

---
 httpx/resilient_client.go |  9 ---------
 httpx/ssrf.go             | 17 +++++++++++++----
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/httpx/resilient_client.go b/httpx/resilient_client.go
index afd66199..8e5b4537 100644
--- a/httpx/resilient_client.go
+++ b/httpx/resilient_client.go
@@ -8,11 +8,8 @@ import (
 	"io"
 	"log"
 	"net/http"
-	"net/http/httptrace"
 	"time"
 
-	"go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace"
-	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/oauth2"
 
@@ -130,12 +127,6 @@ func NewResilientClient(opts ...ResilientOptions) *retryablehttp.Client {
 		o.c.Transport = ifelse(o.ipV6, allowInternalAllowIPv6, allowInternalProhibitIPv6)
 	}
 
-	if o.tracer != nil {
-		o.c.Transport = otelhttp.NewTransport(o.c.Transport, otelhttp.WithClientTrace(func(ctx context.Context) *httptrace.ClientTrace {
-			return otelhttptrace.NewClientTrace(ctx, otelhttptrace.WithoutHeaders(), otelhttptrace.WithoutSubSpans())
-		}))
-	}
-
 	cl := retryablehttp.NewClient()
 	cl.HTTPClient = o.c
 	cl.Logger = o.l
diff --git a/httpx/ssrf.go b/httpx/ssrf.go
index b92d579e..99b16e9e 100644
--- a/httpx/ssrf.go
+++ b/httpx/ssrf.go
@@ -7,11 +7,14 @@ import (
 	"context"
 	"net"
 	"net/http"
+	"net/http/httptrace"
 	"net/netip"
 	"time"
 
 	"code.dny.dev/ssrf"
 	"github.com/gobwas/glob"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 )
 
 var _ http.RoundTripper = (*noInternalIPRoundTripper)(nil)
@@ -64,7 +67,7 @@ func init() {
 		ssrf.WithAnyPort(),
 		ssrf.WithNetworks("tcp4", "tcp6"),
 	).Safe
-	prohibitInternalAllowIPv6 = t
+	prohibitInternalAllowIPv6 = otelTransport(t)
 }
 
 func init() {
@@ -76,7 +79,7 @@ func init() {
 	t.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
 		return d.DialContext(ctx, "tcp4", addr)
 	}
-	prohibitInternalProhibitIPv6 = t
+	prohibitInternalProhibitIPv6 = otelTransport(t)
 }
 
 func init() {
@@ -96,7 +99,7 @@ func init() {
 			netip.MustParsePrefix("fc00::/7"), // Unique Local (RFC 4193)
 		),
 	).Safe
-	allowInternalAllowIPv6 = t
+	allowInternalAllowIPv6 = otelTransport(t)
 }
 
 func init() {
@@ -119,7 +122,7 @@ func init() {
 	t.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
 		return d.DialContext(ctx, "tcp4", addr)
 	}
-	allowInternalProhibitIPv6 = t
+	allowInternalProhibitIPv6 = otelTransport(t)
 }
 
 func newDefaultTransport() (*http.Transport, *net.Dialer) {
@@ -137,3 +140,9 @@ func newDefaultTransport() (*http.Transport, *net.Dialer) {
 		ExpectContinueTimeout: 1 * time.Second,
 	}, &dialer
 }
+
+func otelTransport(t *http.Transport) http.RoundTripper {
+	return otelhttp.NewTransport(t, otelhttp.WithClientTrace(func(ctx context.Context) *httptrace.ClientTrace {
+		return otelhttptrace.NewClientTrace(ctx, otelhttptrace.WithoutHeaders(), otelhttptrace.WithoutSubSpans())
+	}))
+}