Merge branch 'sh-inject-and-log-correlation-di' into 'master'

Generate and log correlation IDs

See merge request gitlab-org/gitlab-shell!400

9 files changed, 67 insertions, 30 deletions

diff --git a/client/client_test.go b/client/client_test.go
index 5c5fda8..d520bbb 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -130,6 +130,7 @@ func testSuccessfulGet(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.Contains(t, entries[0].Message, "status=200")
 		assert.Contains(t, entries[0].Message, "Finished HTTP request")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -155,6 +156,7 @@ func testSuccessfulPost(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.Contains(t, entries[0].Message, "status=200")
 		assert.Contains(t, entries[0].Message, "Finished HTTP request")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -172,6 +174,7 @@ func testMissing(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.Contains(t, entries[0].Message, "status=404")
 		assert.Contains(t, entries[0].Message, "Internal API error")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 
 	t.Run("Missing error for POST", func(t *testing.T) {
@@ -187,6 +190,7 @@ func testMissing(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.Contains(t, entries[0].Message, "status=404")
 		assert.Contains(t, entries[0].Message, "Internal API error")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -219,6 +223,7 @@ func testBrokenRequest(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.NotContains(t, entries[0].Message, "status=")
 		assert.Contains(t, entries[0].Message, "Internal API unreachable")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 
 	t.Run("Broken request for POST", func(t *testing.T) {
@@ -235,6 +240,7 @@ func testBrokenRequest(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.NotContains(t, entries[0].Message, "status=")
 		assert.Contains(t, entries[0].Message, "Internal API unreachable")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index 7fbf63e..0657ca0 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -11,6 +12,7 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
 const (
@@ -57,23 +59,32 @@ func normalizePath(path string) string {
 	return path
 }
 
-func newRequest(method, host, path string, data interface{}) (*http.Request, error) {
+func newRequest(method, host, path string, data interface{}) (*http.Request, string, error) {
 	var jsonReader io.Reader
 	if data != nil {
 		jsonData, err := json.Marshal(data)
 		if err != nil {
-			return nil, err
+			return nil, "", err
 		}
 
 		jsonReader = bytes.NewReader(jsonData)
 	}
 
-	request, err := http.NewRequest(method, host+path, jsonReader)
+	correlationID, err := correlation.RandomID()
+	ctx := context.Background()
+
 	if err != nil {
-		return nil, err
+		log.WithError(err).Warn("unable to generate correlation ID")
+	} else {
+		ctx = correlation.ContextWithCorrelation(ctx, correlationID)
+	}
+
+	request, err := http.NewRequestWithContext(ctx, method, host+path, jsonReader)
+	if err != nil {
+		return nil, "", err
 	}
 
-	return request, nil
+	return request, correlationID, nil
 }
 
 func parseError(resp *http.Response) error {
@@ -100,7 +111,7 @@ func (c *GitlabNetClient) Post(path string, data interface{}) (*http.Response, e
 }
 
 func (c *GitlabNetClient) DoRequest(method, path string, data interface{}) (*http.Response, error) {
-	request, err := newRequest(method, c.httpClient.Host, path, data)
+	request, correlationID, err := newRequest(method, c.httpClient.Host, path, data)
 	if err != nil {
 		return nil, err
 	}
@@ -119,9 +130,10 @@ func (c *GitlabNetClient) DoRequest(method, path string, data interface{}) (*htt
 	start := time.Now()
 	response, err := c.httpClient.Do(request)
 	fields := log.Fields{
-		"method":      method,
-		"url":         request.URL.String(),
-		"duration_ms": time.Since(start) / time.Millisecond,
+		"correlation_id": correlationID,
+		"method":         method,
+		"url":            request.URL.String(),
+		"duration_ms":    time.Since(start) / time.Millisecond,
 	}
 	logger := log.WithFields(fields)
 
diff --git a/client/httpclient.go b/client/httpclient.go
index ff0cc25..63386f7 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -10,6 +10,8 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+
+	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
 const (
@@ -39,8 +41,9 @@ func NewHTTPClient(gitlabURL, caFile, caPath string, selfSignedCert bool, readTi
 		return nil
 	}
 
+
 	c := &http.Client{
-		Transport: transport,
+		Transport: correlation.NewInstrumentedRoundTripper(transport),
 		Timeout:   readTimeout(readTimeoutSeconds),
 	}
 
diff --git a/internal/command/receivepack/gitalycall.go b/internal/command/receivepack/gitalycall.go
index ffe0b6f..0754a3e 100644
--- a/internal/command/receivepack/gitalycall.go
+++ b/internal/command/receivepack/gitalycall.go
@@ -32,11 +32,9 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := context.WithCancel(ctx)
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, request.GitProtocol)
 		defer cancel()
 
-		gc.LogExecution(request.Repository, response, request.GitProtocol)
-
 		rw := c.ReadWriter
 		return client.ReceivePack(ctx, conn, rw.In, rw.Out, rw.ErrOut, request)
 	})
diff --git a/internal/command/receivepack/gitalycall_test.go b/internal/command/receivepack/gitalycall_test.go
index df922d9..8bee484 100644
--- a/internal/command/receivepack/gitalycall_test.go
+++ b/internal/command/receivepack/gitalycall_test.go
@@ -56,4 +56,5 @@ func TestReceivePack(t *testing.T) {
 	require.Contains(t, entries[1].Message, "remote_ip=127.0.0.1")
 	require.Contains(t, entries[1].Message, "gl_key_type=key")
 	require.Contains(t, entries[1].Message, "gl_key_id=123")
+	require.Contains(t, entries[1].Message, "correlation_id=")
 }
diff --git a/internal/command/uploadarchive/gitalycall.go b/internal/command/uploadarchive/gitalycall.go
index 1b792cb..f17ee50 100644
--- a/internal/command/uploadarchive/gitalycall.go
+++ b/internal/command/uploadarchive/gitalycall.go
@@ -24,11 +24,9 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	request := &pb.SSHUploadArchiveRequest{Repository: &response.Gitaly.Repo}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := context.WithCancel(ctx)
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, "")
 		defer cancel()
 
-		gc.LogExecution(request.Repository, response, "")
-
 		rw := c.ReadWriter
 		return client.UploadArchive(ctx, conn, rw.In, rw.Out, rw.ErrOut, request)
 	})
diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index c15ef38..ba0fef2 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -29,11 +29,9 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := context.WithCancel(ctx)
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, request.GitProtocol)
 		defer cancel()
 
-		gc.LogExecution(request.Repository, response, request.GitProtocol)
-
 		rw := c.ReadWriter
 		return client.UploadPack(ctx, conn, rw.In, rw.Out, rw.ErrOut, request)
 	})
diff --git a/internal/gitlabnet/accessverifier/client.go b/internal/gitlabnet/accessverifier/client.go
index 234e854..00b9d76 100644
--- a/internal/gitlabnet/accessverifier/client.go
+++ b/internal/gitlabnet/accessverifier/client.go
@@ -65,6 +65,7 @@ type Response struct {
 	ConsoleMessages  []string      `json:"gl_console_messages"`
 	Who              string
 	StatusCode       int
+	CorrelationID    string
 }
 
 func NewClient(config *config.Config) (*Client, error) {
@@ -109,6 +110,8 @@ func parse(hr *http.Response, args *commandargs.Shell) (*Response, error) {
 	}
 
 	response.StatusCode = hr.StatusCode
+	// We expect the same correlation ID that we sent
+	response.CorrelationID = hr.Header.Get("X-Request-Id")
 
 	return response, nil
 }
diff --git a/internal/handler/exec.go b/internal/handler/exec.go
index e4641b2..3688336 100644
--- a/internal/handler/exec.go
+++ b/internal/handler/exec.go
@@ -15,6 +15,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
+	"gitlab.com/gitlab-org/labkit/correlation"
 	grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
 	"gitlab.com/gitlab-org/labkit/tracing"
 	"google.golang.org/grpc"
@@ -57,9 +58,24 @@ func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error {
 	return err
 }
 
+// PrepareContext wraps a given context with a correlation ID and logs the command to
+// be run.
+func (gc *GitalyCommand) PrepareContext(ctx context.Context, repository *pb.Repository, response *accessverifier.Response, protocol string) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(ctx)
+
+	gc.LogExecution(repository, response, protocol)
+
+	if response.CorrelationID != "" {
+		ctx = correlation.ContextWithCorrelation(ctx, response.CorrelationID)
+	}
+
+	return ctx, cancel
+}
+
 func (gc *GitalyCommand) LogExecution(repository *pb.Repository, response *accessverifier.Response, protocol string) {
 	fields := log.Fields{
 		"command":         gc.ServiceName,
+		"correlation_id":  response.CorrelationID,
 		"gl_project_path": repository.GlProjectPath,
 		"gl_repository":   repository.GlRepository,
 		"user_id":         response.UserId,
@@ -91,19 +107,21 @@ func getConn(gc *GitalyCommand) (*GitalyConn, error) {
 	}
 
 	connOpts := client.DefaultDialOpts
-	if gc.Token != "" {
-		connOpts = append(client.DefaultDialOpts,
-			grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(gc.Token)),
-			grpc.WithStreamInterceptor(
-				grpccorrelation.StreamClientCorrelationInterceptor(
-					grpccorrelation.WithClientName(executable.GitlabShell),
-				),
+	connOpts = append(connOpts,
+		grpc.WithStreamInterceptor(
+			grpccorrelation.StreamClientCorrelationInterceptor(
+				grpccorrelation.WithClientName(executable.GitlabShell),
 			),
-			grpc.WithUnaryInterceptor(
-				grpccorrelation.UnaryClientCorrelationInterceptor(
-					grpccorrelation.WithClientName(executable.GitlabShell),
-				),
+		),
+		grpc.WithUnaryInterceptor(
+			grpccorrelation.UnaryClientCorrelationInterceptor(
+				grpccorrelation.WithClientName(executable.GitlabShell),
 			),
+		))
+
+	if gc.Token != "" {
+		connOpts = append(connOpts,
+			grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(gc.Token)),
 		)
 	}