Merge pull request RustPython#1620 from RustPython/coolreader18/socket-fixes

Add some dns functions to socket

Author: windelbouwman

Cargo.lock

@@ -86,6 +86,7 @@ gethostname = "0.2.0"
 subprocess = "0.1.18"
 num_cpus = "1"
 socket2 = { version = "0.3", features = ["unix"] }
+dns-lookup = "1.0"
 
 [target."cfg(windows)".dependencies.winapi]
 version = "0.3"

vm/Cargo.toml

@@ -13,12 +13,15 @@ use super::os::convert_io_error;
 #[cfg(unix)]
 use super::os::convert_nix_error;
 use crate::function::{OptionalArg, PyFuncArgs};
+use crate::obj::objbytearray::PyByteArrayRef;
 use crate::obj::objbyteinner::PyBytesLike;
 use crate::obj::objbytes::PyBytesRef;
-use crate::obj::objstr::PyStringRef;
+use crate::obj::objstr::{PyString, PyStringRef};
 use crate::obj::objtuple::PyTupleRef;
 use crate::obj::objtype::PyClassRef;
-use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject};
+use crate::pyobject::{
+    Either, IntoPyObject, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+};
 use crate::vm::VirtualMachine;
 
 #[cfg(unix)]
@@ -172,14 +175,22 @@ impl PySocket {
     }
 
     #[pymethod]
-    fn recv(&self, bufsize: usize, vm: &VirtualMachine) -> PyResult {
+    fn recv(&self, bufsize: usize, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
         let mut buffer = vec![0u8; bufsize];
         match self.sock.borrow_mut().read_exact(&mut buffer) {
-            Ok(()) => Ok(vm.ctx.new_bytes(buffer)),
+            Ok(()) => Ok(buffer),
             Err(err) => Err(convert_sock_error(vm, err)),
         }
     }
 
+    #[pymethod]
+    fn recv_into(&self, buf: PyByteArrayRef, vm: &VirtualMachine) -> PyResult<usize> {
+        let mut buffer = buf.inner.borrow_mut();
+        self.sock()
+            .recv(&mut buffer.elements)
+            .map_err(|err| convert_sock_error(vm, err))
+    }
+
     #[pymethod]
     fn recvfrom(&self, bufsize: usize, vm: &VirtualMachine) -> PyResult<(Vec<u8>, AddrTuple)> {
         let mut buffer = vec![0u8; bufsize];
@@ -191,11 +202,20 @@ impl PySocket {
 
     #[pymethod]
     fn send(&self, bytes: PyBytesLike, vm: &VirtualMachine) -> PyResult<usize> {
+        // TODO: use PyBytesLike.with_ref() instead of to_cow()
         self.sock()
             .send(bytes.to_cow().as_ref())
             .map_err(|err| convert_sock_error(vm, err))
     }
 
+    #[pymethod]
+    fn sendall(&self, bytes: PyBytesLike, vm: &VirtualMachine) -> PyResult<()> {
+        self.sock
+            .borrow_mut()
+            .write_all(bytes.to_cow().as_ref())
+            .map_err(|err| convert_sock_error(vm, err))
+    }
+
     #[pymethod]
     fn sendto(&self, bytes: PyBytesLike, address: Address, vm: &VirtualMachine) -> PyResult<()> {
         let addr = get_addr(vm, address)?;
@@ -315,10 +335,14 @@ impl TryFromObject for Address {
         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())?,
-                port: u16::try_from_object(vm, tuple.elements[1].clone())?,
-            })
+            let host = PyStringRef::try_from_object(vm, tuple.elements[0].clone())?;
+            let host = if host.as_str().is_empty() {
+                PyString::from("0.0.0.0").into_ref(vm)
+            } else {
+                host
+            };
+            let port = u16::try_from_object(vm, tuple.elements[1].clone())?;
+            Ok(Address { host, port })
         }
     }
 }
