Respect parent context for Gitaly calls

Without these changes, Gitaly calls would not be linked to a parent
context. This means that they would have an unassociated correlationID,
and Gitaly RPC calls would not be cancel()ed by parent context
cancellation.

Changelog: fixed

13 files changed, 39 insertions, 41 deletions

diff --git a/internal/command/receivepack/gitalycall.go b/internal/command/receivepack/gitalycall.go
index 713f323..06e7b2c 100644
--- a/internal/command/receivepack/gitalycall.go
+++ b/internal/command/receivepack/gitalycall.go
@@ -12,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response) error {
+func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.ReceivePack),
@@ -30,7 +30,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
-	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
 		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
diff --git a/internal/command/receivepack/gitalycall_test.go b/internal/command/receivepack/gitalycall_test.go
index 9a1019d..70f2b2e 100644
--- a/internal/command/receivepack/gitalycall_test.go
+++ b/internal/command/receivepack/gitalycall_test.go
@@ -6,8 +6,8 @@ import (
 	"testing"
 
 	"github.com/sirupsen/logrus"
-
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/labkit/correlation"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
@@ -61,8 +61,9 @@ func TestReceivePack(t *testing.T) {
 		}
 
 		hook := testhelper.SetupLogger()
+		ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
 
-		err := cmd.Execute(context.Background())
+		err := cmd.Execute(ctx)
 		require.NoError(t, err)
 
 		if tc.username != "" {
@@ -80,6 +81,6 @@ 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=")
+		require.Contains(t, entries[1].Message, "correlation_id=a-correlation-id")
 	}
 }
