Display constistently in gitlab-sshd and gitlab-shell

- Use console package to format the errors in gitlab-sshd
- Suppress internal Gitaly errors in client output

4 files changed, 31 insertions, 18 deletions

diff --git a/internal/handler/exec.go b/internal/handler/exec.go
index 0870d2b..3d22284 100644
--- a/internal/handler/exec.go
+++ b/internal/handler/exec.go
@@ -58,13 +58,11 @@ func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHan
 	exitStatus, err := handler(childCtx, conn)
 
 	if err != nil {
-		if grpcstatus.Convert(err).Code() == grpccodes.Unavailable {
-			ctxlog.WithError(fmt.Errorf("RunGitalyCommand: %v", err)).Error("Gitaly is unavailable")
+		ctxlog.WithError(err).WithFields(log.Fields{"exit_status": exitStatus}).Error("Failed to execute Git command")
 
-			return fmt.Errorf("The git server, Gitaly, is not available at this time. Please contact your administrator.")
+		if grpcstatus.Code(err) == grpccodes.Unavailable {
+			return grpcstatus.Error(grpccodes.Unavailable, "The git server, Gitaly, is not available at this time. Please contact your administrator.")
 		}
-
-		ctxlog.WithError(err).WithFields(log.Fields{"exit_status": exitStatus}).Error("Failed to execute Git command")
 	}
 
 	return err
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index 791fe79..ccec88f 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -84,10 +84,8 @@ func TestUnavailableGitalyErr(t *testing.T) {
 		},
 	)
 
-	expectedErr := grpcstatus.Error(grpccodes.Unavailable, "error")
-	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, expectedErr))
-
-	require.EqualError(t, err, "The git server, Gitaly, is not available at this time. Please contact your administrator.")
+	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, grpcstatus.Error(grpccodes.Unavailable, "error")))
+	require.Equal(t, err, grpcstatus.Error(grpccodes.Unavailable, "The git server, Gitaly, is not available at this time. Please contact your administrator."))
 }
 
 func TestRunGitalyCommandMetadata(t *testing.T) {
diff --git a/internal/sshd/session.go b/internal/sshd/session.go
index 831beb8..ac079dc 100644
--- a/internal/sshd/session.go
+++ b/internal/sshd/session.go
@@ -8,11 +8,14 @@ import (
 
 	"gitlab.com/gitlab-org/labkit/log"
 	"golang.org/x/crypto/ssh"
+	grpccodes "google.golang.org/grpc/codes"
+	grpcstatus "google.golang.org/grpc/status"
 
 	shellCmd "gitlab.com/gitlab-org/gitlab-shell/cmd/gitlab-shell/command"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/shared/disallowedcommand"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/console"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
 
@@ -158,10 +161,12 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (uint32, er
 
 	cmd, err := shellCmd.NewWithKey(s.gitlabKeyId, env, s.cfg, rw)
 	if err != nil {
-		if !errors.Is(err, disallowedcommand.Error) {
-			s.toStderr(ctx, "Failed to parse command: %v\n", err.Error())
+		if errors.Is(err, disallowedcommand.Error) {
+			s.toStderr(ctx, "ERROR: Unknown command: %v\n", s.execCmd)
+		} else {
+			s.toStderr(ctx, "ERROR: Failed to parse command: %v\n", err.Error())
 		}
-		s.toStderr(ctx, "Unknown command: %v\n", s.execCmd)
+
 		return 128, err
 	}
 
@@ -169,7 +174,11 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (uint32, er
 	ctxlog.WithFields(log.Fields{"env": env, "command": cmdName}).Info("session: handleShell: executing command")
 
 	if err := cmd.Execute(ctx); err != nil {
-		s.toStderr(ctx, "remote: ERROR: %v\n", err.Error())
+		grpcStatus := grpcstatus.Convert(err)
+		if grpcStatus.Code() != grpccodes.Internal {
+			s.toStderr(ctx, "ERROR: %v\n", grpcStatus.Message())
+		}
+
 		return 1, err
 	}
 
@@ -181,7 +190,7 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (uint32, er
 func (s *session) toStderr(ctx context.Context, format string, args ...interface{}) {
 	out := fmt.Sprintf(format, args...)
 	log.WithContextFields(ctx, log.Fields{"stderr": out}).Debug("session: toStderr: output")
-	fmt.Fprint(s.channel.Stderr(), out)
+	console.DisplayWarningMessage(out, s.channel.Stderr())
 }
 
 func (s *session) exit(ctx context.Context, status uint32) {
diff --git a/internal/sshd/session_test.go b/internal/sshd/session_test.go
index 5bc8e7c..6090696 100644
--- a/internal/sshd/session_test.go
+++ b/internal/sshd/session_test.go
@@ -13,6 +13,7 @@ import (
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/console"
 )
 
 type fakeChannel struct {
@@ -160,14 +161,14 @@ func TestHandleShell(t *testing.T) {
 		{
 			desc:              "fails to parse command",
 			cmd:               `\`,
-			errMsg:            "Failed to parse command: Invalid SSH command: invalid command line string\nUnknown command: \\\n",
+			errMsg:            "ERROR: Failed to parse command: Invalid SSH command: invalid command line string\n",
 			gitlabKeyId:       "root",
 			expectedErrString: "Invalid SSH command: invalid command line string",
 			expectedExitCode:  128,
 		}, {
 			desc:              "specified command is unknown",
 			cmd:               "unknown-command",
-			errMsg:            "Unknown command: unknown-command\n",
+			errMsg:            "ERROR: Unknown command: unknown-command\n",
 			gitlabKeyId:       "root",
 			expectedErrString: "Disallowed command",
 			expectedExitCode:  128,
@@ -175,7 +176,7 @@ func TestHandleShell(t *testing.T) {
 			desc:              "fails to parse command",
 			cmd:               "discover",
 			gitlabKeyId:       "",
-			errMsg:            "remote: ERROR: Failed to get username: who='' is invalid\n",
+			errMsg:            "ERROR: Failed to get username: who='' is invalid\n",
 			expectedErrString: "Failed to get username: who='' is invalid",
 			expectedExitCode:  1,
 		}, {
@@ -208,7 +209,14 @@ func TestHandleShell(t *testing.T) {
 			}
 
 			require.Equal(t, tc.expectedExitCode, exitCode)
-			require.Equal(t, tc.errMsg, out.String())
+
+			formattedErr := &bytes.Buffer{}
+			if tc.errMsg != "" {
+				console.DisplayWarningMessage(tc.errMsg, formattedErr)
+				require.Equal(t, formattedErr.String(), out.String())
+			} else {
+				require.Equal(t, tc.errMsg, out.String())
+			}
 		})
 	}
 }