From c649d0ea3238139f1add055fc719f413f398b175 Mon Sep 17 00:00:00 2001
From: f-bor <42858062+f-bor@users.noreply.github.com>
Date: Fri, 26 Oct 2018 06:19:17 +0200
Subject: [PATCH] enhance recv calls in network_cli (#47345)

* enhance recv calls in network_cli

* updated network_cli test unit

* enhance recv calls in network_cli

* fix mistake

* better timeout management

* remove exception trigger

* test

* test2

* restore exception and timeout

* ganeshrn's way

* correction

* timeout and exception return
---
 lib/ansible/plugins/connection/network_cli.py    | 16 ++++++++++++++--
 .../units/plugins/connection/test_network_cli.py |  2 ++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/lib/ansible/plugins/connection/network_cli.py b/lib/ansible/plugins/connection/network_cli.py
index 18dbd25104..f6a65f6b8f 100644
--- a/lib/ansible/plugins/connection/network_cli.py
+++ b/lib/ansible/plugins/connection/network_cli.py
@@ -182,6 +182,7 @@ import os
 import signal
 import socket
 import traceback
+import time
 from io import BytesIO
 
 from ansible.errors import AnsibleConnectionFailure
@@ -359,6 +360,15 @@ class Connection(NetworkConnectionBase):
                 display.debug("ssh connection has been closed successfully")
         super(Connection, self).close()
 
+    def receive_ssh_data(self, count, timeout):
+        if timeout:
+            start = time.time()
+            while not self._ssh_shell.recv_ready():
+                if time.time() - start >= timeout:
+                    raise AnsibleConnectionFailure("timeout waiting ssh data")
+                time.sleep(0.001)
+        return self._ssh_shell.recv(count)
+
     def receive(self, command=None, prompts=None, answer=None, newline=True, prompt_retry_check=False, check_all=False):
         '''
         Handles receiving of output from command
@@ -376,12 +386,14 @@ class Connection(NetworkConnectionBase):
         buffer_read_timeout = self.get_option('persistent_buffer_read_timeout')
         self._validate_timeout_value(buffer_read_timeout, "persistent_buffer_read_timeout")
 
+        receive_data_timeout = self._ssh_shell.gettimeout()
+
         while True:
             if command_prompt_matched:
                 try:
                     signal.signal(signal.SIGALRM, self._handle_buffer_read_timeout)
                     signal.setitimer(signal.ITIMER_REAL, buffer_read_timeout)
-                    data = self._ssh_shell.recv(256)
+                    data = self.receive_ssh_data(256, receive_data_timeout)
                     signal.alarm(0)
                     # if data is still received on channel it indicates the prompt string
                     # is wrongly matched in between response chunks, continue to read
@@ -395,7 +407,7 @@ class Connection(NetworkConnectionBase):
                 except AnsibleCmdRespRecv:
                     return self._command_response
             else:
-                data = self._ssh_shell.recv(256)
+                data = self.receive_ssh_data(256, receive_data_timeout)
 
             # when a channel stream is closed, received data will be empty
             if not data:
diff --git a/test/units/plugins/connection/test_network_cli.py b/test/units/plugins/connection/test_network_cli.py
index cbd64addd2..a0d5d21186 100644
--- a/test/units/plugins/connection/test_network_cli.py
+++ b/test/units/plugins/connection/test_network_cli.py
@@ -126,6 +126,8 @@ class TestConnectionClass(unittest.TestCase):
 
         mock__shell = MagicMock()
         conn._ssh_shell = mock__shell
+        conn._ssh_shell.recv_ready.return_value = True
+        conn._ssh_shell.gettimeout.return_value = 10
 
         response = b"""device#command
         command response