Skip to content

Commit 233df65

Browse files
authored
Merge pull request #638 from Jakuje/sftp-large
Fix large SFTP reads
2 parents 31eecbd + 09bf9fe commit 233df65

File tree

3 files changed

+59
-15
lines changed

3 files changed

+59
-15
lines changed
+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Fixed reading files over SFTP that go over the pre-defined chunk size.
2+
3+
Prior to this change, the files could end up being corrupted, ending up with the last read chunk written to the file instead of the entire payload.
4+
5+
-- by :user:`Jakuje`

src/pylibsshext/sftp.pyx

+10-10
Original file line numberDiff line numberDiff line change
@@ -91,16 +91,16 @@ cdef class SFTP:
9191
if rf is NULL:
9292
raise LibsshSFTPException("Opening remote file [%s] for read failed with error [%s]" % (remote_file, self._get_sftp_error_str()))
9393

94-
while True:
95-
file_data = sftp.sftp_read(rf, <void *>read_buffer, sizeof(char) * 1024)
96-
if file_data == 0:
97-
break
98-
elif file_data < 0:
99-
sftp.sftp_close(rf)
100-
raise LibsshSFTPException("Reading data from remote file [%s] failed with error [%s]"
101-
% (remote_file, self._get_sftp_error_str()))
102-
103-
with open(local_file, 'wb+') as f:
94+
with open(local_file, 'wb') as f:
95+
while True:
96+
file_data = sftp.sftp_read(rf, <void *>read_buffer, sizeof(char) * 1024)
97+
if file_data == 0:
98+
break
99+
elif file_data < 0:
100+
sftp.sftp_close(rf)
101+
raise LibsshSFTPException("Reading data from remote file [%s] failed with error [%s]"
102+
% (remote_file, self._get_sftp_error_str()))
103+
104104
bytes_written = f.write(read_buffer[:file_data])
105105
if bytes_written and file_data != bytes_written:
106106
sftp.sftp_close(rf)

tests/unit/sftp_test.py

+44-5
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
"""Tests suite for sftp."""
44

5+
import random
6+
import string
57
import uuid
68

79
import pytest
@@ -18,11 +20,21 @@ def sftp_session(ssh_client_session):
1820
del sftp_sess # noqa: WPS420
1921

2022

21-
@pytest.fixture
22-
def transmit_payload():
23-
"""Generate a binary test payload."""
24-
uuid_name = uuid.uuid4()
25-
return 'Hello, {name!s}'.format(name=uuid_name).encode()
23+
@pytest.fixture(
24+
params=(32, 1024 + 1),
25+
ids=('small-payload', 'large-payload'),
26+
)
27+
def transmit_payload(request: pytest.FixtureRequest) -> bytes:
28+
"""Generate binary test payloads of assorted sizes.
29+
30+
The choice 32 is arbitrary small value.
31+
32+
The choice 1024 + 1 is meant to be 1B larger than the chunk size used in
33+
:file:`sftp.pyx` to make sure we excercise at least two rounds of reading/writing.
34+
"""
35+
payload_len = request.param
36+
random_bytes = [ord(random.choice(string.printable)) for _ in range(payload_len)]
37+
return bytes(random_bytes)
2638

2739

2840
@pytest.fixture
@@ -48,6 +60,21 @@ def dst_path(file_paths_pair):
4860
return path
4961

5062

63+
@pytest.fixture
64+
def other_payload():
65+
"""Generate a binary test payload."""
66+
uuid_name = uuid.uuid4()
67+
return 'Original content: {name!s}'.format(name=uuid_name).encode()
68+
69+
70+
@pytest.fixture
71+
def pre_existing_dst_path(dst_path, other_payload):
72+
"""Return a data destination path."""
73+
dst_path.write_bytes(other_payload)
74+
assert dst_path.exists()
75+
return dst_path
76+
77+
5178
def test_make_sftp(sftp_session):
5279
"""Smoke-test SFTP instance creation."""
5380
assert sftp_session
@@ -63,3 +90,15 @@ def test_get(dst_path, src_path, sftp_session, transmit_payload):
6390
"""Check that SFTP file download works."""
6491
sftp_session.get(str(src_path), str(dst_path))
6592
assert dst_path.read_bytes() == transmit_payload
93+
94+
95+
def test_get_existing(pre_existing_dst_path, src_path, sftp_session, transmit_payload):
96+
"""Check that SFTP file download works when target file exists."""
97+
sftp_session.get(str(src_path), str(pre_existing_dst_path))
98+
assert pre_existing_dst_path.read_bytes() == transmit_payload
99+
100+
101+
def test_put_existing(pre_existing_dst_path, src_path, sftp_session, transmit_payload):
102+
"""Check that SFTP file upload works when target file exists."""
103+
sftp_session.put(str(src_path), str(pre_existing_dst_path))
104+
assert pre_existing_dst_path.read_bytes() == transmit_payload

0 commit comments

Comments
 (0)