diff options
-rw-r--r-- | client/client_test.go | 6 | ||||
-rw-r--r-- | internal/command/receivepack/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/command/receivepack/gitalycall_test.go | 1 | ||||
-rw-r--r-- | internal/command/uploadarchive/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/command/uploadpack/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/handler/exec.go | 39 |
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)), ) } |