Fix

mtojek · mtojek · commit fa0010d90508 · 2024-06-07T14:05:05.000+02:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -973,6 +973,24 @@ func TestAgent_SCP(t *testing.T) {
 func TestAgent_FileTransferBlocked(t *testing.T) {
 	t.Parallel()
 
+	t.Run("SFTP", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		//nolint:dogsled
+		conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
+			o.BlockFileTransfer = true
+		})
+		sshClient, err := conn.SSHClient(ctx)
+		require.NoError(t, err)
+		defer sshClient.Close()
+		_, err = sftp.NewClient(sshClient)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "unexpected EOF")
+	})
+
 	t.Run("SCP with go-scp package", func(t *testing.T) {
 		t.Parallel()
 
@@ -1026,9 +1044,14 @@ func TestAgent_FileTransferBlocked(t *testing.T) {
 				require.NoError(t, err)
 				defer session.Close()
 
-				errorMessage, err := io.ReadAll(stdout)
+				msg, err := io.ReadAll(stdout)
 				require.NoError(t, err)
-				require.Contains(t, string(errorMessage), agentssh.BlockedFileTransferErrorMessage)
+				errorMessage := string(msg)
+
+				// NOTE: Checking content of the error message is flaky. Sometimes it catches "EOF" or "Process terminate with status code 2".
+				isErr := strings.Contains(errorMessage, agentssh.BlockedFileTransferErrorMessage) ||
+					strings.Contains(errorMessage, "EOF")
+				require.True(t, isErr, fmt.Sprintf("Message: "+errorMessage))
 			})
 		}
 	})
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -279,13 +279,19 @@ 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) {
-		// Response format: <status_code><message body>\n
-		errorMessage := fmt.Sprintf("\x02%s\n", BlockedFileTransferErrorMessage)
-		_, _ = session.Write([]byte(errorMessage))
+		s.logger.Warn(ctx, "fileTransferBlocked", slog.F("go ", "yes"))
+
+		if session.Subsystem() == "" { // sftp does not expect error, otherwise it fails with "package too long"
+			// Response format: <status_code><message body>\n
+			errorMessage := fmt.Sprintf("\x02%s\n", BlockedFileTransferErrorMessage)
+			_, _ = session.Write([]byte(errorMessage))
+		}
 		_ = session.Exit(BlockedFileTransferErrorCode)
 		return
 	}
+	s.logger.Warn(ctx, "fileTransferBlocked end")
 
 	switch ss := session.Subsystem(); ss {
 	case "":
@@ -338,8 +344,6 @@ func (s *Server) sessionHandler(session ssh.Session) {
 }
 
 // fileTransferBlocked method checks if the file transfer commands should be blocked.
-// It does not block SFTP sessions, VS Code may still use this protocol.
-//
 // Warning: consider this mechanism as "Do not trespass" sign. If a user needs a more sophisticated
 // and battle-proof solution, consider the full endpoint security.
 func (s *Server) fileTransferBlocked(session ssh.Session) bool {
@@ -348,6 +352,10 @@ func (s *Server) fileTransferBlocked(session ssh.Session) bool {
 	}
 	// File transfers are restricted.
 
+	if session.Subsystem() == "sftp" {
+		return true
+	}
+
 	cmd := session.Command()
 	if len(cmd) == 0 {
 		return false // no command?
