From 1a18e75b845078407c35ff39e7d48e074651f373 Mon Sep 17 00:00:00 2001
From: Panos <pkittenis@users.noreply.github.com>
Date: Tue, 22 Mar 2022 19:27:46 +0000
Subject: [PATCH] More scp_send race condition fixes (#341)

* More scp_send race condition fixes - resolves #337
* Updated changelog
---
 Changelog.rst                        |  8 ++++++++
 pssh/clients/native/single.py        |  4 +++-
 tests/native/test_parallel_client.py | 18 +++++++++++-------
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/Changelog.rst b/Changelog.rst
index 7f4ec9df..9dd2592b 100644
--- a/Changelog.rst
+++ b/Changelog.rst
@@ -1,6 +1,14 @@
 Change Log
 ============
 
+2.9.1
++++++
+
+Fixes
+------
+
+* Even more rare race condition would sometimes cause ``scp_send`` to write incomplete files - #337
+
 2.9.0
 +++++
 
diff --git a/pssh/clients/native/single.py b/pssh/clients/native/single.py
index ab8ee7af..672a3353 100644
--- a/pssh/clients/native/single.py
+++ b/pssh/clients/native/single.py
@@ -670,7 +670,9 @@ def _scp_send(self, local_file, remote_file):
             raise SCPError(msg, remote_file, self.host, ex)
         finally:
             self._eagain(chan.flush)
-            chan.close()
+            self._eagain(chan.send_eof)
+            self._eagain(chan.wait_eof)
+            self._eagain(chan.wait_closed)
 
     def _sftp_openfh(self, open_func, remote_file, *args):
         try:
diff --git a/tests/native/test_parallel_client.py b/tests/native/test_parallel_client.py
index 82a70d96..db25153f 100644
--- a/tests/native/test_parallel_client.py
+++ b/tests/native/test_parallel_client.py
@@ -1519,13 +1519,9 @@ def test_scp_send_dir_recurse(self):
             except OSError:
                 pass
 
-    def test_scp_send_larger_files(self):
-        hosts = ['127.0.0.1%s' % (i,) for i in range(1, 3)]
-        servers = [OpenSSHServer(host, port=self.port) for host in hosts]
-        for server in servers:
-            server.start_server()
+    def _scp_larger_files(self, hosts):
         client = ParallelSSHClient(
-            hosts, port=self.port, pkey=self.user_key, num_retries=1, timeout=1,
+            hosts, port=self.port, pkey=self.user_key, num_retries=1, timeout=30,
             pool_size=len(hosts),
         )
         local_filename = 'test_file'
@@ -1537,7 +1533,7 @@ def test_scp_send_larger_files(self):
         remote_file_names = [arg['remote_file'] for arg in copy_args]
         sha = sha256()
         with open(local_filename, 'wb') as file_h:
-            for _ in range(10000):
+            for _ in range(1000):
                 data = os.urandom(1024)
                 file_h.write(data)
                 sha.update(data)
@@ -1570,6 +1566,14 @@ def test_scp_send_larger_files(self):
             except OSError:
                 pass
 
+    def test_scp_send_larger_files(self):
+        hosts = ['127.0.0.1%s' % (i,) for i in range(1, 3)]
+        servers = [OpenSSHServer(host, port=self.port) for host in hosts]
+        for server in servers:
+            server.start_server()
+        for _ in range(20):
+            self._scp_larger_files(hosts)
+
     def test_scp_bad_copy_args(self):
         client = ParallelSSHClient([self.host, self.host])
         copy_args = [{'local_file': 'fake_file', 'remote_file': 'fake_remote_file'}]