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 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()
class UDTEpollTest(unittest.TestCase): 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) def otherConnect(self): sock = UDTSocket() sock.connect("127.0.0.1:7015") self.assertEqual(UDTSocket.Status.CONNECTED, sock.status) sock.send(b"hello")
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: 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()
class UDTSocketTest(unittest.TestCase): 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 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 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 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 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") CONTENTS = "contents123456" 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 sendFile(self, sock): FILE_IN = "/tmp/udtsocket_test_in.txt" with open(FILE_IN, "w") as fw: fw.write(UDTSocketTest.CONTENTS) sock.connect("127.0.0.1:7015") sock.sendfile(FILE_IN) os.remove(FILE_IN)
class ServerUDTManager: def __init__(self, tcp_mode): self.tcp_mode = tcp_mode self.udt_sock = None self.conn = None self.sock = self.get_socket() self.port = self.sock.getsockname()[1] self.nonce = self.generate_nonce() self.size = 0 def open_connection(self): if not self.tcp_mode: self.udt_sock = UDTSocket() self.udt_sock.bind(self.sock.fileno()) self.udt_sock.listen() listening_thread = threading.Thread(target=self.accept_and_verify) listening_thread.start() return (self.port, self.nonce) def get_total_recieved(self): return self.size def accept_and_verify(self): if not self.tcp_mode: self.conn, addr = self.udt_sock.accept() logger.info('Connected by %s', addr) recvd_nonce = bytearray(NONCE_SIZE) self.conn.recv(recvd_nonce) recvd_nonce = str(recvd_nonce) else: self.conn, addr = self.sock.accept() logger.info('Connected by %s', addr) recvd_nonce = self.conn.recv(NONCE_SIZE) if recvd_nonce != self.nonce: fail(format("Received nonce %s doesn't match %s.", recvd_nonce, self.nonce)) logger.debug("Nonce verified.") def receive_data(self, output_file, block_count, file_size): """ Receives data and writes it to disk, stops when it is no longer receiving data. """ def receive_data_threaded(output_file, block_count, file_size): logger.debug("Receiving data...") output_file = open(output_file, "r+") output_file.seek(block_count * CHUNK_SIZE) self.size = block_count * CHUNK_SIZE data = bytearray(CHUNK_SIZE) if not self.tcp_mode: while 1: len_rec = self.conn.recv(data) data = str(data) output_file.write(data[:len_rec]) self.size += len_rec if len_rec == 0 or str(self.size) == str(file_size): break else: while 1: data = self.conn.recv(CHUNK_SIZE) output_file.write(data) self.size += len(data) if len(data) == 0: break logger.debug("Closing file... " + output_file.name) output_file.close() thread = threading.Thread(target=receive_data_threaded, args=(output_file, block_count, file_size)) thread.start() return thread def get_socket(self): """ Opens and returns a socket on an open port. """ s = None if self.tcp_mode: sock_type = socket.SOCK_STREAM else: sock_type = socket.SOCK_DGRAM try: s = socket.socket(socket.AF_INET, sock_type) except socket.error as msg: fail(msg) try: s.bind(('', 0)) if self.tcp_mode: s.listen(1) except socket.error as msg: s.close() fail(str(msg)) return s 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)])
class UDTSocketTest(unittest.TestCase): 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 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 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 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 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") CONTENTS = "contents123456" 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 sendFile(self, sock): FILE_IN = "/tmp/udtsocket_test_in.txt" with open(FILE_IN, "w") as fw: fw.write(UDTSocketTest.CONTENTS) sock.connect("127.0.0.1:7015") sock.sendfile(FILE_IN) os.remove(FILE_IN)