From 6679a1ec31696dbf3acebc5385440bfdd8b3cf55 Mon Sep 17 00:00:00 2001
From: Josh Deprez <jd@buildkite.com>
Date: Mon, 23 Sep 2024 15:33:09 +1000
Subject: [PATCH] Attempt to workaround golang/go#59690

---
 api/auth.go   | 40 ++++++++++++++++++++------
 api/client.go | 79 +++++++++++++++++++++++++++------------------------
 2 files changed, 73 insertions(+), 46 deletions(-)

diff --git a/api/auth.go b/api/auth.go
index 1fb28da103..274d92a2d6 100644
--- a/api/auth.go
+++ b/api/auth.go
@@ -5,11 +5,7 @@ import (
 	"net/http"
 )
 
-type canceler interface {
-	CancelRequest(*http.Request)
-}
-
-// authenticatedTransport manages injection of the API token
+// authenticatedTransport manages injection of the API token.
 type authenticatedTransport struct {
 	// The Token used for authentication. This can either the be
 	// organizations registration token, or the agents access token.
@@ -19,19 +15,45 @@ type authenticatedTransport struct {
 	Delegate http.RoundTripper
 }
 
-// RoundTrip invoked each time a request is made
+// RoundTrip invoked each time a request is made.
 func (t authenticatedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if t.Token == "" {
 		return nil, fmt.Errorf("Invalid token, empty string supplied")
 	}
 
+	// Per net/http#RoundTripper:
+	//
+	// "RoundTrip should not modify the request, except for
+	// consuming and closing the Request's Body. RoundTrip may
+	// read fields of the request in a separate goroutine. Callers
+	// should not mutate or reuse the request until the Response's
+	// Body has been closed."
+	//
+	// But we can pass a _different_ request to t.Delegate.RoundTrip.
+	// req.Clone does a sufficiently deep clone (including Header which we
+	// modify).
+	req = req.Clone(req.Context())
 	req.Header.Set("Authorization", fmt.Sprintf("Token %s", t.Token))
 
 	return t.Delegate.RoundTrip(req)
 }
 
-// CancelRequest cancels an in-flight request by closing its connection.
+// CancelRequest forwards the call to t.Delegate, if it implements CancelRequest
+// itself.
 func (t *authenticatedTransport) CancelRequest(req *http.Request) {
-	cancelableTransport := t.Delegate.(canceler)
-	cancelableTransport.CancelRequest(req)
+	canceler, ok := t.Delegate.(interface{ CancelRequest(*http.Request) })
+	if !ok {
+		return
+	}
+	canceler.CancelRequest(req)
+}
+
+// CloseIdleConnections forwards the call to t.Delegate, if it implements
+// CloseIdleConnections itself.
+func (t *authenticatedTransport) CloseIdleConnections() {
+	closer, ok := t.Delegate.(interface{ CloseIdleConnections() })
+	if !ok {
+		return
+	}
+	closer.CloseIdleConnections()
 }
diff --git a/api/client.go b/api/client.go
index ef58ec814b..1b37b9c741 100644
--- a/api/client.go
+++ b/api/client.go
@@ -10,7 +10,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"net"
 	"net/http"
 	"net/http/httptrace"
 	"net/http/httputil"
@@ -83,44 +82,50 @@ func NewClient(l logger.Logger, conf Config) *Client {
 
 	httpClient := conf.HTTPClient
 
-	if conf.HTTPClient == nil {
-		if conf.DisableHTTP2 {
-			tr := &http.Transport{
-				Proxy:              http.ProxyFromEnvironment,
-				DisableCompression: false,
-				DisableKeepAlives:  false,
-				DialContext: (&net.Dialer{
-					Timeout:   30 * time.Second,
-					KeepAlive: 30 * time.Second,
-				}).DialContext,
-				MaxIdleConns:        100,
-				IdleConnTimeout:     90 * time.Second,
-				TLSHandshakeTimeout: 30 * time.Second,
-				TLSNextProto:        make(map[string]func(authority string, c *tls.Conn) http.RoundTripper),
-			}
-			httpClient = &http.Client{
-				Timeout: 60 * time.Second,
-				Transport: &authenticatedTransport{
-					Token:    conf.Token,
-					Delegate: tr,
-				},
-			}
-		} else {
-			tr := &http2.Transport{
-				ReadIdleTimeout: 30 * time.Second,
-			}
-			if conf.TLSConfig != nil {
-				tr.TLSClientConfig = conf.TLSConfig
-			}
+	if httpClient != nil {
+		return &Client{
+			logger: l,
+			client: httpClient,
+			conf:   conf,
+		}
+	}
 
-			httpClient = &http.Client{
-				Timeout: 60 * time.Second,
-				Transport: &authenticatedTransport{
-					Token:    conf.Token,
-					Delegate: tr,
-				},
-			}
+	var transport http.RoundTripper
+	if conf.DisableHTTP2 {
+ 		tr := http.DefaultTransport.(*http.Transport).Clone()
+		tr.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
+		// The default TLSClientConfig has h2 in NextProtos, so the
+		// negotiated TLS connection will assume h2 support.
+		// see https://github.com/golang/go/issues/50571
+		tr.TLSClientConfig.NextProtos = []string{"http/1.1"}
+
+		// Allow override of TLSConfig
+		if conf.TLSConfig != nil {
+			tr.TLSClientConfig = conf.TLSConfig
+ 		}
+
+		transport = tr
+	} else {
+		// There is still a bug in http2 around reusing dead connections.
+		// This is a workaround. See https://github.com/golang/go/issues/59690
+		tr := &http2.Transport{
+			ReadIdleTimeout: 30 * time.Second,
+ 		}
+
+		// Allow override of TLSConfig
+		if conf.TLSConfig != nil {
+			tr.TLSClientConfig = conf.TLSConfig
 		}
+
+		transport = tr
+	}
+
+	httpClient = &http.Client{
+		Timeout: 60 * time.Second,
+		Transport: &authenticatedTransport{
+			Token:    conf.Token,
+			Delegate: transport,
+		},
 	}
 
 	return &Client{