forked from M-Labs/artiq
1
0
Fork 0

artiq_devtool: more robust port forwarding.

This commit is contained in:
whitequark 2016-11-10 21:05:24 +00:00
parent fbc2420443
commit 8b6418c604
1 changed files with 17 additions and 11 deletions

View File

@ -105,6 +105,7 @@ def main():
elif action == "connect": elif action == "connect":
def forwarder(port): def forwarder(port):
listener = socket.socket() listener = socket.socket()
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(('localhost', port)) listener.bind(('localhost', port))
listener.listen(1) listener.listen(1)
while True: while True:
@ -114,17 +115,22 @@ def main():
remote_stream = get_ssh().get_transport() \ remote_stream = get_ssh().get_transport() \
.open_channel('direct-tcpip', (args.ip, port), peer_addr) .open_channel('direct-tcpip', (args.ip, port), peer_addr)
while True: while True:
r, w, x = select.select([local_stream, remote_stream], [], []) try:
if local_stream in r: r, w, x = select.select([local_stream, remote_stream], [], [])
data = local_stream.recv(1024) if local_stream in r:
if data == b"": data = local_stream.recv(1024)
break if data == b"":
remote_stream.send(data) break
if remote_stream in r: remote_stream.send(data)
data = remote_stream.recv(1024) if remote_stream in r:
if data == b"": data = remote_stream.recv(1024)
break if data == b"":
local_stream.send(data) break
local_stream.send(data)
except Exception as e:
logger.exception("Forward error on port %s", port)
local_stream.close()
remote_stream.close()
for port in (1381, 1382): for port in (1381, 1382):
thread = threading.Thread(target=forwarder, args=(port,), thread = threading.Thread(target=forwarder, args=(port,),