Pass ctx where needed

cmd/check/main.go

@@ -43,7 +43,7 @@ func main() {
 	ctx, finished := command.Setup(executable.Name, config)
 	defer finished()
 
-	if err = cmd.Execute(ctx); err != nil {
+	if ctx, err = cmd.Execute(ctx); err != nil {
 		fmt.Fprintf(readWriter.ErrOut, "%v\n", err)
 		os.Exit(1)
 	}

cmd/gitlab-shell-authorized-keys-check/main.go

@@ -46,7 +46,7 @@ func main() {
 	ctx, finished := command.Setup(executable.Name, config)
 	defer finished()
 
-	if err = cmd.Execute(ctx); err != nil {
+	if ctx, err = cmd.Execute(ctx); err != nil {
 		console.DisplayWarningMessage(err.Error(), readWriter.ErrOut)
 		os.Exit(1)
 	}

cmd/gitlab-shell-authorized-principals-check/main.go

@@ -46,7 +46,7 @@ func main() {
 	ctx, finished := command.Setup(executable.Name, config)
 	defer finished()
 
-	if err = cmd.Execute(ctx); err != nil {
+	if ctx, err = cmd.Execute(ctx); err != nil {
 		console.DisplayWarningMessage(err.Error(), readWriter.ErrOut)
 		os.Exit(1)
 	}

cmd/gitlab-shell/main.go

@@ -76,7 +76,7 @@ func main() {
 	ctxlog.WithFields(log.Fields{"env": env, "command": cmdName}).Info("gitlab-shell: main: executing command")
 	fips.Check()
 
-	if err := cmd.Execute(ctx); err != nil {
+	if _, err := cmd.Execute(ctx); err != nil {
 		ctxlog.WithError(err).Warn("gitlab-shell: main: command execution failed")
 		if grpcstatus.Convert(err).Code() != grpccodes.Internal {
 			console.DisplayWarningMessage(err.Error(), readWriter.ErrOut)

internal/command/authorizedkeys/authorized_keys.go

@@ -18,21 +18,21 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	// Do and return nothing when the expected and actual user don't match.
 	// This can happen when the user in sshd_config doesn't match the user
 	// trying to login. When nothing is printed, the user will be denied access.
 	if c.Args.ExpectedUser != c.Args.ActualUser {
 		// TODO: Log this event once we have a consistent way to log in Go.
 		// See https://gitlab.com/gitlab-org/gitlab-shell/issues/192 for more info.
-		return nil
+		return ctx, nil
 	}
 
 	if err := c.printKeyLine(ctx); err != nil {
-		return err
+		return ctx, err
 	}
 
-	return nil
+	return ctx, nil
 }
 
 func (c *Command) printKeyLine(ctx context.Context) error {

internal/command/authorizedkeys/authorized_keys_test.go

@@ -84,7 +84,7 @@ func TestExecute(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{Out: buffer},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedOutput, buffer.String())

internal/command/authorizedprincipals/authorized_principals.go

@@ -16,12 +16,12 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	if err := c.printPrincipalLines(); err != nil {
-		return err
+		return ctx, err
 	}
 
-	return nil
+	return ctx, nil
 }
 
 func (c *Command) printPrincipalLines() error {

internal/command/authorizedprincipals/authorized_principals_test.go

@@ -42,7 +42,7 @@ func TestExecute(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{Out: buffer},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedOutput, buffer.String())

internal/command/command.go

@@ -9,7 +9,7 @@ import (
 )
 
 type Command interface {
-	Execute(ctx context.Context) error
+	Execute(ctx context.Context) (context.Context, error)
 }
 
 // Setup() initializes tracing from the configuration file and generates a

internal/command/discover/discover.go

@@ -16,10 +16,10 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	response, err := c.getUserInfo(ctx)
 	if err != nil {
-		return fmt.Errorf("Failed to get username: %v", err)
+		return ctx, fmt.Errorf("Failed to get username: %v", err)
 	}
 
 	if response.IsAnonymous() {
@@ -28,7 +28,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		fmt.Fprintf(c.ReadWriter.Out, "Welcome to GitLab, @%s!\n", response.Username)
 	}
 
-	return nil
+	return ctx, nil
 }
 
 func (c *Command) getUserInfo(ctx context.Context) (*discover.Response, error) {

internal/command/discover/discover_test.go

@@ -81,7 +81,7 @@ func TestExecute(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{Out: buffer},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedOutput, buffer.String())
@@ -123,7 +123,7 @@ func TestFailingExecute(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{Out: buffer},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 
 			require.Empty(t, buffer.String())
 			require.EqualError(t, err, tc.expectedError)

internal/command/healthcheck/healthcheck.go

@@ -19,20 +19,20 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	response, err := c.runCheck(ctx)
 	if err != nil {
-		return fmt.Errorf("%v: FAILED - %v", apiMessage, err)
+		return ctx, fmt.Errorf("%v: FAILED - %v", apiMessage, err)
 	}
 
 	fmt.Fprintf(c.ReadWriter.Out, "%v: OK\n", apiMessage)
 
 	if !response.Redis {
-		return fmt.Errorf("%v: FAILED", redisMessage)
+		return ctx, fmt.Errorf("%v: FAILED", redisMessage)
 	}
 
 	fmt.Fprintf(c.ReadWriter.Out, "%v: OK\n", redisMessage)
-	return nil
+	return ctx, nil
 }
 
 func (c *Command) runCheck(ctx context.Context) (*healthcheck.Response, error) {

internal/command/healthcheck/healthcheck_test.go

@@ -53,7 +53,7 @@ func TestExecute(t *testing.T) {
 		ReadWriter: &readwriter.ReadWriter{Out: buffer},
 	}
 
-	err := cmd.Execute(context.Background())
+	_, err := cmd.Execute(context.Background())
 
 	require.NoError(t, err)
 	require.Equal(t, "Internal API available: OK\nRedis available via internal API: OK\n", buffer.String())
@@ -68,7 +68,7 @@ func TestFailingRedisExecute(t *testing.T) {
 		ReadWriter: &readwriter.ReadWriter{Out: buffer},
 	}
 
-	err := cmd.Execute(context.Background())
+	_, err := cmd.Execute(context.Background())
 	require.Error(t, err, "Redis available via internal API: FAILED")
 	require.Equal(t, "Internal API available: OK\n", buffer.String())
 }
@@ -82,7 +82,7 @@ func TestFailingAPIExecute(t *testing.T) {
 		ReadWriter: &readwriter.ReadWriter{Out: buffer},
 	}
 
-	err := cmd.Execute(context.Background())
+	_, err := cmd.Execute(context.Background())
 	require.Empty(t, buffer.String())
 	require.EqualError(t, err, "Internal API available: FAILED - Internal API unreachable")
 }

internal/command/lfsauthenticate/lfsauthenticate.go

@@ -37,10 +37,10 @@ type Payload struct {
 	ExpiresIn int           `json:"expires_in,omitempty"`
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	args := c.Args.SshArgs
 	if len(args) < 3 {
-		return disallowedcommand.Error
+		return ctx, disallowedcommand.Error
 	}
 
 	// e.g. git-lfs-authenticate user/repo.git download
@@ -49,12 +49,12 @@ func (c *Command) Execute(ctx context.Context) error {
 
 	action, err := actionFromOperation(operation)
 	if err != nil {
-		return err
+		return ctx, err
 	}
 
 	accessResponse, err := c.verifyAccess(ctx, action, repo)
 	if err != nil {
-		return err
+		return ctx, err
 	}
 
 	payload, err := c.authenticate(ctx, operation, repo, accessResponse.UserId)
@@ -65,12 +65,12 @@ func (c *Command) Execute(ctx context.Context) error {
 			log.Fields{"operation": operation, "repo": repo, "user_id": accessResponse.UserId},
 		).WithError(err).Debug("lfsauthenticate: execute: LFS authentication failed")
 
-		return nil
+		return ctx, nil
 	}
 
 	fmt.Fprintf(c.ReadWriter.Out, "%s\n", payload)
 
-	return nil
+	return ctx, nil
 }
 
 func actionFromOperation(operation string) (commandargs.CommandType, error) {

internal/command/lfsauthenticate/lfsauthenticate_test.go

@@ -54,7 +54,7 @@ func TestFailedRequests(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 			require.Error(t, err)
 
 			require.Equal(t, tc.expectedOutput, err.Error())
@@ -145,7 +145,7 @@ func TestLfsAuthenticateRequests(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 			require.NoError(t, err)
 
 			require.Equal(t, tc.expectedOutput, output.String())

internal/command/personalaccesstoken/personalaccesstoken.go

@@ -34,10 +34,10 @@ type tokenArgs struct {
 	ExpiresDate string // Calculated, a TTL is passed from command-line.
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	err := c.parseTokenArgs()
 	if err != nil {
-		return err
+		return ctx, err
 	}
 
 	log.WithContextFields(ctx, log.Fields{
@@ -46,13 +46,14 @@ func (c *Command) Execute(ctx context.Context) error {
 
 	response, err := c.getPersonalAccessToken(ctx)
 	if err != nil {
-		return err
+		return ctx, err
 	}
 
 	fmt.Fprint(c.ReadWriter.Out, "Token:   "+response.Token+"\n")
 	fmt.Fprint(c.ReadWriter.Out, "Scopes:  "+strings.Join(response.Scopes, ",")+"\n")
 	fmt.Fprint(c.ReadWriter.Out, "Expires: "+response.ExpiresAt+"\n")
-	return nil
+
+	return ctx, nil
 }
 
 func (c *Command) parseTokenArgs() error {

internal/command/personalaccesstoken/personalaccesstoken_test.go

@@ -167,7 +167,7 @@ func TestExecute(t *testing.T) {
 				ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
 			}
 
-			err := cmd.Execute(context.Background())
+			_, err := cmd.Execute(context.Background())
 
 			if tc.expectedError == "" {
 				require.NoError(t, err)

internal/command/receivepack/gitalycall_test.go

@@ -72,7 +72,7 @@ func TestReceivePack(t *testing.T) {
 				ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
 				ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
-				err := cmd.Execute(ctx)
+				_, err := cmd.Execute(ctx)
 				require.NoError(t, err)
 
 				if tc.username != "" {

internal/command/receivepack/receivepack.go

@@ -18,16 +18,16 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
-func (c *Command) Execute(ctx context.Context) error {
+func (c *Command) Execute(ctx context.Context) (context.Context, error) {
 	args := c.Args.SshArgs
 	if len(args) != 2 {
-		return disallowedcommand.Error
+		return ctx, disallowedcommand.Error
 	}
 
 	repo := args[1]
 	response, err := c.verifyAccess(ctx, repo)
 	if err != nil {
-		return err
+		return ctx, err
 	}
 
 	if response.IsCustomAction() {
@@ -42,7 +42,7 @@ func (c *Command) Execute(ctx context.Context) error {
 				Response:   response,
 			}
 
-			return cmd.Execute(ctx)
+			return ctx, cmd.Execute(ctx)
 		}
 
 		customAction := customaction.Command{
@@ -50,10 +50,10 @@ func (c *Command) Execute(ctx context.Context) error {
 			ReadWriter: c.ReadWriter,
 			EOFSent:    true,
 		}
-		return customAction.Execute(ctx, response)
+		return ctx, customAction.Execute(ctx, response)
 	}
 
-	return c.performGitalyCall(ctx, response)
+	return ctx, c.performGitalyCall(ctx, response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {

internal/command/receivepack/receivepack_test.go

@@ -18,14 +18,15 @@ func TestForbiddenAccess(t *testing.T) {
 	requests := requesthandlers.BuildDisallowedByApiHandlers(t)
 	cmd, _ := setup(t, "disallowed", requests)
 
-	err := cmd.Execute(context.Background())
+	_, err := cmd.Execute(context.Background())
 	require.Equal(t, "Disallowed by API call", err.Error())
 }
 
 func TestCustomReceivePack(t *testing.T) {
 	cmd, output := setup(t, "1", requesthandlers.BuildAllowedWithCustomActionsHandlers(t))
 
-	require.NoError(t, cmd.Execute(context.Background()))
+	_, err := cmd.Execute(context.Background())
+	require.NoError(t, err)
 	require.Equal(t, "customoutput", output.String())
 }