forked from M-Labs/artiq
1
0
Fork 0

artiq_devtool: forward multiple connections.

In particular this is useful for running concurrent arqctl_corelog
and anything else using the management channel, such as
artiq_coreboot.
This commit is contained in:
whitequark 2017-08-23 09:33:07 +00:00
parent 32aac9c70d
commit 3823193060
1 changed files with 29 additions and 23 deletions

View File

@ -112,11 +112,31 @@ def main():
elif action == "connect": elif action == "connect":
transport = client.get_transport() transport = client.get_transport()
def forwarder(port):
def forwarder(local_stream, remote_stream):
try:
while True:
r, _, _ = select.select([local_stream, remote_stream], [], [])
if local_stream in r:
data = local_stream.recv(65535)
if data == b"":
break
remote_stream.sendall(data)
if remote_stream in r:
data = remote_stream.recv(65535)
if data == b"":
break
local_stream.sendall(data)
except Exception as err:
logger.error("Cannot forward on port %s: %s", port, repr(err))
local_stream.close()
remote_stream.close()
def listener(port):
listener = socket.socket() listener = socket.socket()
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(('localhost', port)) listener.bind(('localhost', port))
listener.listen(1) listener.listen(8)
while True: while True:
local_stream, peer_addr = listener.accept() local_stream, peer_addr = listener.accept()
logger.info("Accepting %s:%s and opening SSH channel to %s:%s", logger.info("Accepting %s:%s and opening SSH channel to %s:%s",
@ -128,31 +148,17 @@ def main():
try: try:
remote_stream = \ remote_stream = \
transport.open_channel('direct-tcpip', (args.device, port), peer_addr) transport.open_channel('direct-tcpip', (args.device, port), peer_addr)
except Exception as e: except Exception:
logger.exception("Cannot open channel on port %s", port) logger.exception("Cannot open channel on port %s", port)
continue continue
while True:
try: thread = threading.Thread(target=forwarder, args=(local_stream, remote_stream),
r, _, _ = select.select([local_stream, remote_stream], [], []) name="forward-{}".format(port), daemon=True)
if local_stream in r: thread.start()
data = local_stream.recv(65535)
if data == b"":
break
remote_stream.sendall(data)
if remote_stream in r:
data = remote_stream.recv(65535)
if data == b"":
break
local_stream.sendall(data)
except Exception as e:
logger.exception("Forward error on port %s", port)
break
local_stream.close()
remote_stream.close()
for port in (1380, 1381, 1382, 1383): for port in (1380, 1381, 1382, 1383):
thread = threading.Thread(target=forwarder, args=(port,), thread = threading.Thread(target=listener, args=(port,),
name="port-{}".format(port), daemon=True) name="listen-{}".format(port), daemon=True)
thread.start() thread.start()
logger.info("Connecting to device") logger.info("Connecting to device")