Set logData.WrittenBytes from CountingWriter.N

Author: ashmckenzie

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:    &readwriter.CountingWriter{W: 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,57 +181,70 @@ 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,
 		},
 	}
 
 	url := testserver.StartHttpServer(t, requests)
 
 	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())
 			}
 		})
 	}