From eebbb566c464007573814ab83641aeab696bd70a Mon Sep 17 00:00:00 2001
From: Julien Pivotto <roidelapluie@o11y.eu>
Date: Fri, 8 Jul 2022 15:27:21 +0200
Subject: [PATCH] OAUth2: Respect disable keepalives option; Implement close
 idle connections

Signed-off-by: Julien Pivotto <roidelapluie@o11y.eu>
---
 config/http_config.go | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/config/http_config.go b/config/http_config.go
index 2ce312f6..be47b579 100644
--- a/config/http_config.go
+++ b/config/http_config.go
@@ -632,6 +632,7 @@ type oauth2RoundTripper struct {
 	secret string
 	mtx    sync.RWMutex
 	opts   *httpClientOptions
+	client *http.Client
 }
 
 func NewOAuth2RoundTripper(config *OAuth2, next http.RoundTripper, opts *httpClientOptions) http.RoundTripper {
@@ -680,14 +681,16 @@ func (rt *oauth2RoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
 		var t http.RoundTripper
 		if len(rt.config.TLSConfig.CAFile) == 0 {
 			t = &http.Transport{
-				TLSClientConfig: tlsConfig,
-				Proxy:           http.ProxyURL(rt.config.ProxyURL.URL),
+				TLSClientConfig:   tlsConfig,
+				Proxy:             http.ProxyURL(rt.config.ProxyURL.URL),
+				DisableKeepAlives: !rt.opts.keepAlivesEnabled,
 			}
 		} else {
 			t, err = NewTLSRoundTripper(tlsConfig, rt.config.TLSConfig.CAFile, func(tls *tls.Config) (http.RoundTripper, error) {
 				return &http.Transport{
-					TLSClientConfig: tls,
-					Proxy:           http.ProxyURL(rt.config.ProxyURL.URL),
+					TLSClientConfig:   tls,
+					Proxy:             http.ProxyURL(rt.config.ProxyURL.URL),
+					DisableKeepAlives: !rt.opts.keepAlivesEnabled,
 				}, nil
 			})
 			if err != nil {
@@ -699,7 +702,8 @@ func (rt *oauth2RoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
 			t = NewUserAgentRoundTripper(rt.opts.userAgent, t)
 		}
 
-		ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Transport: t})
+		client := &http.Client{Transport: t}
+		ctx := context.WithValue(context.Background(), oauth2.HTTPClient, client)
 		tokenSource := config.TokenSource(ctx)
 
 		rt.mtx.Lock()
@@ -708,6 +712,10 @@ func (rt *oauth2RoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
 			Base:   rt.next,
 			Source: tokenSource,
 		}
+		if rt.client != nil {
+			rt.client.CloseIdleConnections()
+		}
+		rt.client = client
 		rt.mtx.Unlock()
 	}
 
@@ -718,7 +726,9 @@ func (rt *oauth2RoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
 }
 
 func (rt *oauth2RoundTripper) CloseIdleConnections() {
-	// OAuth2 RT does not support CloseIdleConnections() but the next RT might.
+	if rt.client != nil {
+		rt.client.CloseIdleConnections()
+	}
 	if ci, ok := rt.next.(closeIdler); ok {
 		ci.CloseIdleConnections()
 	}