Don't drop/lose data when our inner timeout goges off.

This commit is contained in:
Jeff Forcier 2014-02-13 13:44:27 -08:00
parent 244e09f57a
commit 4402f67fa6
1 changed files with 8 additions and 9 deletions

View File

@ -52,6 +52,7 @@ class ProxyCommand(object):
self.cmd = shlsplit(command_line) self.cmd = shlsplit(command_line)
self.process = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE) self.process = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
self.timeout = None self.timeout = None
self.buffer = []
def send(self, content): def send(self, content):
""" """
@ -83,23 +84,21 @@ class ProxyCommand(object):
""" """
try: try:
start = datetime.now() start = datetime.now()
read = [] while len(self.buffer) < size:
while len(read) < size:
if self.timeout is not None: if self.timeout is not None:
elapsed = (datetime.now() - start).microseconds elapsed = (datetime.now() - start).microseconds
timeout = self.timeout * 1000 * 1000 # to microseconds timeout = self.timeout * 1000 * 1000 # to microseconds
# Unsure why the 'default' timeout is too short here -
# causes us to raise before e.g. the SSH banner is read,
# probably generally awful for slow connections.
# Try inflating it some.
timeout = timeout * 2
if elapsed >= timeout: if elapsed >= timeout:
raise socket.timeout() raise socket.timeout()
r, w, x = select([self.process.stdout], [], [], 0.0) r, w, x = select([self.process.stdout], [], [], 0.0)
if r and r[0] == self.process.stdout: if r and r[0] == self.process.stdout:
b = os.read(self.process.stdout.fileno(), 1) b = os.read(self.process.stdout.fileno(), 1)
read.append(b) # Store in class-level buffer for persistence across
result = ''.join(read) # timeouts; this makes us act more like a real socket
# (where timeouts don't actually drop data.)
self.buffer.append(b)
result = ''.join(self.buffer)
self.buffer = []
return result return result
except IOError, e: except IOError, e:
raise ProxyCommandFailure(' '.join(self.cmd), e.strerror) raise ProxyCommandFailure(' '.join(self.cmd), e.strerror)