Merge branch 'fix/gb/gitlab-api-reuse-http-connections' into 'master'

Ensure that we always close GitLab API response body

See merge request gitlab-org/gitlab-pages!275

2 files changed, 45 insertions, 0 deletions

diff --git a/internal/source/gitlab/client/client.go b/internal/source/gitlab/client/client.go
index bfe425b9..3a805ce9 100644
--- a/internal/source/gitlab/client/client.go
+++ b/internal/source/gitlab/client/client.go
@@ -85,6 +85,15 @@ func (gc *Client) GetLookup(ctx context.Context, host string) api.Lookup {
 		return api.Lookup{Name: host}
 	}
 
+	// ensure that entire response body has been read and close it, to make it
+	// possible to reuse HTTP connection. In case of a JSON being invalid and
+	// larger than 512 bytes, the response body will not be closed properly, thus
+	// we need to close it manually in every case.
+	defer func() {
+		io.Copy(ioutil.Discard, resp.Body)
+		resp.Body.Close()
+	}()
+
 	lookup := api.Lookup{Name: host}
 	lookup.Error = json.NewDecoder(resp.Body).Decode(&lookup.Domain)
 
diff --git a/internal/source/gitlab/client/client_test.go b/internal/source/gitlab/client/client_test.go
index 58729153..02684860 100644
--- a/internal/source/gitlab/client/client_test.go
+++ b/internal/source/gitlab/client/client_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
 	"testing"
 	"time"
 
@@ -20,6 +21,41 @@ const (
 	defaultJWTTokenExpiry    = 30 * time.Second
 )
 
+func TestConnectionReuse(t *testing.T) {
+	mux := http.NewServeMux()
+
+	mux.HandleFunc("/api/v4/internal/pages", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+
+		// we want to test for an invalid JSON that is larger than 512 bytes
+		b := make([]byte, 513)
+		for i := range b {
+			b[i] = 'x'
+		}
+
+		w.Write(b)
+	})
+
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	client := defaultClient(t, server.URL)
+	reused := make(chan bool, 2)
+
+	trace := &httptrace.ClientTrace{
+		GotConn: func(connInfo httptrace.GotConnInfo) {
+			reused <- connInfo.Reused
+		},
+	}
+
+	ctx := httptrace.WithClientTrace(context.Background(), trace)
+	client.GetLookup(ctx, "group.gitlab.io")
+	client.GetLookup(ctx, "group.gitlab.io")
+
+	require.False(t, <-reused)
+	require.True(t, <-reused)
+}
+
 func TestNewValidBaseURL(t *testing.T) {
 	_, err := NewClient("https://gitlab.com", secretKey(t), defaultClientConnTimeout, defaultJWTTokenExpiry)
 	require.NoError(t, err)