review: TLS test case, comments, span start extraction

Author: pawbana

passthrough.go

@@ -1,6 +1,7 @@
 package aibridge
 
 import (
+	"context"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -58,10 +59,7 @@ func newPassthroughRouter(prov provider.Provider, logger slog.Logger, m *metrics
 			m.PassthroughCount.WithLabelValues(prov.Name(), r.URL.Path, r.Method).Add(1)
 		}
 
-		ctx, span := tracer.Start(r.Context(), "Passthrough", trace.WithAttributes(
-			attribute.String(tracing.PassthroughURL, r.URL.String()),
-			attribute.String(tracing.PassthroughMethod, r.Method),
-		))
+		ctx, span := startSpan(r, tracer)
 		defer span.End()
 
 		proxy.ServeHTTP(w, r.WithContext(ctx))
@@ -73,6 +71,10 @@ func newPassthroughRouter(prov provider.Provider, logger slog.Logger, m *metrics
 func rewritePassthroughRequest(pr *httputil.ProxyRequest, provBaseURL *url.URL, prov provider.Provider) {
 	pr.SetURL(provBaseURL)
 
+	// Rewrite sets "X-Forwarded-For" to just last hop (clients IP address).
+	// To preserve old Director behavior pr.In "X-Forwarded-For" header
+	// values need to be copied manually.
+	// https://pkg.go.dev/net/http/httputil#ProxyRequest.SetXForwarded
 	if prior, ok := pr.In.Header["X-Forwarded-For"]; ok {
 		pr.Out.Header["X-Forwarded-For"] = append([]string(nil), prior...)
 	}
@@ -90,22 +92,26 @@ func rewritePassthroughRequest(pr *httputil.ProxyRequest, provBaseURL *url.URL,
 	prov.InjectAuthHeader(&pr.Out.Header)
 }
 
-// newInvalidBaseURLHandler returns a handler that always returns 502 because
-// the provider's base URL is invalid.
+// newInvalidBaseURLHandler returns a handler that always returns 502
+// when the provider's base URL is invalid.
 func newInvalidBaseURLHandler(prov provider.Provider, logger slog.Logger, m *metrics.Metrics, tracer trace.Tracer, baseURLErr error) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
+		ctx, span := startSpan(r, tracer)
+		defer span.End()
+
 		if m != nil {
 			m.PassthroughCount.WithLabelValues(prov.Name(), r.URL.Path, r.Method).Add(1)
 		}
 
-		ctx, span := tracer.Start(r.Context(), "Passthrough", trace.WithAttributes(
-			attribute.String(tracing.PassthroughURL, r.URL.String()),
-			attribute.String(tracing.PassthroughMethod, r.Method),
-		))
-		defer span.End()
-
 		logger.Warn(ctx, "invalid provider base URL", slog.Error(baseURLErr))
 		http.Error(w, "invalid provider base URL", http.StatusBadGateway)
 		span.SetStatus(codes.Error, "invalid provider base URL: "+baseURLErr.Error())
 	}
 }
+
+func startSpan(r *http.Request, tracer trace.Tracer) (context.Context, trace.Span) {
+	return tracer.Start(r.Context(), "Passthrough", trace.WithAttributes(
+		attribute.String(tracing.PassthroughURL, r.URL.String()),
+		attribute.String(tracing.PassthroughMethod, r.Method),
+	))
+}

passthrough_test.go

@@ -140,7 +140,6 @@ func TestRewritePassthroughRequest(t *testing.T) {
 		reqRemoteAddr string
 		reqHeaders    http.Header
 		reqTLS        bool
-		baseURL       string
 		provider      *testutil.MockProvider
 		expectURL     string
 		expectHeaders http.Header
@@ -149,7 +148,6 @@ func TestRewritePassthroughRequest(t *testing.T) {
 			name:          "sets_upstream_url_and_forwarded_headers_from_client_peer",
 			reqPath:       "http://client-host/chat?stream=true",
 			reqRemoteAddr: "1.1.1.1:1111",
-			baseURL:       "https://upstream-host/base",
 			provider:      &testutil.MockProvider{URL: "https://upstream-host/base"},
 			expectURL:     "https://upstream-host/base/chat?stream=true",
 			expectHeaders: http.Header{
@@ -164,7 +162,6 @@ func TestRewritePassthroughRequest(t *testing.T) {
 			reqPath:       "http://client-host/chat",
 			reqRemoteAddr: "1.1.1.1:1111",
 			reqHeaders:    http.Header{"User-Agent": {"custom-agent/1.0"}},
-			baseURL:       "https://upstream-host/base",
 			provider:      &testutil.MockProvider{URL: "https://upstream-host/base"},
 			expectURL:     "https://upstream-host/base/chat",
 			expectHeaders: http.Header{
@@ -178,7 +175,6 @@ func TestRewritePassthroughRequest(t *testing.T) {
 			name:          "injects_auth_header",
 			reqPath:       "http://client-host/chat",
 			reqRemoteAddr: "1.1.1.1:1111",
-			baseURL:       "https://upstream-host/base",
 			provider: &testutil.MockProvider{
 				URL: "https://upstream-host/base",
 				InjectAuthHeaderFunc: func(h *http.Header) {
@@ -201,7 +197,6 @@ func TestRewritePassthroughRequest(t *testing.T) {
 			reqHeaders: http.Header{
 				"X-Forwarded-For": {"2.2.2.2, 3.3.3.3"},
 			},
-			baseURL:   "https://upstream-host/base",
 			provider:  &testutil.MockProvider{URL: "https://upstream-host/base"},
 			expectURL: "https://upstream-host/base/chat",
 			expectHeaders: http.Header{
@@ -211,14 +206,27 @@ func TestRewritePassthroughRequest(t *testing.T) {
 				"User-Agent":        {"aibridge"},
 			},
 		},
+		{
+			name:          "tls_request_sets_forwarded_proto_to_https",
+			reqPath:       "http://client-host/chat",
+			reqRemoteAddr: "1.1.1.1:1111",
+			reqTLS:        true,
+			provider:      &testutil.MockProvider{URL: "https://upstream-host/base"},
+			expectURL:     "https://upstream-host/base/chat",
+			expectHeaders: http.Header{
+				"X-Forwarded-Host":  {"client-host"},
+				"X-Forwarded-Proto": {"https"},
+				"X-Forwarded-For":   {"1.1.1.1"},
+				"User-Agent":        {"aibridge"},
+			},
+		},
 		{
 			name:          "omits_forwarded_for_when_remote_addr_is_not_parseable",
 			reqPath:       "http://client-host/chat",
 			reqRemoteAddr: "not-a-socket-address",
 			reqHeaders: http.Header{
 				"X-Forwarded-For": {"1.1.1.1"},
 			},
-			baseURL:   "https://upstream-host/base",
 			provider:  &testutil.MockProvider{URL: "https://upstream-host/base"},
 			expectURL: "https://upstream-host/base/chat",
 			expectHeaders: http.Header{
@@ -239,7 +247,7 @@ func TestRewritePassthroughRequest(t *testing.T) {
 			if tc.reqTLS {
 				r.TLS = &tls.ConnectionState{}
 			}
-			provBaseURL, err := url.Parse(tc.baseURL)
+			provBaseURL, err := url.Parse(tc.provider.URL)
 			assert.NoError(t, err)
 
 			pr := &httputil.ProxyRequest{