diff --git a/internal/command/receivepack/receivepack.go b/internal/command/receivepack/receivepack.go
index 4d5c686..27e4b72 100644
--- a/internal/command/receivepack/receivepack.go
+++ b/internal/command/receivepack/receivepack.go
@@ -38,7 +38,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	return c.performGitalyCall(response)
+	return c.performGitalyCall(ctx, response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {
diff --git a/internal/command/uploadarchive/gitalycall.go b/internal/command/uploadarchive/gitalycall.go
index 05bb794..5b66bbb 100644
--- a/internal/command/uploadarchive/gitalycall.go
+++ b/internal/command/uploadarchive/gitalycall.go
@@ -12,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response) error {
+func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.UploadArchive),
@@ -23,7 +23,7 @@ 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) {
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
 		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
diff --git a/internal/command/uploadarchive/gitalycall_test.go b/internal/command/uploadarchive/gitalycall_test.go
index 03223e9..b89b749 100644
--- a/internal/command/uploadarchive/gitalycall_test.go
+++ b/internal/command/uploadarchive/gitalycall_test.go
@@ -6,8 +6,8 @@ import (
 	"testing"
 
 	"github.com/sirupsen/logrus"
-
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/labkit/correlation"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
@@ -36,8 +36,9 @@ func TestUploadPack(t *testing.T) {
 	}
 
 	hook := testhelper.SetupLogger()
+	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
 
-	err := cmd.Execute(context.Background())
+	err := cmd.Execute(ctx)
 	require.NoError(t, err)
 
 	require.Equal(t, "UploadArchive: "+repo, output.String())
@@ -50,4 +51,5 @@ func TestUploadPack(t *testing.T) {
 	require.Contains(t, entries[1].Message, "command=git-upload-archive")
 	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=a-correlation-id")
 }
diff --git a/internal/command/uploadarchive/uploadarchive.go b/internal/command/uploadarchive/uploadarchive.go
index 178b42b..3c86be3 100644
--- a/internal/command/uploadarchive/uploadarchive.go
+++ b/internal/command/uploadarchive/uploadarchive.go
@@ -28,7 +28,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		return err
 	}
 
-	return c.performGitalyCall(response)
+	return c.performGitalyCall(ctx, response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {
diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index 5a8a605..a263fc4 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -12,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response) error {
+func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.UploadPack),
@@ -27,7 +27,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
-	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
 		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
index d945764..bedc6e6 100644
--- a/internal/command/uploadpack/gitalycall_test.go
+++ b/internal/command/uploadpack/gitalycall_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/labkit/correlation"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
@@ -35,8 +36,9 @@ func TestUploadPack(t *testing.T) {
 	}
 
 	hook := testhelper.SetupLogger()
+	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
 
-	err := cmd.Execute(context.Background())
+	err := cmd.Execute(ctx)
 	require.NoError(t, err)
 
 	require.Equal(t, "UploadPack: "+repo, output.String())
@@ -48,6 +50,7 @@ func TestUploadPack(t *testing.T) {
 		require.Contains(t, entries[1].Message, "command=git-upload-pack")
 		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=a-correlation-id")
 		return true
 	}, time.Second, time.Millisecond)
 
diff --git a/internal/command/uploadpack/uploadpack.go b/internal/command/uploadpack/uploadpack.go
index fca3823..82e4d6e 100644
--- a/internal/command/uploadpack/uploadpack.go
+++ b/internal/command/uploadpack/uploadpack.go
@@ -38,7 +38,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	return c.performGitalyCall(response)
+	return c.performGitalyCall(ctx, response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {
diff --git a/internal/gitlabnet/accessverifier/client.go b/internal/gitlabnet/accessverifier/client.go
index 9389257..6272c8b 100644
--- a/internal/gitlabnet/accessverifier/client.go
+++ b/internal/gitlabnet/accessverifier/client.go
@@ -65,7 +65,6 @@ type Response struct {
 	ConsoleMessages  []string      `json:"gl_console_messages"`
 	Who              string
 	StatusCode       int
-	CorrelationID    string
 }
 
 func NewClient(config *config.Config) (*Client, error) {
@@ -110,8 +109,6 @@ 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 ac59dab..d11e1c3 100644
--- a/internal/handler/exec.go
+++ b/internal/handler/exec.go
@@ -45,16 +45,15 @@ type GitalyCommand struct {
 // RunGitalyCommand provides a bootstrap for Gitaly commands executed
 // through GitLab-Shell. It ensures that logging, tracing and other
 // common concerns are configured before executing the `handler`.
-func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error {
-	gitalyConn, err := getConn(gc)
-
+func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHandlerFunc) error {
+	gitalyConn, err := getConn(ctx, gc)
 	if err != nil {
 		return err
 	}
 
-	_, err = handler(gitalyConn.ctx, gitalyConn.conn)
+	defer gitalyConn.close()
 
-	gitalyConn.close()
+	_, err = handler(gitalyConn.ctx, gitalyConn.conn)
 
 	return err
 }
@@ -63,12 +62,7 @@ func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error {
 // be run.
 func (gc *GitalyCommand) PrepareContext(ctx context.Context, repository *pb.Repository, response *accessverifier.Response, env sshenv.Env) (context.Context, context.CancelFunc) {
 	ctx, cancel := context.WithCancel(ctx)
-
-	gc.LogExecution(repository, response, env)
-
-	if response.CorrelationID != "" {
-		ctx = correlation.ContextWithCorrelation(ctx, response.CorrelationID)
-	}
+	gc.LogExecution(ctx, repository, response, env)
 
 	md, ok := metadata.FromOutgoingContext(ctx)
 	if !ok {
@@ -84,10 +78,10 @@ func (gc *GitalyCommand) PrepareContext(ctx context.Context, repository *pb.Repo
 	return ctx, cancel
 }
 
-func (gc *GitalyCommand) LogExecution(repository *pb.Repository, response *accessverifier.Response, env sshenv.Env) {
+func (gc *GitalyCommand) LogExecution(ctx context.Context, repository *pb.Repository, response *accessverifier.Response, env sshenv.Env) {
 	fields := log.Fields{
 		"command":         gc.ServiceName,
-		"correlation_id":  response.CorrelationID,
+		"correlation_id":  correlation.ExtractFromContext(ctx),
 		"gl_project_path": repository.GlProjectPath,
 		"gl_repository":   repository.GlRepository,
 		"user_id":         response.UserId,
@@ -113,7 +107,7 @@ func withOutgoingMetadata(ctx context.Context, features map[string]string) conte
 	return metadata.NewOutgoingContext(ctx, md)
 }
 
-func getConn(gc *GitalyCommand) (*GitalyConn, error) {
+func getConn(ctx context.Context, gc *GitalyCommand) (*GitalyConn, error) {
 	if gc.Address == "" {
 		return nil, fmt.Errorf("no gitaly_address given")
 	}
@@ -158,10 +152,10 @@ func getConn(gc *GitalyCommand) (*GitalyConn, error) {
 		tracing.WithConnectionString(gc.Config.GitlabTracing),
 	)
 
-	ctx, finished := tracing.ExtractFromEnv(context.Background())
-	ctx = withOutgoingMetadata(ctx, gc.Features)
+	childCtx, finished := tracing.ExtractFromEnv(ctx)
+	childCtx = withOutgoingMetadata(childCtx, gc.Features)
 
-	conn, err := client.Dial(gc.Address, connOpts)
+	conn, err := client.DialContext(childCtx, gc.Address, connOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -172,5 +166,5 @@ func getConn(gc *GitalyCommand) (*GitalyConn, error) {
 		conn.Close()
 	}
 
-	return &GitalyConn{ctx: ctx, conn: conn, close: finish}, nil
+	return &GitalyConn{ctx: childCtx, conn: conn, close: finish}, nil
 }
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index 915bf5a..6f84709 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -30,18 +30,18 @@ func TestRunGitalyCommand(t *testing.T) {
 		Address: "tcp://localhost:9999",
 	}
 
-	err := cmd.RunGitalyCommand(makeHandler(t, nil))
+	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, nil))
 	require.NoError(t, err)
 
 	expectedErr := errors.New("error")
-	err = cmd.RunGitalyCommand(makeHandler(t, expectedErr))
+	err = cmd.RunGitalyCommand(context.Background(), makeHandler(t, expectedErr))
 	require.Equal(t, err, expectedErr)
 }
 
 func TestMissingGitalyAddress(t *testing.T) {
 	cmd := GitalyCommand{Config: &config.Config{}}
 
-	err := cmd.RunGitalyCommand(makeHandler(t, nil))
+	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, nil))
 	require.EqualError(t, err, "no gitaly_address given")
 }
 
@@ -70,7 +70,7 @@ func TestGetConnMetadata(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			conn, err := getConn(tt.gc)
+			conn, err := getConn(context.Background(), tt.gc)
 			require.NoError(t, err)
 
 			md, exists := metadata.FromOutgoingContext(conn.ctx)
diff --git a/internal/sshd/sshd.go b/internal/sshd/sshd.go
index a9e797b..470c069 100644
--- a/internal/sshd/sshd.go
+++ b/internal/sshd/sshd.go
@@ -17,6 +17,7 @@ import (
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/authorizedkeys"
+	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
 func Run(cfg *config.Config) error {
@@ -104,7 +105,7 @@ func handleConn(cfg *config.Config, sshCfg *ssh.ServerConfig, nconn net.Conn) {
 		}
 	}()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(correlation.ContextWithCorrelation(context.Background(), correlation.SafeRandomID()))
 	defer cancel()
 
 	sconn, chans, reqs, err := ssh.NewServerConn(nconn, sshCfg)