From b6908181ff26d221975dabde254bfc464e4cf021 Mon Sep 17 00:00:00 2001
From: Mihai Parparita <mihai@tailscale.com>
Date: Fri, 3 Feb 2023 16:37:33 -0800
Subject: [PATCH] net/tshttpproxy: more directly use Transport proxy CONNECT
 hooks

GetProxyConnectHeader (golang/go#41048) was upstreamed in Go 1.16 and
OnProxyConnectResponse (golang/go#54299) in Go 1.20, thus we no longer
need to guard their use by the tailscale_go build tag.

Updates #7123

Signed-off-by: Mihai Parparita <mihai@tailscale.com>
---
 net/tshttpproxy/tshttpproxy.go        | 31 ++++++++++++++-----
 net/tshttpproxy/tshttpproxy_future.go | 44 ---------------------------
 2 files changed, 24 insertions(+), 51 deletions(-)
 delete mode 100644 net/tshttpproxy/tshttpproxy_future.go

diff --git a/net/tshttpproxy/tshttpproxy.go b/net/tshttpproxy/tshttpproxy.go
index 32f6e18ad..91864fc44 100644
--- a/net/tshttpproxy/tshttpproxy.go
+++ b/net/tshttpproxy/tshttpproxy.go
@@ -6,6 +6,9 @@
 package tshttpproxy
 
 import (
+	"context"
+	"fmt"
+	"log"
 	"net/http"
 	"net/url"
 	"os"
@@ -90,15 +93,29 @@ func GetAuthHeader(u *url.URL) (string, error) {
 	return "", nil
 }
 
-var condSetTransportGetProxyConnectHeader func(*http.Transport)
+const proxyAuthHeader = "Proxy-Authorization"
 
 // SetTransportGetProxyConnectHeader sets the provided Transport's
-// GetProxyConnectHeader field, if the current build of Go supports
-// it.
-//
-// See https://github.com/golang/go/issues/41048.
+// GetProxyConnectHeader field, and adds logging of the received response.
 func SetTransportGetProxyConnectHeader(tr *http.Transport) {
-	if f := condSetTransportGetProxyConnectHeader; f != nil {
-		f(tr)
+	tr.GetProxyConnectHeader = func(ctx context.Context, proxyURL *url.URL, target string) (http.Header, error) {
+		v, err := GetAuthHeader(proxyURL)
+		if err != nil {
+			log.Printf("failed to get proxy Auth header for %v; ignoring: %v", proxyURL, err)
+			return nil, nil
+		}
+		if v == "" {
+			return nil, nil
+		}
+		return http.Header{proxyAuthHeader: []string{v}}, nil
+	}
+	tr.OnProxyConnectResponse = func(ctx context.Context, proxyURL *url.URL, connectReq *http.Request, res *http.Response) error {
+		auth := connectReq.Header.Get(proxyAuthHeader)
+		const truncLen = 20
+		if len(auth) > truncLen {
+			auth = fmt.Sprintf("%s...(%d total bytes)", auth[:truncLen], len(auth))
+		}
+		log.Printf("tshttpproxy: CONNECT response from %v for target %q (auth %q): %v", proxyURL, connectReq.Host, auth, res.Status)
+		return nil
 	}
 }
diff --git a/net/tshttpproxy/tshttpproxy_future.go b/net/tshttpproxy/tshttpproxy_future.go
deleted file mode 100644
index 56c9eca15..000000000
--- a/net/tshttpproxy/tshttpproxy_future.go
+++ /dev/null
@@ -1,44 +0,0 @@
-// Copyright (c) Tailscale Inc & AUTHORS
-// SPDX-License-Identifier: BSD-3-Clause
-
-//go:build tailscale_go
-
-// We want to use https://github.com/golang/go/issues/41048 but it's only in the
-// Tailscale Go tree for now. Hence the build tag above.
-
-package tshttpproxy
-
-import (
-	"context"
-	"fmt"
-	"log"
-	"net/http"
-	"net/url"
-)
-
-const proxyAuthHeader = "Proxy-Authorization"
-
-func init() {
-	condSetTransportGetProxyConnectHeader = func(tr *http.Transport) {
-		tr.GetProxyConnectHeader = func(ctx context.Context, proxyURL *url.URL, target string) (http.Header, error) {
-			v, err := GetAuthHeader(proxyURL)
-			if err != nil {
-				log.Printf("failed to get proxy Auth header for %v; ignoring: %v", proxyURL, err)
-				return nil, nil
-			}
-			if v == "" {
-				return nil, nil
-			}
-			return http.Header{proxyAuthHeader: []string{v}}, nil
-		}
-		tr.OnProxyConnectResponse = func(ctx context.Context, proxyURL *url.URL, connectReq *http.Request, res *http.Response) error {
-			auth := connectReq.Header.Get(proxyAuthHeader)
-			const truncLen = 20
-			if len(auth) > truncLen {
-				auth = fmt.Sprintf("%s...(%d total bytes)", auth[:truncLen], len(auth))
-			}
-			log.Printf("tshttpproxy: CONNECT response from %v for target %q (auth %q): %v", proxyURL, connectReq.Host, auth, res.Status)
-			return nil
-		}
-	}
-}