From e5d1d11c3eaab4dd4114a8c7cda7c54a8417a61e Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 2 Mar 2019 13:46:45 +0200 Subject: [PATCH 1/5] Add UDP to socket --- tests/snippets/stdlib_socket.py | 39 +++++++++++++++------ vm/src/stdlib/socket.rs | 61 +++++++++++++++++++++++++-------- 2 files changed, 76 insertions(+), 24 deletions(-) diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py index e0ea6d168..2a3c3a94e 100644 --- a/tests/snippets/stdlib_socket.py +++ b/tests/snippets/stdlib_socket.py @@ -1,6 +1,10 @@ import socket from testutils import assertRaises +MESSAGE_A = b'aaaa' +MESSAGE_B= b'bbbbb' + +# TCP listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) listener.bind(("127.0.0.1", 0)) @@ -10,16 +14,12 @@ connector = socket.socket(socket.AF_INET, socket.SOCK_STREAM) connector.connect(("127.0.0.1", listener.getsockname()[1])) connection = listener.accept()[0] -message_a = b'aaaa' -message_b = b'bbbbb' - -connector.send(message_a) -connection.send(message_b) -recv_a = connection.recv(len(message_a)) -recv_b = connector.recv(len(message_b)) -assert recv_a == message_a -assert recv_b == message_b - +connector.send(MESSAGE_A) +connection.send(MESSAGE_B) +recv_a = connection.recv(len(MESSAGE_A)) +recv_b = connector.recv(len(MESSAGE_B)) +assert recv_a == MESSAGE_A +assert recv_b == MESSAGE_B connection.close() connector.close() listener.close() @@ -35,3 +35,22 @@ with assertRaises(TypeError): s.bind((888, 8888)) s.close() + +# UDP +sock1 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) +sock1.bind(("127.0.0.1", 0)) + +sock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) +sock2.bind(("127.0.0.1", 0)) + +sock1.connect(("127.0.0.1", sock2.getsockname()[1])) +sock2.connect(("127.0.0.1", sock1.getsockname()[1])) + +sock1.send(MESSAGE_A) +sock2.send(MESSAGE_B) +recv_a = sock2.recv(len(MESSAGE_A)) +recv_b = sock1.recv(len(MESSAGE_B)) +assert recv_a == MESSAGE_A +assert recv_b == MESSAGE_B +sock1.close() +sock2.close() diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index 30a1bbdc5..e85b22d81 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -2,7 +2,7 @@ use std::cell::RefCell; use std::io; use std::io::Read; use std::io::Write; -use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::net::{SocketAddr, TcpListener, TcpStream, UdpSocket}; use std::ops::DerefMut; use crate::obj::objbytes; @@ -53,7 +53,7 @@ impl SocketKind { enum Connection { TcpListener(TcpListener), TcpStream(TcpStream), - // UdpSocket(UdpSocket), + UdpSocket(UdpSocket), } impl Connection { @@ -67,6 +67,7 @@ impl Connection { fn local_addr(&self) -> io::Result { match self { Connection::TcpListener(con) => con.local_addr(), + Connection::UdpSocket(con) => con.local_addr(), _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), } } @@ -76,6 +77,7 @@ impl Read for Connection { fn read(&mut self, buf: &mut [u8]) -> io::Result { match self { Connection::TcpStream(con) => con.read(buf), + Connection::UdpSocket(con) => con.recv(buf), _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), } } @@ -85,6 +87,7 @@ impl Write for Connection { fn write(&mut self, buf: &[u8]) -> io::Result { match self { Connection::TcpStream(con) => con.write(buf), + Connection::UdpSocket(con) => con.send(buf), _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), } } @@ -153,12 +156,27 @@ fn socket_connect(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { let mut socket = get_socket(zelf); - if let Ok(stream) = TcpStream::connect(address_string) { - socket.con = Some(Connection::TcpStream(stream)); - Ok(vm.get_none()) - } else { - // TODO: Socket error - Err(vm.new_type_error("socket failed".to_string())) + match socket.socket_kind { + SocketKind::Stream => { + if let Ok(stream) = TcpStream::connect(address_string) { + socket.con = Some(Connection::TcpStream(stream)); + Ok(vm.get_none()) + } else { + // TODO: Socket error + Err(vm.new_type_error("socket failed".to_string())) + } + } + SocketKind::Dgram => { + if let Some(Connection::UdpSocket(con)) = &socket.con { + match con.connect(address_string) { + Ok(_) => Ok(vm.get_none()), + // TODO: Socket error + Err(_) => Err(vm.new_type_error("socket failed".to_string())), + } + } else { + Err(vm.new_type_error("".to_string())) + } + } } } @@ -173,12 +191,25 @@ fn socket_bind(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { let mut socket = get_socket(zelf); - if let Ok(stream) = TcpListener::bind(address_string) { - socket.con = Some(Connection::TcpListener(stream)); - Ok(vm.get_none()) - } else { - // TODO: Socket error - Err(vm.new_type_error("socket failed".to_string())) + match socket.socket_kind { + SocketKind::Stream => { + if let Ok(stream) = TcpListener::bind(address_string) { + socket.con = Some(Connection::TcpListener(stream)); + Ok(vm.get_none()) + } else { + // TODO: Socket error + Err(vm.new_type_error("socket failed".to_string())) + } + } + SocketKind::Dgram => { + if let Ok(dgram) = UdpSocket::bind(address_string) { + socket.con = Some(Connection::UdpSocket(dgram)); + Ok(vm.get_none()) + } else { + // TODO: Socket error + Err(vm.new_type_error("socket failed".to_string())) + } + } } } @@ -325,6 +356,8 @@ pub fn mk_module(ctx: &PyContext) -> PyObjectRef { ctx.new_int(SocketKind::Stream as i32), ); + ctx.set_attr(&py_mod, "SOCK_DGRAM", ctx.new_int(SocketKind::Dgram as i32)); + let socket = { let socket = ctx.new_class("socket", ctx.object()); ctx.set_attr(&socket, "__new__", ctx.new_rustfunc(socket_new)); From 86b60faa61d80458e6304db5b6ad596da52f4183 Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 2 Mar 2019 17:12:13 +0200 Subject: [PATCH 2/5] Add socket.{sendto, recvfrom} --- tests/snippets/stdlib_socket.py | 6 ++- vm/src/stdlib/socket.rs | 87 +++++++++++++++++++++++++++++---- 2 files changed, 82 insertions(+), 11 deletions(-) diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py index 2a3c3a94e..af29b451c 100644 --- a/tests/snippets/stdlib_socket.py +++ b/tests/snippets/stdlib_socket.py @@ -41,8 +41,12 @@ sock1 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) sock1.bind(("127.0.0.1", 0)) sock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) -sock2.bind(("127.0.0.1", 0)) +sock2.sendto(MESSAGE_A, sock1.getsockname()) +(recv_a, addr) = sock1.recvfrom(len(MESSAGE_A)) +assert recv_a == MESSAGE_A + +sock2.bind(("127.0.0.1", 0)) sock1.connect(("127.0.0.1", sock2.getsockname()[1])) sock2.connect(("127.0.0.1", sock1.getsockname()[1])) diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index e85b22d81..37dc8a6fa 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -71,6 +71,13 @@ impl Connection { _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), } } + + fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { + match self { + Connection::UdpSocket(con) => con.recv_from(buf), + _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), + } + } } impl Read for Connection { @@ -298,6 +305,34 @@ fn socket_recv(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { Ok(vm.ctx.new_bytes(buffer)) } +fn socket_recvfrom(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { + arg_check!( + vm, + args, + required = [(zelf, None), (bufsize, Some(vm.ctx.int_type()))] + ); + + let mut socket = get_socket(zelf); + + let mut buffer = vec![0u8; objint::get_value(bufsize).to_usize().unwrap()]; + let ret = match socket.con { + Some(ref mut v) => v.recv_from(&mut buffer), + None => return Err(vm.new_type_error("".to_string())), + }; + + let addr = match ret { + Ok((_size, addr)) => addr, + _ => return Err(vm.new_type_error("".to_string())), + }; + + let elements = RefCell::new(vec![vm.ctx.new_bytes(buffer), get_addr_tuple(vm, addr)?]); + + Ok(PyObject::new( + PyObjectPayload::Sequence { elements }, + vm.ctx.tuple_type(), + )) +} + fn socket_send(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, @@ -313,6 +348,34 @@ fn socket_send(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { Ok(vm.get_none()) } +fn socket_sendto(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { + arg_check!( + vm, + args, + required = [ + (zelf, None), + (bytes, Some(vm.ctx.bytes_type())), + (address, Some(vm.ctx.tuple_type())) + ] + ); + let address_string = get_address_string(vm, address)?; + + let socket = get_socket(zelf); + + match socket.socket_kind { + SocketKind::Dgram => { + // We can't do sendto without bind in std::net::UdpSocket + if let Ok(dgram) = UdpSocket::bind("0.0.0.0:0") { + if let Ok(_) = dgram.send_to(&objbytes::get_value(&bytes), address_string) { + return Ok(vm.get_none()); + } + } + Err(vm.new_type_error("socket failed".to_string())) + } + _ => Err(vm.new_not_implemented_error("".to_string())), + } +} + fn socket_close(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(zelf, None)]); @@ -331,20 +394,22 @@ fn socket_getsockname(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { }; match addr { - Ok(addr) => { - let port = vm.ctx.new_int(addr.port()); - let ip = vm.ctx.new_str(addr.ip().to_string()); - let elements = RefCell::new(vec![ip, port]); - - Ok(PyObject::new( - PyObjectPayload::Sequence { elements }, - vm.ctx.tuple_type(), - )) - } + Ok(addr) => get_addr_tuple(vm, addr), _ => Err(vm.new_type_error("".to_string())), } } +fn get_addr_tuple(vm: &mut VirtualMachine, addr: SocketAddr) -> PyResult { + let port = vm.ctx.new_int(addr.port()); + let ip = vm.ctx.new_str(addr.ip().to_string()); + let elements = RefCell::new(vec![ip, port]); + + Ok(PyObject::new( + PyObjectPayload::Sequence { elements }, + vm.ctx.tuple_type(), + )) +} + pub fn mk_module(ctx: &PyContext) -> PyObjectRef { let py_mod = ctx.new_module(&"socket".to_string(), ctx.new_scope(None)); @@ -369,6 +434,8 @@ pub fn mk_module(ctx: &PyContext) -> PyObjectRef { ctx.set_attr(&socket, "listen", ctx.new_rustfunc(socket_listen)); ctx.set_attr(&socket, "close", ctx.new_rustfunc(socket_close)); ctx.set_attr(&socket, "getsockname", ctx.new_rustfunc(socket_getsockname)); + ctx.set_attr(&socket, "sendto", ctx.new_rustfunc(socket_sendto)); + ctx.set_attr(&socket, "recvfrom", ctx.new_rustfunc(socket_recvfrom)); socket }; ctx.set_attr(&py_mod, "socket", socket.clone()); From 9b5ba068afc88a439999f884bb8e8f4ddbc9ca0f Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 2 Mar 2019 17:21:36 +0200 Subject: [PATCH 3/5] Return addr in accept --- tests/snippets/stdlib_socket.py | 3 ++- vm/src/stdlib/socket.rs | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py index af29b451c..54ae180c3 100644 --- a/tests/snippets/stdlib_socket.py +++ b/tests/snippets/stdlib_socket.py @@ -12,7 +12,8 @@ listener.listen(1) connector = socket.socket(socket.AF_INET, socket.SOCK_STREAM) connector.connect(("127.0.0.1", listener.getsockname()[1])) -connection = listener.accept()[0] +(connection, addr) = listener.accept() +assert addr == connector.getsockname() connector.send(MESSAGE_A) connection.send(MESSAGE_B) diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index 37dc8a6fa..f4b88b85f 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -68,7 +68,7 @@ impl Connection { match self { Connection::TcpListener(con) => con.local_addr(), Connection::UdpSocket(con) => con.local_addr(), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), + Connection::TcpStream(con) => con.local_addr(), } } @@ -263,8 +263,8 @@ fn socket_accept(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { None => return Err(vm.new_type_error("".to_string())), }; - let tcp_stream = match ret { - Ok((socket, _addr)) => socket, + let (tcp_stream, addr) = match ret { + Ok((socket, addr)) => (socket, addr), _ => return Err(vm.new_type_error("".to_string())), }; @@ -281,7 +281,7 @@ fn socket_accept(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { zelf.typ(), ); - let elements = RefCell::new(vec![sock_obj, vm.get_none()]); + let elements = RefCell::new(vec![sock_obj, get_addr_tuple(vm, addr)?]); Ok(PyObject::new( PyObjectPayload::Sequence { elements }, From f09f75ac8d5c1190aa6d7d89ecf652f786e62bc0 Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 2 Mar 2019 17:45:26 +0200 Subject: [PATCH 4/5] Do implicit bind in socket.sendto --- tests/snippets/stdlib_socket.py | 28 +++++++++++++++++++++------- vm/src/stdlib/socket.rs | 33 ++++++++++++++++++++++++++------- 2 files changed, 47 insertions(+), 14 deletions(-) diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py index 54ae180c3..f984515ea 100644 --- a/tests/snippets/stdlib_socket.py +++ b/tests/snippets/stdlib_socket.py @@ -44,18 +44,32 @@ sock1.bind(("127.0.0.1", 0)) sock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) sock2.sendto(MESSAGE_A, sock1.getsockname()) -(recv_a, addr) = sock1.recvfrom(len(MESSAGE_A)) +(recv_a, addr1) = sock1.recvfrom(len(MESSAGE_A)) assert recv_a == MESSAGE_A -sock2.bind(("127.0.0.1", 0)) -sock1.connect(("127.0.0.1", sock2.getsockname()[1])) -sock2.connect(("127.0.0.1", sock1.getsockname()[1])) +sock2.sendto(MESSAGE_B, sock1.getsockname()) +(recv_b, addr2) = sock1.recvfrom(len(MESSAGE_B)) +assert recv_b == MESSAGE_B +assert addr1[0] == addr2[0] +assert addr1[1] == addr2[1] + +sock2.close() + +sock3 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) +sock3.bind(("127.0.0.1", 0)) +sock3.sendto(MESSAGE_A, sock1.getsockname()) +(recv_a, addr) = sock1.recvfrom(len(MESSAGE_A)) +assert recv_a == MESSAGE_A +assert addr == sock3.getsockname() + +sock1.connect(("127.0.0.1", sock3.getsockname()[1])) +sock3.connect(("127.0.0.1", sock1.getsockname()[1])) sock1.send(MESSAGE_A) -sock2.send(MESSAGE_B) -recv_a = sock2.recv(len(MESSAGE_A)) +sock3.send(MESSAGE_B) +recv_a = sock3.recv(len(MESSAGE_A)) recv_b = sock1.recv(len(MESSAGE_B)) assert recv_a == MESSAGE_A assert recv_b == MESSAGE_B sock1.close() -sock2.close() +sock3.close() diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index f4b88b85f..d0ff61b37 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -2,7 +2,7 @@ use std::cell::RefCell; use std::io; use std::io::Read; use std::io::Write; -use std::net::{SocketAddr, TcpListener, TcpStream, UdpSocket}; +use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket}; use std::ops::DerefMut; use crate::obj::objbytes; @@ -78,6 +78,13 @@ impl Connection { _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), } } + + fn send_to(&self, buf: &[u8], addr: A) -> io::Result { + match self { + Connection::UdpSocket(con) => con.send_to(buf, addr), + _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), + } + } } impl Read for Connection { @@ -360,17 +367,29 @@ fn socket_sendto(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { ); let address_string = get_address_string(vm, address)?; - let socket = get_socket(zelf); + let mut socket = get_socket(zelf); match socket.socket_kind { SocketKind::Dgram => { - // We can't do sendto without bind in std::net::UdpSocket - if let Ok(dgram) = UdpSocket::bind("0.0.0.0:0") { - if let Ok(_) = dgram.send_to(&objbytes::get_value(&bytes), address_string) { - return Ok(vm.get_none()); + match socket.con { + Some(ref mut v) => { + if let Ok(_) = v.send_to(&objbytes::get_value(&bytes), address_string) { + Ok(vm.get_none()) + } else { + Err(vm.new_type_error("socket failed".to_string())) + } + } + None => { + // Doing implicit bind + if let Ok(dgram) = UdpSocket::bind("0.0.0.0:0") { + if let Ok(_) = dgram.send_to(&objbytes::get_value(&bytes), address_string) { + socket.con = Some(Connection::UdpSocket(dgram)); + return Ok(vm.get_none()); + } + } + Err(vm.new_type_error("socket failed".to_string())) } } - Err(vm.new_type_error("socket failed".to_string())) } _ => Err(vm.new_not_implemented_error("".to_string())), } From 283b331323178e8a6c6f8f184d0a7e46d90b354a Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sun, 3 Mar 2019 18:22:20 +0200 Subject: [PATCH 5/5] Use vm.ctx.new_tuple --- vm/src/stdlib/socket.rs | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index d0ff61b37..7626cc182 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -288,12 +288,9 @@ fn socket_accept(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { zelf.typ(), ); - let elements = RefCell::new(vec![sock_obj, get_addr_tuple(vm, addr)?]); + let addr_tuple = get_addr_tuple(vm, addr)?; - Ok(PyObject::new( - PyObjectPayload::Sequence { elements }, - vm.ctx.tuple_type(), - )) + Ok(vm.ctx.new_tuple(vec![sock_obj, addr_tuple])) } fn socket_recv(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { @@ -332,12 +329,9 @@ fn socket_recvfrom(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { _ => return Err(vm.new_type_error("".to_string())), }; - let elements = RefCell::new(vec![vm.ctx.new_bytes(buffer), get_addr_tuple(vm, addr)?]); + let addr_tuple = get_addr_tuple(vm, addr)?; - Ok(PyObject::new( - PyObjectPayload::Sequence { elements }, - vm.ctx.tuple_type(), - )) + Ok(vm.ctx.new_tuple(vec![vm.ctx.new_bytes(buffer), addr_tuple])) } fn socket_send(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { @@ -421,12 +415,8 @@ fn socket_getsockname(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { fn get_addr_tuple(vm: &mut VirtualMachine, addr: SocketAddr) -> PyResult { let port = vm.ctx.new_int(addr.port()); let ip = vm.ctx.new_str(addr.ip().to_string()); - let elements = RefCell::new(vec![ip, port]); - Ok(PyObject::new( - PyObjectPayload::Sequence { elements }, - vm.ctx.tuple_type(), - )) + Ok(vm.ctx.new_tuple(vec![ip, port])) } pub fn mk_module(ctx: &PyContext) -> PyObjectRef {