Merge branch 'id-define-standard-client' into 'main'

Define Do function for Gitlab net client

See merge request https://gitlab.com/gitlab-org/gitlab-shell/-/merge_requests/715

Merged-by: Ash McKenzie <amckenzie@gitlab.com>
Approved-by: Alejandro Rodríguez <alejandro@gitlab.com>
Approved-by: Ash McKenzie <amckenzie@gitlab.com>
Reviewed-by: Igor Drozdov <idrozdov@gitlab.com>
Reviewed-by: Ash McKenzie <amckenzie@gitlab.com>
Co-authored-by: Igor Drozdov <idrozdov@gitlab.com>

3 files changed, 77 insertions, 37 deletions

diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index 352196e..9794da3 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -12,8 +12,6 @@ import (
 
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/hashicorp/go-retryablehttp"
-
-	"gitlab.com/gitlab-org/labkit/log"
 )
 
 const (
@@ -107,7 +105,11 @@ func newRequest(ctx context.Context, method, host, path string, data interface{}
 	return request, nil
 }
 
-func parseError(resp *http.Response) error {
+func parseError(resp *http.Response, respErr error) error {
+	if respErr != nil {
+		return &ApiError{"Internal API unreachable"}
+	}
+
 	if resp.StatusCode >= 200 && resp.StatusCode <= 399 {
 		return nil
 	}
@@ -129,6 +131,15 @@ func (c *GitlabNetClient) Post(ctx context.Context, path string, data interface{
 	return c.DoRequest(ctx, http.MethodPost, normalizePath(path), data)
 }
 
+func (c *GitlabNetClient) Do(request *http.Request) (*http.Response, error) {
+	response, err := c.httpClient.RetryableHTTP.HTTPClient.Do(request)
+	if err := parseError(response, err); err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}
+
 func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, data interface{}) (*http.Response, error) {
 	request, err := newRequest(ctx, method, c.httpClient.Host, path, data)
 	if err != nil {
@@ -152,43 +163,13 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 	}
 	request.Header.Set(apiSecretHeaderName, tokenString)
 
-	originalRemoteIP, ok := ctx.Value(OriginalRemoteIPContextKey{}).(string)
-	if ok {
-		request.Header.Add("X-Forwarded-For", originalRemoteIP)
-	}
-
 	request.Header.Add("Content-Type", "application/json")
 	request.Header.Add("User-Agent", c.userAgent)
-	request.Close = true
-
-	start := time.Now()
 
 	response, err := c.httpClient.RetryableHTTP.Do(request)
-	fields := log.Fields{
-		"method":      method,
-		"url":         request.URL.String(),
-		"duration_ms": time.Since(start) / time.Millisecond,
-	}
-	logger := log.WithContextFields(ctx, fields)
-
-	if err != nil {
-		logger.WithError(err).Error("Internal API unreachable")
-		return nil, &ApiError{"Internal API unreachable"}
-	}
-
-	if response != nil {
-		logger = logger.WithField("status", response.StatusCode)
-	}
-	if err := parseError(response); err != nil {
-		logger.WithError(err).Error("Internal API error")
+	if err := parseError(response, err); err != nil {
 		return nil, err
 	}
 
-	if response.ContentLength >= 0 {
-		logger = logger.WithField("content_length_bytes", response.ContentLength)
-	}
-
-	logger.Info("Finished HTTP request")
-
 	return response, nil
 }
diff --git a/client/httpclient.go b/client/httpclient.go
index 9b57add..05ed4ef 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -14,8 +14,6 @@ import (
 	"time"
 
 	"github.com/hashicorp/go-retryablehttp"
-	"gitlab.com/gitlab-org/labkit/correlation"
-	"gitlab.com/gitlab-org/labkit/tracing"
 )
 
 const (
@@ -121,7 +119,7 @@ func NewHTTPClientWithOpts(gitlabURL, gitlabRelativeURLRoot, caFile, caPath stri
 	c.RetryWaitMax = hcc.retryWaitMax
 	c.RetryWaitMin = hcc.retryWaitMin
 	c.Logger = nil
-	c.HTTPClient.Transport = correlation.NewInstrumentedRoundTripper(tracing.NewRoundTripper(transport))
+	c.HTTPClient.Transport = newTransport(transport)
 	c.HTTPClient.Timeout = readTimeout(readTimeoutSeconds)
 
 	client := &HttpClient{RetryableHTTP: c, Host: host}
diff --git a/client/transport.go b/client/transport.go
new file mode 100644
index 0000000..e082cfa
--- /dev/null
+++ b/client/transport.go
@@ -0,0 +1,61 @@
+package client
+
+import (
+	"net/http"
+	"time"
+
+	"gitlab.com/gitlab-org/labkit/correlation"
+	"gitlab.com/gitlab-org/labkit/log"
+	"gitlab.com/gitlab-org/labkit/tracing"
+)
+
+type transport struct {
+	next http.RoundTripper
+}
+
+func (rt *transport) RoundTrip(request *http.Request) (*http.Response, error) {
+	ctx := request.Context()
+
+	originalRemoteIP, ok := ctx.Value(OriginalRemoteIPContextKey{}).(string)
+	if ok {
+		request.Header.Add("X-Forwarded-For", originalRemoteIP)
+	}
+	request.Close = true
+	request.Header.Add("User-Agent", defaultUserAgent)
+
+	start := time.Now()
+
+	response, err := rt.next.RoundTrip(request)
+
+	fields := log.Fields{
+		"method":      request.Method,
+		"url":         request.URL.String(),
+		"duration_ms": time.Since(start) / time.Millisecond,
+	}
+	logger := log.WithContextFields(ctx, fields)
+
+	if err != nil {
+		logger.WithError(err).Error("Internal API unreachable")
+		return response, err
+	}
+
+	logger = logger.WithField("status", response.StatusCode)
+
+	if response.StatusCode >= 400 {
+		logger.WithError(err).Error("Internal API error")
+		return response, err
+	}
+
+	if response.ContentLength >= 0 {
+		logger = logger.WithField("content_length_bytes", response.ContentLength)
+	}
+
+	logger.Info("Finished HTTP request")
+
+	return response, nil
+}
+
+func newTransport(next http.RoundTripper) http.RoundTripper {
+	t := &transport{next: next}
+	return correlation.NewInstrumentedRoundTripper(tracing.NewRoundTripper(t))
+}