From 2052676927842ce543fcf6268e194441d651bb17 Mon Sep 17 00:00:00 2001
From: Alan Hamlett <alan.hamlett@gmail.com>
Date: Wed, 19 May 2021 20:10:36 -0700
Subject: [PATCH] Control http.Client creation instead of using
 http.DefaultClient

---
 cmd/legacy/heartbeat/heartbeat.go |  3 +--
 cmd/legacy/today/today.go         |  3 +--
 cmd/legacy/todaygoal/todaygoal.go |  3 +--
 go.mod                            |  2 +-
 go.sum                            |  2 ++
 pkg/api/api.go                    | 13 ++++++++++---
 pkg/api/goal_test.go              | 10 +++++-----
 pkg/api/heartbeat_test.go         |  8 ++++----
 pkg/api/option_test.go            | 23 +++++++----------------
 pkg/api/summary_test.go           | 12 ++++++------
 pkg/api/transport.go              | 18 ++++++++++++++++++
 11 files changed, 56 insertions(+), 41 deletions(-)
 create mode 100644 pkg/api/transport.go

diff --git a/cmd/legacy/heartbeat/heartbeat.go b/cmd/legacy/heartbeat/heartbeat.go
index 0320708d..ee87e83f 100644
--- a/cmd/legacy/heartbeat/heartbeat.go
+++ b/cmd/legacy/heartbeat/heartbeat.go
@@ -3,7 +3,6 @@ package heartbeat
 import (
 	"errors"
 	"fmt"
-	"net/http"
 	"os"
 	"strings"
 
@@ -130,7 +129,7 @@ func SendHeartbeats(v *viper.Viper, queueFilepath string) error {
 		clientOpts = append(clientOpts, api.WithUserAgentUnknownPlugin())
 	}
 
-	c := api.NewClient(params.API.URL, http.DefaultClient, clientOpts...)
+	c := api.NewClient(params.API.URL, clientOpts...)
 
 	heartbeats := []heartbeat.Heartbeat{
 		heartbeat.New(
diff --git a/cmd/legacy/today/today.go b/cmd/legacy/today/today.go
index 41e01aff..47556211 100644
--- a/cmd/legacy/today/today.go
+++ b/cmd/legacy/today/today.go
@@ -3,7 +3,6 @@ package today
 import (
 	"errors"
 	"fmt"
-	"net/http"
 	"os"
 	"strings"
 	"time"
@@ -123,7 +122,7 @@ func Summary(v *viper.Viper) (string, error) {
 		url = params.API.URL
 	}
 
-	c := api.NewClient(url, http.DefaultClient, opts...)
+	c := api.NewClient(url, opts...)
 
 	now := time.Now()
 	todayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
diff --git a/cmd/legacy/todaygoal/todaygoal.go b/cmd/legacy/todaygoal/todaygoal.go
index 6fa2403d..d9e8f039 100644
--- a/cmd/legacy/todaygoal/todaygoal.go
+++ b/cmd/legacy/todaygoal/todaygoal.go
@@ -3,7 +3,6 @@ package todaygoal
 import (
 	"errors"
 	"fmt"
-	"net/http"
 	"os"
 	"regexp"
 	"strings"
@@ -126,7 +125,7 @@ func Goal(v *viper.Viper) (string, error) {
 		url = params.API.URL
 	}
 
-	c := api.NewClient(url, http.DefaultClient, opts...)
+	c := api.NewClient(url, opts...)
 
 	goal, err := c.Goal(params.GoalID)
 	if err != nil {
diff --git a/go.mod b/go.mod
index ecada400..caa0185c 100644
--- a/go.mod
+++ b/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/sirupsen/logrus v1.8.1
 	github.com/slongfield/pyfmt v0.0.0-20180124071345-020a7cb18bca
 	github.com/spf13/cobra v1.1.1
-	github.com/spf13/jwalterweatherman v1.0.0
+	github.com/spf13/jwalterweatherman v1.1.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.7.1
 	github.com/stretchr/testify v1.7.0
diff --git a/go.sum b/go.sum
index 7f0528a6..0f3d397f 100644
--- a/go.sum
+++ b/go.sum
@@ -194,6 +194,8 @@ github.com/spf13/cobra v1.1.1 h1:KfztREH0tPxJJ+geloSLaAkaPkr4ki2Er5quFV1TDo4=
 github.com/spf13/cobra v1.1.1/go.mod h1:WnodtKOvamDL/PwE2M4iKs8aMDBZ5Q5klgD3qfVJQMI=
 github.com/spf13/jwalterweatherman v1.0.0 h1:XHEdyB+EcvlqZamSM4ZOMGlc93t6AcsBEu9Gc1vn7yk=
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
+github.com/spf13/jwalterweatherman v1.1.0 h1:ue6voC5bR5F8YxI5S67j9i582FU4Qvo2bmqnqMYADFk=
+github.com/spf13/jwalterweatherman v1.1.0/go.mod h1:aNWZUN0dPAAO/Ljvb5BEdw96iTZ0EXowPYD95IqWIGo=
 github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
diff --git a/pkg/api/api.go b/pkg/api/api.go
index c26a05a9..8ac565aa 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -1,6 +1,9 @@
 package api
 
-import "net/http"
+import (
+	"net/http"
+	"time"
+)
 
 // BaseURL is the base url of the wakatime api.
 const BaseURL = "https://api.wakatime.com/api"
@@ -25,10 +28,14 @@ type Client struct {
 }
 
 // NewClient creates a new Client. Any number of Options can be provided.
-func NewClient(baseURL string, client *http.Client, opts ...Option) *Client {
+func NewClient(baseURL string, opts ...Option) *Client {
 	c := &Client{
 		baseURL: baseURL,
-		client:  client,
+		client: &http.Client{
+			Transport:     NewTransport(),
+			CheckRedirect: nil, // defaults to following up to 10 redirects
+			Timeout:       30 * time.Second,
+		},
 		doFunc: func(c *Client, req *http.Request) (*http.Response, error) {
 			req.Header.Set("Accept", "application/json")
 			return c.client.Do(req)
diff --git a/pkg/api/goal_test.go b/pkg/api/goal_test.go
index d55cf522..4b696085 100644
--- a/pkg/api/goal_test.go
+++ b/pkg/api/goal_test.go
@@ -39,7 +39,7 @@ func TestClient_Goal(t *testing.T) {
 			require.NoError(t, err)
 		})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	goal, err := c.Goal("00000000-0000-4000-8000-000000000000")
 
 	require.NoError(t, err)
@@ -65,7 +65,7 @@ func TestClient_GoalWithTimeout(t *testing.T) {
 		})
 
 	opts := []api.Option{api.WithTimeout(20 * time.Millisecond)}
-	c := api.NewClient(u, http.DefaultClient, opts...)
+	c := api.NewClient(u, opts...)
 	_, err := c.Goal("00000000-0000-4000-8000-000000000000")
 	require.Error(t, err)
 
@@ -93,7 +93,7 @@ func TestClient_Goal_Err(t *testing.T) {
 			w.WriteHeader(http.StatusInternalServerError)
 		})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	_, err := c.Goal("00000000-0000-4000-8000-000000000000")
 
 	var apierr api.Err
@@ -114,7 +114,7 @@ func TestClient_Goal_ErrAuth(t *testing.T) {
 			w.WriteHeader(http.StatusUnauthorized)
 		})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	_, err := c.Goal("00000000-0000-4000-8000-000000000000")
 
 	var autherr api.ErrAuth
@@ -124,7 +124,7 @@ func TestClient_Goal_ErrAuth(t *testing.T) {
 }
 
 func TestClient_Goal_ErrRequest(t *testing.T) {
-	c := api.NewClient("invalid-url", http.DefaultClient)
+	c := api.NewClient("invalid-url")
 	_, err := c.Goal("00000000-0000-4000-8000-000000000000")
 
 	var reqerr api.ErrRequest
diff --git a/pkg/api/heartbeat_test.go b/pkg/api/heartbeat_test.go
index 8e0a1e65..a8d2263f 100644
--- a/pkg/api/heartbeat_test.go
+++ b/pkg/api/heartbeat_test.go
@@ -55,7 +55,7 @@ func TestClient_Send(t *testing.T) {
 				require.NoError(t, err)
 			})
 
-			c := api.NewClient(url, http.DefaultClient)
+			c := api.NewClient(url)
 			results, err := c.Send(testHeartbeats())
 			require.NoError(t, err)
 
@@ -115,7 +115,7 @@ func TestClient_Send_Err(t *testing.T) {
 		w.WriteHeader(http.StatusInternalServerError)
 	})
 
-	c := api.NewClient(url, http.DefaultClient)
+	c := api.NewClient(url)
 	_, err := c.Send(testHeartbeats())
 
 	var errapi api.Err
@@ -136,7 +136,7 @@ func TestClient_Send_ErrAuth(t *testing.T) {
 		w.WriteHeader(http.StatusUnauthorized)
 	})
 
-	c := api.NewClient(url, http.DefaultClient)
+	c := api.NewClient(url)
 	_, err := c.Send(testHeartbeats())
 
 	var errauth api.ErrAuth
@@ -147,7 +147,7 @@ func TestClient_Send_ErrAuth(t *testing.T) {
 }
 
 func TestClient_Send_ErrRequest(t *testing.T) {
-	c := api.NewClient("invalid-url", http.DefaultClient)
+	c := api.NewClient("invalid-url")
 	_, err := c.Send(testHeartbeats())
 
 	var errreq api.ErrRequest
diff --git a/pkg/api/option_test.go b/pkg/api/option_test.go
index 9bf5b46e..48c6d244 100644
--- a/pkg/api/option_test.go
+++ b/pkg/api/option_test.go
@@ -53,7 +53,7 @@ func TestOption_WithAuth(t *testing.T) {
 			req, err := http.NewRequest(http.MethodGet, url, nil)
 			require.NoError(t, err)
 
-			c := api.NewClient("", httpClient(), []api.Option{withAuth}...)
+			c := api.NewClient("", []api.Option{withAuth}...)
 			resp, err := c.Do(req)
 			require.NoError(t, err)
 
@@ -81,7 +81,7 @@ func TestOption_WithHostname(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	require.NoError(t, err)
 
-	c := api.NewClient("", httpClient(), opts...)
+	c := api.NewClient("", opts...)
 	resp, err := c.Do(req)
 	require.NoError(t, err)
 
@@ -129,7 +129,7 @@ func TestOption_WithNTLM(t *testing.T) {
 			req, err := http.NewRequest(http.MethodGet, url, nil)
 			require.NoError(t, err)
 
-			c := api.NewClient("", &http.Client{}, []api.Option{withNTLM}...)
+			c := api.NewClient("", []api.Option{withNTLM}...)
 			resp, err := c.Do(req)
 			require.NoError(t, err)
 
@@ -186,7 +186,7 @@ func TestOption_WithNTLMRequestRetry(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	require.NoError(t, err)
 
-	c := api.NewClient("", &http.Client{}, []api.Option{withNTLMRetry}...)
+	c := api.NewClient("", []api.Option{withNTLMRetry}...)
 	resp, err := c.Do(req)
 	require.NoError(t, err)
 
@@ -213,7 +213,7 @@ func TestOption_WithProxy(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, "http://example.org", nil)
 	require.NoError(t, err)
 
-	c := api.NewClient("", httpClient(), opts...)
+	c := api.NewClient("", opts...)
 	resp, err := c.Do(req)
 	require.NoError(t, err)
 
@@ -248,7 +248,7 @@ func TestOption_WithUserAgent(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	require.NoError(t, err)
 
-	c := api.NewClient("", httpClient(), opts...)
+	c := api.NewClient("", opts...)
 	resp, err := c.Do(req)
 	require.NoError(t, err)
 
@@ -283,7 +283,7 @@ func TestOption_WithUserAgentUnknownPlugin(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	require.NoError(t, err)
 
-	c := api.NewClient("", httpClient(), opts...)
+	c := api.NewClient("", opts...)
 	resp, err := c.Do(req)
 	require.NoError(t, err)
 
@@ -291,12 +291,3 @@ func TestOption_WithUserAgentUnknownPlugin(t *testing.T) {
 
 	assert.Eventually(t, func() bool { return numCalls == 1 }, time.Second, 50*time.Millisecond)
 }
-
-func httpClient() *http.Client {
-	transport := http.DefaultTransport.(*http.Transport).Clone()
-	transport.DisableKeepAlives = true
-
-	return &http.Client{
-		Transport: transport,
-	}
-}
diff --git a/pkg/api/summary_test.go b/pkg/api/summary_test.go
index 6ba28da4..1d8bf5ea 100644
--- a/pkg/api/summary_test.go
+++ b/pkg/api/summary_test.go
@@ -49,7 +49,7 @@ func TestClient_Summaries(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	summaries, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
@@ -87,7 +87,7 @@ func TestClient_SummariesByCategory(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	summaries, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
@@ -138,7 +138,7 @@ func TestClient_SummariesWithTimeout(t *testing.T) {
 	})
 
 	opts := []api.Option{api.WithTimeout(20 * time.Millisecond)}
-	c := api.NewClient(u, http.DefaultClient, opts...)
+	c := api.NewClient(u, opts...)
 	_, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
@@ -168,7 +168,7 @@ func TestClient_Summaries_Err(t *testing.T) {
 		w.WriteHeader(http.StatusInternalServerError)
 	})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	_, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
@@ -191,7 +191,7 @@ func TestClient_Summaries_ErrAuth(t *testing.T) {
 		w.WriteHeader(http.StatusUnauthorized)
 	})
 
-	c := api.NewClient(u, http.DefaultClient)
+	c := api.NewClient(u)
 	_, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
@@ -204,7 +204,7 @@ func TestClient_Summaries_ErrAuth(t *testing.T) {
 }
 
 func TestClient_Summaries_ErrRequest(t *testing.T) {
-	c := api.NewClient("invalid-url", http.DefaultClient)
+	c := api.NewClient("invalid-url")
 	_, err := c.Summaries(
 		time.Date(2020, time.April, 1, 0, 0, 0, 0, time.UTC),
 		time.Date(2020, time.April, 2, 0, 0, 0, 0, time.UTC),
diff --git a/pkg/api/transport.go b/pkg/api/transport.go
new file mode 100644
index 00000000..44ed3bfc
--- /dev/null
+++ b/pkg/api/transport.go
@@ -0,0 +1,18 @@
+package api
+
+import (
+	"net/http"
+	"time"
+)
+
+// NewTransport creates a new http.Transport, with default 30 second TLS timeout.
+func NewTransport() *http.Transport {
+	return &http.Transport{
+		Proxy:               nil,
+		TLSHandshakeTimeout: 30 * time.Second,
+		MaxIdleConns:        1,
+		MaxIdleConnsPerHost: 1,
+		MaxConnsPerHost:     1,
+		ForceAttemptHTTP2:   true,
+	}
+}