Skip to content

Commit dae3db6

Browse files
committed
fix: Handle stderr properly in SSH transport
Instead of reading a small amount of data in a blocking way from the stdout and stderr streams, enable non-blocking I/O and attempt to read as much data as possible. In addition, make a convenience function to handle the stream being closed and timeouts so that we don't need to store a temp variable in the loop. This prevents the situation that lead to issue #74 where the stderr data is written to the output XML if stdout is already closed. Fixes #74
1 parent 3d13ebd commit dae3db6

File tree

2 files changed

+37
-20
lines changed

2 files changed

+37
-20
lines changed

src/itoolkit/transport/ssh.py

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
"""
2424
from .base import XmlServiceTransport
2525

26+
import socket
27+
2628
__all__ = [
2729
'SshTransport'
2830
]
@@ -75,33 +77,45 @@ def _call(self, tk):
7577
"""
7678
command = "/QOpenSys/pkgs/bin/xmlservice-cli"
7779
stdin, stdout, stderr = self.conn.exec_command(command)
80+
channel = stdout.channel
81+
7882
xml_in = tk.xml_in()
7983
stdin.write(xml_in.encode())
80-
stdin.flush()
81-
stdin.channel.shutdown_write()
84+
stdin.close()
85+
channel.shutdown_write()
86+
87+
# Disable blocking I/O
88+
# chan.settimeout(0.0) is equivalent to chan.setblocking(0)
89+
# https://docs.paramiko.org/en/stable/api/channel.html#paramiko.channel.Channel.settimeout
90+
channel.settimeout(0.0)
8291

8392
# rather than doing all this loop-de-loop, we could instead just use
8493
# a single call to stdout.readlines(), but there is a remote
8594
# possibility that the process is hanging writing to a filled-up
86-
# stderr pipe. So, we read a little from both until we're all done
95+
# stderr pipe. So, we read from both until we're all done
8796
err_out = b""
8897
xml_out = b""
89-
blockSize = 64 # arbitrary
90-
while not stderr.closed or not stdout.closed:
91-
if not stderr.closed:
92-
newData = stderr.read(blockSize)
93-
if not newData:
94-
stderr.close() # reaching EOF doesn't implicitly close
95-
else:
96-
err_out += newData
97-
if not stdout.closed:
98-
newData = stdout.read(blockSize)
99-
if not newData:
100-
stdout.close() # reaching EOF doesn't implicitly close
101-
else:
102-
xml_out += newData
103-
stdout.channel.close()
104-
stderr.channel.close()
98+
99+
# Convenience wrapper for reading data from stdout/stderr
100+
# Returns empty binary string if EOF *or* timeout (no data)
101+
# Closes file on EOF
102+
def read_data(f):
103+
if f.closed:
104+
return b""
105+
106+
try:
107+
data = f.read()
108+
if not data:
109+
f.close()
110+
return data
111+
except socket.timeout:
112+
return b""
113+
114+
while not all((stdout.closed, stderr.closed)):
115+
xml_out += read_data(stdout)
116+
err_out += read_data(stderr)
117+
118+
channel.close()
105119
return xml_out
106120

107121
def _close(self):

tests/test_unit_transport_ssh.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,10 @@ def test_ssh_transport_minimal(mocker):
3030
assert isinstance(out, (bytes, str))
3131

3232
command = "/QOpenSys/pkgs/bin/xmlservice-cli"
33-
ssh_client.exec_command.assert_called_once_with(command)
33+
ssh_client.exec_command.assert_called_once()
34+
35+
args = ssh_client.exec_command.call_args
36+
assert args[0] == (command,)
3437

3538

3639
def test_ssh_transport_raises_when_closed(mocker):

0 commit comments

Comments
 (0)