Populate correlation ID to gRPC requestsh-inject-and-log-correlation-di

Refactor common code into PrepareContext

6 files changed, 38 insertions, 20 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/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/handler/exec.go b/internal/handler/exec.go
index 5b87c2b..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,6 +58,20 @@ 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,
@@ -92,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)),
 		)
 	}