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")
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()