def testFiles(self):
FILE_OUT = "/tmp/udtsocket_test_out.txt"
sock1 = UDTSocket()
sock2 = UDTSocket()
sock1.bind("0.0.0.0:7015")
sock1.listen()
other_thread = threading.Thread(target=self.sendFile, args=(sock2, ))
other_thread.start()
sock, _ = sock1.accept()
sock.recvfile(FILE_OUT, 0, len(UDTSocketTest.CONTENTS))
with open(FILE_OUT, "r") as fr:
self.assertEqual(UDTSocketTest.CONTENTS, fr.read())
os.remove(FILE_OUT)
other_thread.join()
def testNoBlock(self):
other_thread = threading.Thread(target=self.otherConnectNoBlock)
other_thread.start()
time.sleep(0.1)
self.socket = UDTSocket(type=socket.SOCK_DGRAM)
self.assertEqual(socket.SOCK_DGRAM, self.socket.type)
self.socket.UDT_RCVSYN = False
self.assertFalse(self.socket.UDT_RCVSYN)
self.socket.bind(("0.0.0.0", 7014))
self.socket.listen()
sock = None
while sock is None:
try:
sock, _ = self.socket.accept()
except UDTException as e:
self.assertEqual(UDTException.EASYNCRCV, e.error_code)
msg = bytearray(5)
msg[0] = 0
while msg[0] == 0:
try:
sock.recvmsg(msg)
except UDTException as e:
self.assertEqual(UDTException.EASYNCRCV, e.error_code)
self.assertEqual(b"hello", msg)
other_thread.join()
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 testRecvSend(self):
self.socket = UDTSocket()
self.assertEqual(UDTSocket.Status.INIT, self.socket.status)
self.socket.bind("0.0.0.0:7013")
self.assertEqual(UDTSocket.Status.OPENED, self.socket.status)
self.assertEqual(("0.0.0.0", 7013), self.socket.address)
self.socket.listen()
self.assertEqual(UDTSocket.Status.LISTENING, self.socket.status)
other_thread = threading.Thread(target=self.otherConnect)
other_thread.start()
sock, _ = self.socket.accept()
self.assertEqual(UDTSocket.Status.CONNECTED, sock.status)
self.assertEqual(socket.AF_INET, sock.family)
self.assertEqual(socket.SOCK_STREAM, sock.type)
self.assertEqual(("127.0.0.1", 7013), sock.address)
self.assertEqual("127.0.0.1", sock.peer_address[0][0:9])
msg = bytearray(5)
sock.recv(msg)
self.assertEqual(b"hello", msg)
msg = b"12345"
sock.recv(msg)
self.assertEqual(b"hello", msg)
buf = bytearray(6)
msg = memoryview(buf)[1:]
sock.recv(msg)
self.assertEqual(b"hello", msg)
other_thread.join()
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 testOptions(self):
self.socket = UDTSocket()
self.assertEqual(socket.AF_INET, self.socket.family)
self.assertEqual(socket.SOCK_STREAM, self.socket.type)
self.assertTrue(self.socket.UDT_SNDSYN)
self.assertTrue(self.socket.UDT_RCVSYN)
self.assertEqual(65536, self.socket.UDP_SNDBUF)
self.assertEqual(12288000, self.socket.UDP_RCVBUF)
self.assertEqual(12058624, self.socket.UDT_SNDBUF)
self.assertEqual(12058624, self.socket.UDT_RCVBUF)
self.assertEqual(25600, self.socket.UDT_FC)
self.socket.UDP_SNDBUF = 1024001
self.assertEqual(1024001, self.socket.UDP_SNDBUF)
def testEpoll(self):
self.socket = UDTSocket()
self.socket.bind("0.0.0.0:7015")
self.socket.listen()
other_thread = threading.Thread(target=self.otherConnect)
other_thread.start()
sock, _ = self.socket.accept()
poll = UDTEpoll()
poll.add(sock)
rs = []
while len(rs) == 0:
rs, ws, _, _ = poll.wait()
self.assertEqual(sock, rs[0])
self.assertEqual(sock, ws[0])
msg = bytearray(5)
sock.recv(msg)
import socket
from sys import argv
import time
from udt4py import UDTSocket, UDTException
if __name__ == "__main__":
SIZE = 15 * 1000 * 1000
msg = bytearray(SIZE)
sock_type = argv[1][0:3]
argv[1] = argv[1][6:]
HOST, PORT = argv[1].split(':')
addr = (HOST, int(PORT))
N = int(argv[2])
if sock_type == "udt":
sock = UDTSocket()
sock.UDP_SNDBUF = 512 * 1024
sock.UDP_RCVBUF = 2 * 1024 * 1024
sock.UDT_SNDBUF = (15 * 1000 + 1000) * 1000
sock.UDT_RCVBUF = (15 * 1000 + 1000) * 1000
elif sock_type == "tcp":
sock = socket.socket()
else:
raise Exception("Socket type \"%s\" is not supported" % sock_type)
if argv[1].startswith("0.0.0.0"):
sock.bind(addr)
sock.listen(1)
peer, _ = sock.accept()
start = time.time()
all = 0
if sock_type == "udt":
while all < N * SIZE:
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