@@ -368,6 +392,85 @@ fn socket_htonl(host: u32, vm: &VirtualMachine) -> PyResult {
     Ok(vm.new_int(host.to_be()))
 }
 
+#[derive(FromArgs)]
+struct GAIOptions {
+    #[pyarg(positional_only)]
+    host: Option<PyStringRef>,
+    #[pyarg(positional_only)]
+    port: Option<Either<PyStringRef, i32>>,
+
+    #[pyarg(positional_only, default = "0")]
+    family: i32,
+    #[pyarg(positional_only, default = "0")]
+    ty: i32,
+    #[pyarg(positional_only, default = "0")]
+    proto: i32,
+    #[pyarg(positional_only, default = "0")]
+    flags: i32,
+}
+
+fn socket_getaddrinfo(opts: GAIOptions, vm: &VirtualMachine) -> PyResult {
+    let hints = dns_lookup::AddrInfoHints {
+        socktype: opts.ty,
+        protocol: opts.proto,
+        address: opts.family,
+        flags: opts.flags,
+    };
+
+    let host = opts.host.as_ref().map(|s| s.as_str());
+    let port = opts.port.as_ref().map(|p| -> std::borrow::Cow<str> {
+        match p {
+            Either::A(ref s) => s.as_str().into(),
+            Either::B(i) => i.to_string().into(),
+        }
+    });
+    let port = port.as_ref().map(|p| p.as_ref());
+
+    let addrs = dns_lookup::getaddrinfo(host, port, Some(hints)).map_err(|err| {
+        let error_type = vm.class("_socket", "gaierror");
+        vm.new_exception(error_type, io::Error::from(err).to_string())
+    })?;
+
+    let list = addrs
+        .map(|ai| {
+            ai.map(|ai| {
+                vm.ctx.new_tuple(vec![
+                    vm.new_int(ai.address),
+                    vm.new_int(ai.socktype),
+                    vm.new_int(ai.protocol),
+                    match ai.canonname {
+                        Some(s) => vm.new_str(s),
+                        None => vm.get_none(),
+                    },
+                    get_addr_tuple(ai.sockaddr).into_pyobject(vm).unwrap(),
+                ])
+            })
+        })
+        .collect::<io::Result<Vec<_>>>()
+        .map_err(|e| convert_sock_error(vm, e))?;
+    Ok(vm.ctx.new_list(list))
+}
+
+fn socket_gethostbyaddr(
+    addr: PyStringRef,
+    vm: &VirtualMachine,
+) -> PyResult<(String, PyObjectRef, PyObjectRef)> {
+    // TODO: figure out how to do this properly
+    let ai = dns_lookup::getaddrinfo(Some(addr.as_str()), None, None)
+        .map_err(|e| convert_sock_error(vm, e.into()))?
+        .next()
+        .unwrap()
+        .map_err(|e| convert_sock_error(vm, e))?;
+    let (hostname, _) =
+        dns_lookup::getnameinfo(&ai.sockaddr, 0).map_err(|e| convert_sock_error(vm, e.into()))?;
+    Ok((
+        hostname,
+        vm.ctx.new_list(vec![]),
+        vm.ctx
+            .new_list(vec![vm.new_str(ai.sockaddr.ip().to_string())]),
+    ))
+}
+
 fn get_addr<T, I>(vm: &VirtualMachine, addr: T) -> PyResult<socket2::SockAddr>
 where
     T: ToSocketAddrs<Iter = I>,
@@ -467,6 +570,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "gethostname" => ctx.new_rustfunc(socket_gethostname),
         "htonl" => ctx.new_rustfunc(socket_htonl),
         "getdefaulttimeout" => ctx.new_rustfunc(|vm: &VirtualMachine| vm.get_none()),
+        "getaddrinfo" => ctx.new_rustfunc(socket_getaddrinfo),
+        "gethostbyaddr" => ctx.new_rustfunc(socket_gethostbyaddr),
     });
 
     extend_module_platform_specific(vm, &module);