use std::cell::RefCell;
use std::io;
use std::io::Read;
use std::io::Write;
use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket};
use crate::obj::objbytes::PyBytesRef;
use crate::obj::objint::PyIntRef;
use crate::obj::objstr::PyStringRef;
use crate::obj::objtuple::PyTupleRef;
use crate::pyobject::{PyObjectRef, PyRef, PyResult, PyValue, TryFromObject};
use crate::vm::VirtualMachine;
use crate::obj::objtype::PyClassRef;
use num_traits::ToPrimitive;
#[derive(Debug, Copy, Clone)]
enum AddressFamily {
Unix = 1,
Inet = 2,
Inet6 = 3,
}
impl TryFromObject for AddressFamily {
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
match i32::try_from_object(vm, obj)? {
1 => Ok(AddressFamily::Unix),
2 => Ok(AddressFamily::Inet),
3 => Ok(AddressFamily::Inet6),
value => Err(vm.new_os_error(format!("Unknown address family value: {}", value))),
}
}
}
#[derive(Debug, Copy, Clone)]
enum SocketKind {
Stream = 1,
Dgram = 2,
}
impl TryFromObject for SocketKind {
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
match i32::try_from_object(vm, obj)? {
1 => Ok(SocketKind::Stream),
2 => Ok(SocketKind::Dgram),
value => Err(vm.new_os_error(format!("Unknown socket kind value: {}", value))),
}
}
}
#[derive(Debug)]
enum Connection {
TcpListener(TcpListener),
TcpStream(TcpStream),
UdpSocket(UdpSocket),
}
impl Connection {
fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)> {
match self {
Connection::TcpListener(con) => con.accept(),
_ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")),
}
}
fn local_addr(&self) -> io::Result<SocketAddr> {
match self {
Connection::TcpListener(con) => con.local_addr(),
Connection::UdpSocket(con) => con.local_addr(),
Connection::TcpStream(con) => con.local_addr(),
}
}
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!")),
}
}
fn send_to<A: ToSocketAddrs>(&self, buf: &[u8], addr: A) -> io::Result<usize> {
match self {
Connection::UdpSocket(con) => con.send_to(buf, addr),
_ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")),
}
}
#[cfg(unix)]
fn fileno(&self) -> i64 {
use std::os::unix::io::AsRawFd;
let raw_fd = match self {
Connection::TcpListener(con) => con.as_raw_fd(),
Connection::UdpSocket(con) => con.as_raw_fd(),
Connection::TcpStream(con) => con.as_raw_fd(),
};
i64::from(raw_fd)
}
#[cfg(windows)]
fn fileno(&self) -> i64 {
use std::os::windows::io::AsRawSocket;
let raw_fd = match self {
Connection::TcpListener(con) => con.as_raw_socket(),
Connection::UdpSocket(con) => con.as_raw_socket(),
Connection::TcpStream(con) => con.as_raw_socket(),
};
raw_fd as i64
}
#[cfg(all(not(unix), not(windows)))]
fn fileno(&self) -> i64 {
unimplemented!();
}
}
impl Read for Connection {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
Connection::TcpStream(con) => con.read(buf),
Connection::UdpSocket(con) => con.recv(buf),
_ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")),
}
}
}
impl Write for Connection {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
match self {
Connection::TcpStream(con) => con.write(buf),
Connection::UdpSocket(con) => con.send(buf),
_ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")),
}
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[derive(Debug)]
pub struct Socket {
address_family: AddressFamily,
socket_kind: SocketKind,
con: RefCell<Option<Connection>>,
}
impl PyValue for Socket {
fn class(vm: &VirtualMachine) -> PyClassRef {
vm.class("socket", "socket")
}
}
impl Socket {
fn new(address_family: AddressFamily, socket_kind: SocketKind) -> Socket {
Socket {
address_family,
socket_kind,
con: RefCell::new(None),
}
}
}
type SocketRef = PyRef<Socket>;
impl SocketRef {
fn new(
cls: PyClassRef,
family: AddressFamily,
kind: SocketKind,
vm: &VirtualMachine,
) -> PyResult<SocketRef> {
Socket::new(family, kind).into_ref_with_type(vm, cls)
}
fn connect(self, address: Address, vm: &VirtualMachine) -> PyResult<()> {
let address_string = address.get_address_string();
match self.socket_kind {
SocketKind::Stream => match TcpStream::connect(address_string) {
Ok(stream) => {
self.con.borrow_mut().replace(Connection::TcpStream(stream));
Ok(())
}
Err(s) => Err(vm.new_os_error(s.to_string())),
},
SocketKind::Dgram => {
if let Some(Connection::UdpSocket(con)) = self.con.borrow().as_ref() {
match con.connect(address_string) {
Ok(_) => Ok(()),
Err(s) => Err(vm.new_os_error(s.to_string())),
}
} else {
Err(vm.new_type_error("".to_string()))
}
}
}
}
fn bind(self, address: Address, vm: &VirtualMachine) -> PyResult<()> {
let address_string = address.get_address_string();
match self.socket_kind {
SocketKind::Stream => match TcpListener::bind(address_string) {
Ok(stream) => {
self.con
.borrow_mut()
.replace(Connection::TcpListener(stream));
Ok(())
}
Err(s) => Err(vm.new_os_error(s.to_string())),
},
SocketKind::Dgram => match UdpSocket::bind(address_string) {
Ok(dgram) => {
self.con.borrow_mut().replace(Connection::UdpSocket(dgram));
Ok(())
}
Err(s) => Err(vm.new_os_error(s.to_string())),
},
}
}
fn listen(self, _num: PyIntRef, _vm: &VirtualMachine) {}
fn accept(self, vm: &VirtualMachine) -> PyResult {
let ret = match self.con.borrow_mut().as_mut() {
Some(v) => v.accept(),
None => return Err(vm.new_type_error("".to_string())),
};
let (tcp_stream, addr) = match ret {
Ok((socket, addr)) => (socket, addr),
Err(s) => return Err(vm.new_os_error(s.to_string())),
};
let socket = Socket {
address_family: self.address_family,
socket_kind: self.socket_kind,
con: RefCell::new(Some(Connection::TcpStream(tcp_stream))),
}
.into_ref(vm);
let addr_tuple = get_addr_tuple(vm, addr)?;
Ok(vm.ctx.new_tuple(vec![socket.into_object(), addr_tuple]))
}
fn recv(self, bufsize: PyIntRef, vm: &VirtualMachine) -> PyResult {
let mut buffer = vec![0u8; bufsize.as_bigint().to_usize().unwrap()];
match self.con.borrow_mut().as_mut() {
Some(v) => match v.read_exact(&mut buffer) {
Ok(_) => (),
Err(s) => return Err(vm.new_os_error(s.to_string())),
},
None => return Err(vm.new_type_error("".to_string())),
};
Ok(vm.ctx.new_bytes(buffer))
}
fn recvfrom(self, bufsize: PyIntRef, vm: &VirtualMachine) -> PyResult {
let mut buffer = vec![0u8; bufsize.as_bigint().to_usize().unwrap()];
let ret = match self.con.borrow().as_ref() {
Some(v) => v.recv_from(&mut buffer),
None => return Err(vm.new_type_error("".to_string())),
};
let addr = match ret {
Ok((_size, addr)) => addr,
Err(s) => return Err(vm.new_os_error(s.to_string())),
};
let addr_tuple = get_addr_tuple(vm, addr)?;
Ok(vm.ctx.new_tuple(vec![vm.ctx.new_bytes(buffer), addr_tuple]))
}
fn send(self, bytes: PyBytesRef, vm: &VirtualMachine) -> PyResult<()> {
match self.con.borrow_mut().as_mut() {
Some(v) => match v.write(&bytes) {
Ok(_) => (),
Err(s) => return Err(vm.new_os_error(s.to_string())),
},
None => return Err(vm.new_type_error("".to_string())),
};
Ok(())
}
fn sendto(self, bytes: PyBytesRef, address: Address, vm: &VirtualMachine) -> PyResult<()> {
let address_string = address.get_address_string();
match self.socket_kind {
SocketKind::Dgram => {
if let Some(v) = self.con.borrow().as_ref() {
return match v.send_to(&bytes, address_string) {
Ok(_) => Ok(()),
Err(s) => Err(vm.new_os_error(s.to_string())),
};
}
match UdpSocket::bind("0.0.0.0:0") {
Ok(dgram) => match dgram.send_to(&bytes, address_string) {
Ok(_) => {
self.con.borrow_mut().replace(Connection::UdpSocket(dgram));
Ok(())
}
Err(s) => Err(vm.new_os_error(s.to_string())),
},
Err(s) => Err(vm.new_os_error(s.to_string())),
}
}
_ => Err(vm.new_not_implemented_error("".to_string())),
}
}
fn close(self, _vm: &VirtualMachine) {
self.con.borrow_mut().take();
}
fn fileno(self, vm: &VirtualMachine) -> PyResult {
let fileno = match self.con.borrow_mut().as_mut() {
Some(v) => v.fileno(),
None => return Err(vm.new_type_error("".to_string())),
};
Ok(vm.ctx.new_int(fileno))
}
fn getsockname(self, vm: &VirtualMachine) -> PyResult {
let addr = match self.con.borrow().as_ref() {
Some(v) => v.local_addr(),
None => return Err(vm.new_type_error("".to_string())),
};
match addr {
Ok(addr) => get_addr_tuple(vm, addr),
Err(s) => Err(vm.new_os_error(s.to_string())),
}
}
}
struct Address {
host: String,
port: usize,
}
impl Address {
fn get_address_string(self) -> String {
format!("{}:{}", self.host, self.port.to_string())
}
}
impl TryFromObject for Address {
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
let tuple = PyTupleRef::try_from_object(vm, obj)?;
if tuple.elements.len() != 2 {
Err(vm.new_type_error("Address tuple should have only 2 values".to_string()))
} else {
Ok(Address {
host: PyStringRef::try_from_object(vm, tuple.elements[0].clone())?
.value
.to_string(),
port: PyIntRef::try_from_object(vm, tuple.elements[1].clone())?
.as_bigint()
.to_usize()
.unwrap(),
})
}
}
}
fn get_addr_tuple(vm: &VirtualMachine, addr: SocketAddr) -> PyResult {
let port = vm.ctx.new_int(addr.port());
let ip = vm.ctx.new_str(addr.ip().to_string());
Ok(vm.ctx.new_tuple(vec![ip, port]))
}
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
let ctx = &vm.ctx;
let socket = py_class!(ctx, "socket", ctx.object(), {
"__new__" => ctx.new_rustfunc(SocketRef::new),
"connect" => ctx.new_rustfunc(SocketRef::connect),
"recv" => ctx.new_rustfunc(SocketRef::recv),
"send" => ctx.new_rustfunc(SocketRef::send),
"bind" => ctx.new_rustfunc(SocketRef::bind),
"accept" => ctx.new_rustfunc(SocketRef::accept),
"listen" => ctx.new_rustfunc(SocketRef::listen),
"close" => ctx.new_rustfunc(SocketRef::close),
"getsockname" => ctx.new_rustfunc(SocketRef::getsockname),
"sendto" => ctx.new_rustfunc(SocketRef::sendto),
"recvfrom" => ctx.new_rustfunc(SocketRef::recvfrom),
"fileno" => ctx.new_rustfunc(SocketRef::fileno),
});
py_module!(vm, "socket", {
"AF_INET" => ctx.new_int(AddressFamily::Inet as i32),
"SOCK_STREAM" => ctx.new_int(SocketKind::Stream as i32),
"SOCK_DGRAM" => ctx.new_int(SocketKind::Dgram as i32),
"socket" => socket,
})
}