diff options
author | Nick Thomas <nick@gitlab.com> | 2021-05-10 10:22:25 +0000 |
---|---|---|
committer | Nick Thomas <nick@gitlab.com> | 2021-05-10 10:22:25 +0000 |
commit | ef42fd615b58778e1d830d9c58691b39cb66e378 (patch) | |
tree | 878560b26c408d9bee14ed771f8c5c70e22ef1dd | |
parent | c6b47a56f553ee15175186dbc705c0d2ee59c0e1 (diff) | |
parent | d79d4777a88fcbf8f44771df76c4a6f1d3baa58b (diff) | |
download | gitlab-shell-ef42fd615b58778e1d830d9c58691b39cb66e378.tar.gz |
Merge branch '501-gitaly-respect-parent-context' into 'main'
Respect parent context for Gitaly calls
See merge request gitlab-org/gitlab-shell!469
-rw-r--r-- | internal/command/receivepack/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/command/receivepack/gitalycall_test.go | 7 | ||||
-rw-r--r-- | internal/command/receivepack/receivepack.go | 2 | ||||
-rw-r--r-- | internal/command/uploadarchive/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/command/uploadarchive/gitalycall_test.go | 6 | ||||
-rw-r--r-- | internal/command/uploadarchive/uploadarchive.go | 2 | ||||
-rw-r--r-- | internal/command/uploadpack/gitalycall.go | 4 | ||||
-rw-r--r-- | internal/command/uploadpack/gitalycall_test.go | 5 | ||||
-rw-r--r-- | internal/command/uploadpack/uploadpack.go | 2 | ||||
-rw-r--r-- | internal/gitlabnet/accessverifier/client.go | 3 | ||||
-rw-r--r-- | internal/handler/exec.go | 30 | ||||
-rw-r--r-- | internal/handler/exec_test.go | 8 | ||||
-rw-r--r-- | internal/sshd/sshd.go | 3 |
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 0b1c56b..723d655 100644 --- a/internal/handler/exec.go +++ b/internal/handler/exec.go @@ -44,16 +44,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 } @@ -62,12 +61,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 { @@ -83,10 +77,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, @@ -112,7 +106,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") } @@ -152,10 +146,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 } @@ -166,5 +160,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) |