From 04ed44cf7ad16abfb74512b190f7bc35ba13c7f4 Mon Sep 17 00:00:00 2001
From: Adam Fisk <afisk@getlantern.org>
Date: Wed, 20 Nov 2024 11:04:44 -0700
Subject: [PATCH] Add more granular otel instrumentation

---
 fronted.go | 96 +++++++++++++++++++++++++++++++-----------------------
 1 file changed, 55 insertions(+), 41 deletions(-)

diff --git a/fronted.go b/fronted.go
index f208e8c..300127a 100644
--- a/fronted.go
+++ b/fronted.go
@@ -257,59 +257,70 @@ func (f *fronted) RoundTripHijack(req *http.Request) (*http.Response, net.Conn,
 			op.FailIf(err)
 			return nil, nil, err
 		}
-		provider := f.providerFor(m)
-		if provider == nil {
-			log.Debugf("Skipping masquerade with disabled/unknown provider '%s'", m.getProviderID())
-			masqueradeGood(false)
-			continue
-		}
-		frontedHost := provider.Lookup(originHost)
-		if frontedHost == "" {
-			// this error is not the masquerade's fault in particular
-			// so it is returned as good.
-			conn.Close()
-			masqueradeGood(true)
-			err := fmt.Errorf("no domain fronting mapping for '%s'. Please add it to provider_map.yaml or equivalent for %s",
-				m.getProviderID(), originHost)
-			op.FailIf(err)
-			return nil, nil, err
-		}
-		log.Debugf("Translated origin %s -> %s for provider %s...", originHost, frontedHost, m.getProviderID())
-
-		reqi, err := cloneRequestWith(req, frontedHost, getBody())
-		if err != nil {
-			return nil, nil, op.FailIf(log.Errorf("Failed to copy http request with origin translated to %v?: %v", frontedHost, err))
-		}
-
-		// don't clobber/confuse Connection header on Upgrade requests.
-		disableKeepAlives := true
-		if strings.EqualFold(reqi.Header.Get("Connection"), "upgrade") {
-			disableKeepAlives = false
-		}
-
-		tr := frontedHTTPTransport(conn, disableKeepAlives)
-		resp, err := tr.RoundTrip(reqi)
-		if err != nil {
-			log.Debugf("Could not complete request: %v", err)
-			masqueradeGood(false)
-			continue
-		}
 
-		err = provider.ValidateResponse(resp)
+		resp, conn, err := f.requestWithConn(req, conn, m, originHost, getBody, masqueradeGood)
 		if err != nil {
 			log.Debugf("Could not complete request: %v", err)
-			resp.Body.Close()
-			masqueradeGood(false)
 			continue
 		}
 
-		masqueradeGood(true)
 		return resp, conn, nil
 	}
 
 	return nil, nil, op.FailIf(errors.New("could not complete request even with retries"))
 }
 
+func (f *fronted) requestWithConn(req *http.Request, conn net.Conn, m MasqueradeInterface, originHost string, getBody func() io.ReadCloser, masqueradeGood func(bool) bool) (*http.Response, net.Conn, error) {
+	op := ops.Begin("request_with_conn")
+	defer op.End()
+	provider := f.providerFor(m)
+	if provider == nil {
+		log.Debugf("Skipping masquerade with disabled/unknown provider '%s'", m.getProviderID())
+		masqueradeGood(false)
+		return nil, nil, op.FailIf(log.Errorf("Skipping masquerade with disabled/unknown provider '%s'", m.getProviderID()))
+	}
+	frontedHost := provider.Lookup(originHost)
+	if frontedHost == "" {
+		// this error is not the masquerade's fault in particular
+		// so it is returned as good.
+		conn.Close()
+		masqueradeGood(true)
+		err := fmt.Errorf("no domain fronting mapping for '%s'. Please add it to provider_map.yaml or equivalent for %s",
+			m.getProviderID(), originHost)
+		op.FailIf(err)
+		return nil, nil, err
+	}
+	log.Debugf("Translated origin %s -> %s for provider %s...", originHost, frontedHost, m.getProviderID())
+
+	reqi, err := cloneRequestWith(req, frontedHost, getBody())
+	if err != nil {
+		return nil, nil, op.FailIf(log.Errorf("Failed to copy http request with origin translated to %v?: %v", frontedHost, err))
+	}
+	disableKeepAlives := true
+	if strings.EqualFold(reqi.Header.Get("Connection"), "upgrade") {
+		disableKeepAlives = false
+	}
+
+	tr := frontedHTTPTransport(conn, disableKeepAlives)
+	resp, err := tr.RoundTrip(reqi)
+	if err != nil {
+		log.Debugf("Could not complete request: %v", err)
+		masqueradeGood(false)
+		return nil, nil, err
+	}
+
+	err = provider.ValidateResponse(resp)
+	if err != nil {
+		log.Debugf("Could not complete request: %v", err)
+		resp.Body.Close()
+		masqueradeGood(false)
+		return nil, nil, err
+	}
+
+	masqueradeGood(true)
+	return resp, conn, nil
+}
+
 // Dial dials out using all available masquerades until one succeeds.
 func (f *fronted) dialAll(ctx context.Context) (net.Conn, MasqueradeInterface, func(bool) bool, error) {
 	conn, m, masqueradeGood, err := f.dialAllWith(ctx, f.masquerades)
@@ -317,6 +328,7 @@ func (f *fronted) dialAll(ctx context.Context) (net.Conn, MasqueradeInterface, f
 }
 
 func (f *fronted) dialAllWith(ctx context.Context, masquerades sortedMasquerades) (net.Conn, MasqueradeInterface, func(bool) bool, error) {
+	defer func(op ops.Op) { op.End() }(ops.Begin("dial_all_with"))
 	// never take more than a minute trying to find a dialer
 	ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
 	defer cancel()
@@ -491,6 +503,8 @@ type directTransport struct {
 }
 
 func (ddf *directTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
+	defer func(op ops.Op) { op.End() }(ops.Begin("direct_transport_roundtrip"))
+
 	// The connection is already encrypted by domain fronting.  We need to rewrite URLs starting
 	// with "https://" to "http://", lest we get an error for doubling up on TLS.