def otherConnect(self):
sock = UDTSocket()
sock.connect("127.0.0.1:7013")
self.assertEqual(UDTSocket.Status.CONNECTED, sock.status)
sock.send(b"hello")
sock.send(b"hello")
sock.send(b"hello")
def otherConnect(self):
sock = UDTSocket()
sock.connect("127.0.0.1:7013")
self.assertEqual(UDTSocket.Status.CONNECTED, sock.status)
sock.send(b"hello")
sock.send(b"hello")
sock.send(b"hello")
def otherConnectNoBlock(self):
sock = UDTSocket(type=socket.SOCK_DGRAM)
sock.UDT_SNDSYN = False
self.assertFalse(sock.UDT_SNDSYN)
while sock.status != UDTSocket.Status.CONNECTED:
try:
sock.connect("127.0.0.1:7014")
except UDTException as e:
self.assertEqual(UDTException.EASYNCRCV, e.error_code)
self.assertEqual(UDTSocket.Status.CONNECTED, sock.status)
sock.sendmsg(b"hello")
def otherConnectNoBlock(self):
sock = UDTSocket(type=socket.SOCK_DGRAM)
sock.UDT_SNDSYN = False
self.assertFalse(sock.UDT_SNDSYN)
while sock.status != UDTSocket.Status.CONNECTED:
try:
sock.connect("127.0.0.1:7014")
except UDTException as e:
self.assertEqual(UDTException.EASYNCRCV, e.error_code)
self.assertEqual(UDTSocket.Status.CONNECTED, sock.status)
sock.sendmsg(b"hello")
sock.bind(addr)
sock.listen(1)
peer, _ = sock.accept()
start = time.time()
all = 0
if sock_type == "udt":
while all < N * SIZE:
all += peer.recv(msg)
else:
while all < N * SIZE:
msg = peer.recv(4096)
all += len(msg)
peer.send(b'0')
peer.close()
else:
sock.connect(addr)
start = time.time()
for i in range(N):
sent = 0
while (sent < len(msg)):
chunk = sock.send(memoryview(msg)[0:len(msg)-sent])
sent += chunk
if sock_type == "udt":
bye = bytearray(1)
sock.recv(bye)
else:
sock.recv(1)
sock.close()
delta = time.time() - start
print("%.2f sec" % delta)
sock.bind(addr)
sock.listen(1)
peer, _ = sock.accept()
start = time.time()
all = 0
if sock_type == "udt":
while all < N * SIZE:
all += peer.recv(msg)
else:
while all < N * SIZE:
msg = peer.recv(4096)
all += len(msg)
peer.send(b'0')
peer.close()
else:
sock.connect(addr)
start = time.time()
for i in range(N):
sent = 0
while (sent < len(msg)):
chunk = sock.send(memoryview(msg)[0:len(msg) - sent])
sent += chunk
if sock_type == "udt":
bye = bytearray(1)
sock.recv(bye)
else:
sock.recv(1)
sock.close()
delta = time.time() - start
print("%.2f sec" % delta)
def pair(
pairing_name: str,
api_key: str = 'serverlessnetworkingfreetrial',
local_port: int = 10000,
remote_port: int = 10000,
natpunch_timeout: int = 30,
natpunch_server: str = 'services.serverlesstech.net/natpunch'
) -> UDTSocket:
"""
Connect to a remote networking peer.
Performs both NAT punching and rendezvous. The result is a new UDTSocket instance
ready to perform reliable messaging and file transfers.
Parameters:
pairing_name (str): Virtual address; both sides must use the same pairing_name to connect. Required argument.
api_key (str): API key to pass to NAT punch web socket. Optional, defaults to 'serverlessnetworkingfreetrial' (which is only valid if natpunch_server is set to the default)
local_port (int): UDP port number to use on this side of the connection. Optional, defaults to 10000.
remote_port (int): UDP port number to use for the remote side of the connection. Optional, defaults to 10000.
natpunch_timeout (int): Seconds to wait attempting to NAT punch/pair. Optional, defaults to 30. Note that rendezvous has a separate, udt-specified timeout.
natpunch_server (str): Host and path portion of URL to use for websocket NAT punch operation. Optional, defaults to services.serverlesstech.net/natpunch.
Returns:
UDTSocket: A new UDTSocket instance representing the p2p connection for pairing_name, or None if NAT punching failed. Clients should
check the status of the socket to ensure it's in a CONNECTED state before proceeding to use it.
"""
# Create a version of the websocket client class that handles AWS sigv4
# authorization by overriding the 'write_http_request' method with the
# logic to construct an x-amzn-auth header at the last possible moment.
class WebSocketSigv4ClientProtocol(websockets.WebSocketClientProtocol):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def write_http_request(self, path: str, headers) -> None:
# Intercept the GET that initiates the websocket protocol at the point where
# all of its 'real' headers have been constructed. Add in the sigv4 header AWS needs.
credentials = Credentials(os.environ['AWS_ACCESS_KEY_ID'],
os.environ['AWS_SECRET_ACCESS_KEY'],
os.environ['AWS_SESSION_TOKEN'])
sigv4 = SigV4Auth(credentials, 'execute-api',
os.environ['AWS_REGION'])
request = AWSRequest(method='GET',
url='https://' + natpunch_server)
sigv4.add_auth(request)
prepped = request.prepare()
headers['Authorization'] = prepped.headers['Authorization']
headers['X-Amz-Date'] = prepped.headers['X-Amz-Date']
headers['x-amz-security-token'] = prepped.headers[
'x-amz-security-token']
# Run the original code with the added sigv4 auth header now included:
super().write_http_request(path, headers)
async def natpunch():
if (not 'AWS_ACCESS_KEY_ID' in os.environ):
raise Exception(
'missing environment variable(s) required for signing',
'AWS_ACCESS_KEY_ID not present')
if (not 'AWS_SECRET_ACCESS_KEY' in os.environ):
raise Exception(
'missing environment variable(s) required for signing',
'AWS_SECRET_ACCESS_KEY not present')
if (not 'AWS_SESSION_TOKEN' in os.environ):
raise Exception(
'missing environment variable(s) required for signing',
'AWS_SESSION_TOKEN not present')
if (not 'AWS_REGION' in os.environ):
raise Exception(
'missing environment variable(s) required for signing',
'AWS_REGION not present')
async with websockets.connect(
'wss://' + natpunch_server,
create_protocol=WebSocketSigv4ClientProtocol,
extra_headers={'x-api-key': api_key}) as websocket:
msg_as_string = json.dumps({
"action": "pair",
"pairing_name": pairing_name
})
await websocket.send(msg_as_string)
try:
result = await asyncio.wait_for(websocket.recv(),
timeout=natpunch_timeout)
return json.loads(result)['SourceIP']
except asyncio.TimeoutError:
return None
remote_ip = asyncio.run(natpunch())
if (not remote_ip):
return None
usock = UDTSocket()
usock.UDT_MSS = 9000
usock.UDT_RENDEZVOUS = True
usock.bind(('0.0.0.0', local_port))
#print('Trying to connect to ' + remote_ip, flush=True)
usock.connect((remote_ip, remote_port))
return usock
class ClientUDTManager:
def __init__(self, server_controller, hostname, tcp_mode):
self.server_controller = server_controller
self.socket = None
self.hostname = hostname
self.port = None
self.nonce = None
self.tcp_mode = tcp_mode
self.server_udt_manager = None
def connect(self):
self.server_udt_manager = self.server_controller.root.get_udt_manager()(self.tcp_mode)
self.port, self.nonce = self.server_udt_manager.open_connection()
self.connect_to_server()
self.send_nonce()
def send_file(self, file_src, file_dest, block_count, file_size):
self.server_udt_manager.receive_data(file_dest, block_count, file_size)
self.send_data(file_src, block_count)
def connect_to_server(self):
"""
Connects to the provided host and port returning a socket object.
"""
if self.tcp_mode:
logger.debug("TCP mode")
sock_type = socket.SOCK_STREAM
for res in socket.getaddrinfo(self.hostname, self.port, socket.AF_UNSPEC, sock_type):
af, socktype, proto, canonname, sa = res
try:
self.socket = socket.socket(af, socktype, proto)
except socket.error:
self.socket = None
continue
try:
self.socket.connect(sa)
except socket.error:
self.socket.close()
# No need to log error here, some errors are expected
self.socket = None
continue
break
else:
self.socket = UDTSocket()
self.socket.connect((socket.gethostbyname(self.hostname), self.port))
if self.socket is None:
fail('Could not connect to' + self.hostname)
def send_nonce(self):
if not self.tcp_mode:
self.socket.send(bytearray(self.nonce))
else:
self.socket.sendall(self.nonce)
def get_total_recieved(self):
if self.server_udt_manager is None:
return 0
return self.server_udt_manager.get_total_recieved()
def send_data(self, file_src, block_count=0):
"""
Opens the file at the number of blocks passed in and uses that along with
the other parameters to send a file to the host at the specified port.
"""
f = open(file_src, 'r')
f.seek(block_count * CHUNK_SIZE)
data = f.read(CHUNK_SIZE)
byteData = bytearray(CHUNK_SIZE)
while data:
if not self.tcp_mode:
byteData[0:CHUNK_SIZE] = data
self.send_chunk(byteData)
else:
self.socket.sendall(data)
data = f.read(CHUNK_SIZE)
logger.debug("Data sent.")
self.socket.close()
f.close()
def send_chunk(self, data):
size = self.socket.send(data)
if not size == len(data):
self.send_chunk(data[size:])
def generate_nonce(self, length=NONCE_SIZE):
"""Generate pseudorandom number. Ripped from google."""
return ''.join([str(random.randint(0, 9)) for i in range(length)])
def __del__(self):
if(self.socket):
self.socket.close()
