Merge branch '648-gitlab-sshd-should-include-data-transfer-bytes-in-logs-3' into 'main'

Resolve "GitLab sshd should include data transfer bytes in logs"

Closes #648

See merge request https://gitlab.com/gitlab-org/gitlab-shell/-/merge_requests/831



Merged-by: Igor Drozdov <idrozdov@gitlab.com>
Approved-by: Igor Drozdov <idrozdov@gitlab.com>
Reviewed-by: Igor Drozdov <idrozdov@gitlab.com>
Reviewed-by: Ash McKenzie <amckenzie@gitlab.com>
Co-authored-by: Ash McKenzie <amckenzie@gitlab.com>

cmd/check/main.go

@@ -23,7 +23,7 @@ func main() {
 	command.CheckForVersionFlag(os.Args, Version, BuildTime)
 
 	readWriter := &readwriter.ReadWriter{
-		Out:    os.Stdout,
+		Out:    &readwriter.CountingWriter{W: os.Stdout},
 		In:     os.Stdin,
 		ErrOut: os.Stderr,
 	}

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

@@ -24,7 +24,7 @@ func main() {
 	command.CheckForVersionFlag(os.Args, Version, BuildTime)
 
 	readWriter := &readwriter.ReadWriter{
-		Out:    os.Stdout,
+		Out:    &readwriter.CountingWriter{W: os.Stdout},
 		In:     os.Stdin,
 		ErrOut: os.Stderr,
 	}

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

@@ -24,7 +24,7 @@ func main() {
 	command.CheckForVersionFlag(os.Args, Version, BuildTime)
 
 	readWriter := &readwriter.ReadWriter{
-		Out:    os.Stdout,
+		Out:    &readwriter.CountingWriter{W: os.Stdout},
 		In:     os.Stdin,
 		ErrOut: os.Stderr,
 	}

cmd/gitlab-shell/main.go

@@ -32,7 +32,7 @@ func main() {
 	command.CheckForVersionFlag(os.Args, Version, BuildTime)
 
 	readWriter := &readwriter.ReadWriter{
-		Out:    os.Stdout,
+		Out:    &readwriter.CountingWriter{W: os.Stdout},
 		In:     os.Stdin,
 		ErrOut: os.Stderr,
 	}

internal/command/command.go

@@ -22,8 +22,9 @@ type LogMetadata struct {
 }
 
 type LogData struct {
-	Username string      `json:"username"`
-	Meta     LogMetadata `json:"meta"`
+	Username     string      `json:"username"`
+	WrittenBytes int64       `json:"written_bytes"`
+	Meta         LogMetadata `json:"meta"`
 }
 
 func CheckForVersionFlag(osArgs []string, version, buildTime string) {
@@ -87,7 +88,8 @@ func NewLogData(project, username string) LogData {
 	}
 
 	return LogData{
-		Username: username,
+		Username:     username,
+		WrittenBytes: 0,
 		Meta: LogMetadata{
 			Project:       project,
 			RootNamespace: rootNameSpace,

internal/command/readwriter/readwriter.go

 package readwriter
 
-import "io"
+import (
+	"io"
+)
 
 type ReadWriter struct {
 	Out    io.Writer
 	In     io.Reader
 	ErrOut io.Writer
 }
+
+// CountingWriter wraps an io.Writer and counts all the writes. Accessing
+// the count N is not thread-safe.
+type CountingWriter struct {
+	W io.Writer
+	N int64
+}
+
+func (cw *CountingWriter) Write(p []byte) (int, error) {
+	n, err := cw.W.Write(p)
+	cw.N += int64(n)
+	return n, err
+}

internal/command/readwriter/readwriter_test.go

+package readwriter
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCountingWriter_Write(t *testing.T) {
+	testString := []byte("test string")
+	buffer := &bytes.Buffer{}
+
+	cw := &CountingWriter{
+		W: buffer,
+	}
+
+	n, err := cw.Write(testString)
+
+	require.NoError(t, err)
+	require.Equal(t, 11, n)
+	require.Equal(t, int64(11), cw.N)
+
+	cw.Write(testString)
+	require.Equal(t, int64(22), cw.N)
+}

internal/sshd/session.go

@@ -167,8 +167,10 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (context.Co
 		NamespacePath:      s.namespace,
 	}
 
+	countingWriter := &readwriter.CountingWriter{W: s.channel}
+
 	rw := &readwriter.ReadWriter{
-		Out:    s.channel,
+		Out:    countingWriter,
 		In:     s.channel,
 		ErrOut: s.channel.Stderr(),
 	}
@@ -183,6 +185,7 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (context.Co
 	} else {
 		cmd, err = shellCmd.NewWithKey(s.gitlabKeyId, env, s.cfg, rw)
 	}
+
 	if err != nil {
 		if errors.Is(err, disallowedcommand.Error) {
 			s.toStderr(ctx, "ERROR: Unknown command: %v\n", s.execCmd)
@@ -202,6 +205,12 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) (context.Co
 	metrics.SshdSessionEstablishedDuration.Observe(establishSessionDuration)
 
 	ctxWithLogData, err := cmd.Execute(ctx)
+
+	logData := extractDataFromContext(ctxWithLogData)
+	logData.WrittenBytes = countingWriter.N
+
+	ctxWithLogData = context.WithValue(ctx, "logData", logData)
+
 	if err != nil {
 		grpcStatus := grpcstatus.Convert(err)
 		if grpcStatus.Code() != grpccodes.Internal {

internal/sshd/session_test.go

@@ -18,6 +18,7 @@ import (
 
 type fakeChannel struct {
 	stdErr             io.ReadWriter
+	stdOut             io.ReadWriter
 	sentRequestName    string
 	sentRequestPayload []byte
 }
@@ -27,7 +28,7 @@ func (f *fakeChannel) Read(data []byte) (int, error) {
 }
 
 func (f *fakeChannel) Write(data []byte) (int, error) {
-	return 0, nil
+	return f.stdOut.Write(data)
 }
 
 func (f *fakeChannel) Close() error {
@@ -145,8 +146,9 @@ func TestHandleExec(t *testing.T) {
 				},
 			}
 			for _, s := range sessions {
-				out := &bytes.Buffer{}
-				f := &fakeChannel{stdErr: out}
+				stdErr := &bytes.Buffer{}
+				stdOut := &bytes.Buffer{}
+				f := &fakeChannel{stdErr: stdErr, stdOut: stdOut}
 				r := &ssh.Request{Payload: tc.payload}
 
 				s.channel = f
@@ -163,12 +165,14 @@ func TestHandleExec(t *testing.T) {
 
 func TestHandleShell(t *testing.T) {
 	testCases := []struct {
-		desc              string
-		cmd               string
-		errMsg            string
-		gitlabKeyId       string
-		expectedErrString string
-		expectedExitCode  uint32
+		desc                 string
+		cmd                  string
+		errMsg               string
+		gitlabKeyId          string
+		expectedOutString    string
+		expectedErrString    string
+		expectedExitCode     uint32
+		expectedWrittenBytes int64
 	}{
 		{
 			desc:              "fails to parse command",
@@ -177,27 +181,32 @@ func TestHandleShell(t *testing.T) {
 			gitlabKeyId:       "root",
 			expectedErrString: "Invalid SSH command: invalid command line string",
 			expectedExitCode:  128,
-		}, {
+		},
+		{
 			desc:              "specified command is unknown",
 			cmd:               "unknown-command",
 			errMsg:            "ERROR: Unknown command: unknown-command\n",
 			gitlabKeyId:       "root",
 			expectedErrString: "Disallowed command",
 			expectedExitCode:  128,
-		}, {
+		},
+		{
 			desc:              "fails to parse command",
 			cmd:               "discover",
 			gitlabKeyId:       "",
 			errMsg:            "ERROR: Failed to get username: who='' is invalid\n",
 			expectedErrString: "Failed to get username: who='' is invalid",
 			expectedExitCode:  1,
-		}, {
-			desc:              "fails to parse command",
-			cmd:               "discover",
-			errMsg:            "",
-			gitlabKeyId:       "root",
-			expectedErrString: "",
-			expectedExitCode:  0,
+		},
+		{
+			desc:                 "parses command",
+			cmd:                  "discover",
+			errMsg:               "",
+			gitlabKeyId:          "root",
+			expectedOutString:    "Welcome to GitLab, @test-user!\n",
+			expectedErrString:    "",
+			expectedExitCode:     0,
+			expectedWrittenBytes: 31,
 		},
 	}
 
@@ -205,29 +214,37 @@ func TestHandleShell(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			out := &bytes.Buffer{}
+			stdOut := &bytes.Buffer{}
+			stdErr := &bytes.Buffer{}
 			s := &session{
 				gitlabKeyId: tc.gitlabKeyId,
 				execCmd:     tc.cmd,
-				channel:     &fakeChannel{stdErr: out},
+				channel:     &fakeChannel{stdErr: stdErr, stdOut: stdOut},
 				cfg:         &config.Config{GitlabUrl: url},
 			}
 			r := &ssh.Request{}
 
-			_, exitCode, err := s.handleShell(context.Background(), r)
+			ctxWithLogData, exitCode, err := s.handleShell(context.Background(), r)
+
+			logData := extractDataFromContext(ctxWithLogData)
+
+			if tc.expectedOutString != "" {
+				require.Equal(t, tc.expectedOutString, stdOut.String())
+			}
 
 			if tc.expectedErrString != "" {
 				require.Equal(t, tc.expectedErrString, err.Error())
 			}
 
 			require.Equal(t, tc.expectedExitCode, exitCode)
+			require.Equal(t, tc.expectedWrittenBytes, logData.WrittenBytes)
 
 			formattedErr := &bytes.Buffer{}
 			if tc.errMsg != "" {
 				console.DisplayWarningMessage(tc.errMsg, formattedErr)
-				require.Equal(t, formattedErr.String(), out.String())
+				require.Equal(t, formattedErr.String(), stdErr.String())
 			} else {
-				require.Equal(t, tc.errMsg, out.String())
+				require.Equal(t, tc.errMsg, stdErr.String())
 			}
 		})
 	}

internal/sshd/sshd.go

@@ -217,8 +217,9 @@ func (s *Server) handleConn(ctx context.Context, nconn net.Conn) {
 	logData := extractDataFromContext(ctxWithLogData)
 
 	ctxlog.WithFields(log.Fields{
-		"duration_s": time.Since(started).Seconds(),
-		"meta":       logData.Meta,
+		"duration_s":    time.Since(started).Seconds(),
+		"written_bytes": logData.WrittenBytes,
+		"meta":          logData.Meta,
 	}).Info("access: finish")
 }