fix

mtojek · mtojek · commit 94926bd8d0bb · 2024-06-07T14:37:58.000+02:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1010,7 +1010,7 @@ func TestAgent_FileTransferBlocked(t *testing.T) {
 		tempFile := filepath.Join(t.TempDir(), "scp")
 		err = scpClient.CopyFile(context.Background(), strings.NewReader("hello world"), tempFile, "0755")
 		require.Error(t, err)
-		require.Contains(t, err.Error(), agentssh.BlockedFileTransferErrorMessage)
+		assertFileTransferBlocked(t, err.Error())
 	})
 
 	t.Run("Forbidden commands", func(t *testing.T) {
@@ -1046,18 +1046,20 @@ func TestAgent_FileTransferBlocked(t *testing.T) {
 
 				msg, err := io.ReadAll(stdout)
 				require.NoError(t, err)
-				errorMessage := string(msg)
-
-				// NOTE: Checking content of the error message is flaky. It can catch: "File transfer has been disabled", "EOF", or "Process exited with status 2".
-				isErr := strings.Contains(errorMessage, agentssh.BlockedFileTransferErrorMessage) ||
-					strings.Contains(errorMessage, "EOF") ||
-					strings.Contains(errorMessage, "Process exited with status 2")
-				require.True(t, isErr, fmt.Sprintf("Message: "+errorMessage))
+				assertFileTransferBlocked(t, string(msg))
 			})
 		}
 	})
 }
 
+func assertFileTransferBlocked(t *testing.T, errorMessage string) {
+	// NOTE: Checking content of the error message is flaky. It can catch: "File transfer has been disabled", "EOF", or "Process exited with status 2".
+	isErr := strings.Contains(errorMessage, agentssh.BlockedFileTransferErrorMessage) ||
+		strings.Contains(errorMessage, "EOF") ||
+		strings.Contains(errorMessage, "Process exited with status 2")
+	require.True(t, isErr, fmt.Sprintf("Message: "+errorMessage))
+}
+
 func TestAgent_EnvironmentVariables(t *testing.T) {
 	t.Parallel()
 	key := "EXAMPLE"
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -279,9 +279,8 @@ func (s *Server) sessionHandler(session ssh.Session) {
 		extraEnv = append(extraEnv, fmt.Sprintf("DISPLAY=:%d.0", x11.ScreenNumber))
 	}
 
-	s.logger.Warn(ctx, "fileTransferBlocked", slog.F("session", session))
 	if s.fileTransferBlocked(session) {
-		s.logger.Warn(ctx, "fileTransferBlocked", slog.F("go ", "yes"))
+		s.logger.Warn(ctx, "file transfer blocked", slog.F("session_subsystem", session.Subsystem()), slog.F("raw_command", session.RawCommand()))
 
 		if session.Subsystem() == "" { // sftp does not expect error, otherwise it fails with "package too long"
 			// Response format: <status_code><message body>\n
@@ -291,7 +290,6 @@ func (s *Server) sessionHandler(session ssh.Session) {
 		_ = session.Exit(BlockedFileTransferErrorCode)
 		return
 	}
-	s.logger.Warn(ctx, "fileTransferBlocked end")
 
 	switch ss := session.Subsystem(); ss {
 	case "":
