Merge pull request RustPython#1617 from RustPython/coolreader18/os-stuff

windelbouwman · web-flow · commit 63f0d27b158b · 2019-12-21T12:31:59.000+01:00
Fix some various issues with `os` and `binascii`
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/tests/snippets/stdlib_binascii.py b/tests/snippets/stdlib_binascii.py
@@ -1,5 +1,5 @@
 import binascii
-from testutils import assert_raises
+from testutils import assert_raises, assert_equal
 
 
 # hexlify tests
@@ -25,8 +25,7 @@
 assert uh(b"ABCDEF") == b"\xab\xcd\xef"
 assert binascii.a2b_hex(b"6161") == b"aa"
 
-# unhexlify on strings not supported yet
-# assert uh("abcd") == b"\xab\xcd"
+assert uh("abcd") == b"\xab\xcd"
 
 with assert_raises(ValueError):
     uh(b"a")  # Odd-length string
@@ -36,4 +35,18 @@
 
 assert binascii.crc32(b"hello world") == 222957957
 assert binascii.crc32(b"hello world", 555555) == 1216827162
-assert binascii.crc32(b"goodbye interesting world",777777) == 1885538403
+assert binascii.crc32(b"goodbye interesting world", 777777) == 1885538403
+
+enc_b64 = binascii.b2a_base64
+
+assert_equal(enc_b64(b"213"), b"MjEz\n")
+assert_equal(enc_b64(b"213", newline=False), b"MjEz")
+
+dec_b64 = binascii.a2b_base64
+
+assert_equal(dec_b64(b"UnVzdFB5dGhvbg==\n"), b"RustPython")
+assert_equal(
+    dec_b64(b"4pii8J+QoyAg4ZaH8J2TpPCdlYrRguKTn/CdlZDwnZWl5Y2Ez4PwnZSrICDimazwn5Gj\n"),
+    "☢🐣  ᖇ𝓤𝕊тⓟ𝕐𝕥卄σ𝔫  ♬👣".encode(),
+)
+
diff --git a/tests/snippets/stdlib_os.py b/tests/snippets/stdlib_os.py
@@ -125,6 +125,17 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 	assert os.read(fd, len(CONTENT2) + len(CONTENT3)) == CONTENT2 + CONTENT3
 	os.close(fd)
 
+	assert not os.isatty(fd)
+
+  # TODO: get os.lseek working on windows
+	if os.name != 'nt':
+		fd = os.open(fname3, 0)
+		assert os.read(fd, len(CONTENT2)) == CONTENT2
+		assert os.read(fd, len(CONTENT3)) == CONTENT3
+		os.lseek(fd, len(CONTENT2), os.SEEK_SET)
+		assert os.read(fd, len(CONTENT3)) == CONTENT3
+		os.close(fd)
+
 	os.rename(fname3, fname)
 	assert os.path.exists(fname3) == False
 	assert os.path.exists(fname) == True
@@ -383,11 +394,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     assert os.system('{') != 0
 
     for arg in [None, 1, 1.0, TabError]:
-        try:
-            os.system(arg)
-        except TypeError:
-            pass
-        else:
-            raise AssertionError(f'os.system failed to raise TypeError with arg {arg}')
+        assert_raises(TypeError, os.system, arg)
 
 
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -67,6 +67,7 @@ wtf8 = "0.0.3"
 arr_macro = "0.1.2"
 csv = "1.1.1"
 paste = "0.1"
+base64 = "0.11"
 
 flame = { version = "0.2", optional = true }
 flamer = { version = "0.3", optional = true }
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -1418,16 +1418,13 @@ fn split_slice_reverse<'a>(slice: &'a [u8], sep: &[u8], maxsplit: i32) -> Vec<&'
 pub enum PyBytesLike {
     Bytes(PyBytesRef),
     Bytearray(PyByteArrayRef),
-    Vec(Vec<u8>),
 }
 
 impl TryFromObject for PyBytesLike {
     fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
         match_class!(match obj {
             b @ PyBytes => Ok(PyBytesLike::Bytes(b)),
             b @ PyByteArray => Ok(PyBytesLike::Bytearray(b)),
-            m @ PyMemoryView => Ok(PyBytesLike::Vec(m.get_obj_value().unwrap())),
-            l @ PyList => Ok(PyBytesLike::Vec(l.get_byte_inner(vm)?.elements)),
             obj => Err(vm.new_type_error(format!(
                 "a bytes-like object is required, not {}",
                 obj.class()
@@ -1441,7 +1438,14 @@ impl PyBytesLike {
         match self {
             PyBytesLike::Bytes(b) => b.get_value().into(),
             PyBytesLike::Bytearray(b) => b.inner.borrow().elements.clone().into(),
-            PyBytesLike::Vec(b) => b.as_slice().into(),
+        }
+    }
+
+    #[inline]
+    pub fn with_ref<R>(&self, f: impl FnOnce(&[u8]) -> R) -> R {
+        match self {
+            PyBytesLike::Bytes(b) => f(b.get_value()),
+            PyBytesLike::Bytearray(b) => f(&b.inner.borrow().elements),
         }
     }
 }
diff --git a/vm/src/stdlib/binascii.rs b/vm/src/stdlib/binascii.rs
@@ -1,8 +1,52 @@
 use crate::function::OptionalArg;
-use crate::obj::objbytes::PyBytesRef;
-use crate::pyobject::{PyObjectRef, PyResult};
+use crate::obj::objbytearray::{PyByteArray, PyByteArrayRef};
+use crate::obj::objbyteinner::PyBytesLike;
+use crate::obj::objbytes::{PyBytes, PyBytesRef};
+use crate::obj::objstr::{PyString, PyStringRef};
+use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
 use crate::vm::VirtualMachine;
+
 use crc::{crc32, Hasher32};
+use itertools::Itertools;
+
+enum SerializedData {
+    Bytes(PyBytesRef),
+    Buffer(PyByteArrayRef),
+    Ascii(PyStringRef),
+}
+
+impl TryFromObject for SerializedData {
+    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+        match_class!(match obj {
+            b @ PyBytes => Ok(SerializedData::Bytes(b)),
+            b @ PyByteArray => Ok(SerializedData::Buffer(b)),
+            a @ PyString => {
+                if a.as_str().is_ascii() {
+                    Ok(SerializedData::Ascii(a))
+                } else {
+                    Err(vm.new_value_error(
+                        "string argument should contain only ASCII characters".to_string(),
+                    ))
+                }
+            }
+            obj => Err(vm.new_type_error(format!(
+                "argument should be bytes, buffer or ASCII string, not '{}'",
+                obj.class().name,
+            ))),
+        })
+    }
+}
+
+impl SerializedData {
+    #[inline]
+    pub fn with_ref<R>(&self, f: impl FnOnce(&[u8]) -> R) -> R {
+        match self {
+            SerializedData::Bytes(b) => f(b.get_value()),
+            SerializedData::Buffer(b) => f(&b.inner.borrow().elements),
+            SerializedData::Ascii(a) => f(a.as_str().as_bytes()),
+        }
+    }
+}
 
 fn hex_nibble(n: u8) -> u8 {
     match n {
@@ -12,15 +56,15 @@ fn hex_nibble(n: u8) -> u8 {
     }
 }
 
-fn binascii_hexlify(data: PyBytesRef, vm: &VirtualMachine) -> PyResult {
-    let bytes = data.get_value();
-    let mut hex = Vec::<u8>::with_capacity(bytes.len() * 2);
-    for b in bytes.iter() {
-        hex.push(hex_nibble(b >> 4));
-        hex.push(hex_nibble(b & 0xf));
-    }
-
-    Ok(vm.ctx.new_bytes(hex))
+fn binascii_hexlify(data: PyBytesLike, _vm: &VirtualMachine) -> Vec<u8> {
+    data.with_ref(|bytes| {
+        let mut hex = Vec::<u8>::with_capacity(bytes.len() * 2);
+        for b in bytes.iter() {
+            hex.push(hex_nibble(b >> 4));
+            hex.push(hex_nibble(b & 0xf));
+        }
+        hex
+    })
 }
 
 fn unhex_nibble(c: u8) -> Option<u8> {
@@ -32,37 +76,66 @@ fn unhex_nibble(c: u8) -> Option<u8> {
     }
 }
 
-fn binascii_unhexlify(hexstr: PyBytesRef, vm: &VirtualMachine) -> PyResult {
-    // TODO: allow 'str' hexstrings as well
-    let hex_bytes = hexstr.get_value();
-    if hex_bytes.len() % 2 != 0 {
-        return Err(vm.new_value_error("Odd-length string".to_string()));
-    }
+fn binascii_unhexlify(data: SerializedData, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+    data.with_ref(|hex_bytes| {
+        if hex_bytes.len() % 2 != 0 {
+            return Err(vm.new_value_error("Odd-length string".to_string()));
+        }
 
-    let mut unhex = Vec::<u8>::with_capacity(hex_bytes.len() / 2);
-    for i in (0..hex_bytes.len()).step_by(2) {
-        let n1 = unhex_nibble(hex_bytes[i]);
-        let n2 = unhex_nibble(hex_bytes[i + 1]);
-        if let (Some(n1), Some(n2)) = (n1, n2) {
-            unhex.push(n1 << 4 | n2);
-        } else {
-            return Err(vm.new_value_error("Non-hexadecimal digit found".to_string()));
+        let mut unhex = Vec::<u8>::with_capacity(hex_bytes.len() / 2);
+        for (n1, n2) in hex_bytes.iter().tuples() {
+            if let (Some(n1), Some(n2)) = (unhex_nibble(*n1), unhex_nibble(*n2)) {
+                unhex.push(n1 << 4 | n2);
+            } else {
+                return Err(vm.new_value_error("Non-hexadecimal digit found".to_string()));
+            }
         }
-    }
 
-    Ok(vm.ctx.new_bytes(unhex))
+        Ok(unhex)
+    })
 }
 
-fn binascii_crc32(data: PyBytesRef, value: OptionalArg<u32>, vm: &VirtualMachine) -> PyResult {
-    let bytes = data.get_value();
-    let crc = value.unwrap_or(0u32);
+fn binascii_crc32(data: SerializedData, value: OptionalArg<u32>, vm: &VirtualMachine) -> PyResult {
+    let crc = value.unwrap_or(0);
 
     let mut digest = crc32::Digest::new_with_initial(crc32::IEEE, crc);
-    digest.write(&bytes);
+    data.with_ref(|bytes| digest.write(&bytes));
 
     Ok(vm.ctx.new_int(digest.sum32()))
 }
 
+#[derive(FromArgs)]
+struct NewlineArg {
+    #[pyarg(keyword_only, default = "true")]
+    newline: bool,
+}
+
+/// trim a newline from the end of the bytestring, if it exists
+fn trim_newline(b: &[u8]) -> &[u8] {
+    if b.ends_with(b"\n") {
+        &b[..b.len() - 1]
+    } else {
+        b
+    }
+}
+
+fn binascii_a2b_base64(s: SerializedData, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+    s.with_ref(|b| base64::decode(trim_newline(b)))
+        .map_err(|err| vm.new_value_error(format!("error decoding base64: {}", err)))
+}
+
+fn binascii_b2a_base64(
+    data: PyBytesLike,
+    NewlineArg { newline }: NewlineArg,
+    _vm: &VirtualMachine,
+) -> Vec<u8> {
+    let mut encoded = data.with_ref(base64::encode).into_bytes();
+    if newline {
+        encoded.push(b'\n');
+    }
+    encoded
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -72,5 +145,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "unhexlify" => ctx.new_rustfunc(binascii_unhexlify),
         "a2b_hex" => ctx.new_rustfunc(binascii_unhexlify),
         "crc32" => ctx.new_rustfunc(binascii_crc32),
+        "a2b_base64" => ctx.new_rustfunc(binascii_a2b_base64),
+        "b2a_base64" => ctx.new_rustfunc(binascii_b2a_base64),
     })
 }
diff --git a/vm/src/stdlib/errno.rs b/vm/src/stdlib/errno.rs
@@ -169,6 +169,7 @@ const ERROR_CODES: &[(&str, i32)] = &[
     e!(EINPROGRESS),
     e!(ENXIO),
     e!(ECANCELED),
+    e!(EWOULDBLOCK),
     e!(cfg(not(windows)), EOWNERDEAD),
     e!(cfg(not(windows)), ENOTRECOVERABLE),
     e!(cfg(windows), WSAEAFNOSUPPORT),
diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -244,7 +244,8 @@ impl PyBytesIORef {
     }
 
     fn write(self, data: PyBytesLike, vm: &VirtualMachine) -> PyResult<u64> {
-        match self.buffer(vm)?.write(data.to_cow().as_ref()) {
+        let mut buffer = self.buffer(vm)?;
+        match data.with_ref(|b| buffer.write(b)) {
             Some(value) => Ok(value),
             None => Err(vm.new_type_error("Error Writing Bytes".to_string())),
         }
@@ -358,11 +359,10 @@ fn io_base_readline(
     let read = vm.get_attribute(instance, "read")?;
     while size < 0 || res.len() < size as usize {
         let read_res = PyBytesLike::try_from_object(vm, vm.invoke(&read, vec![vm.new_int(1)])?)?;
-        let b = read_res.to_cow();
-        if b.is_empty() {
+        if read_res.with_ref(|b| b.is_empty()) {
             break;
         }
-        res.extend_from_slice(b.as_ref());
+        read_res.with_ref(|b| res.extend_from_slice(b));
         if res.ends_with(b"\n") {
             break;
         }
@@ -643,8 +643,8 @@ mod fileio {
     ) -> PyResult<usize> {
         let mut handle = fio_get_fileno(&instance, vm)?;
 
-        let len = handle
-            .write(obj.to_cow().as_ref())
+        let len = obj
+            .with_ref(|b| handle.write(b))
             .map_err(|e| os::convert_io_error(vm, e))?;
 
         fio_set_fileno(&instance, handle, vm)?;
diff --git a/vm/src/stdlib/multiprocessing.rs b/vm/src/stdlib/multiprocessing.rs
@@ -11,10 +11,7 @@ use winapi::um::winsock2::{self, SOCKET};
 fn multiprocessing_closesocket(socket: usize, vm: &VirtualMachine) -> PyResult<()> {
     let res = unsafe { winsock2::closesocket(socket as SOCKET) };
     if res == 0 {
-        Err(super::os::convert_io_error(
-            vm,
-            std::io::Error::last_os_error(),
-        ))
+        Err(super::os::errno_err(vm))
     } else {
         Ok(())
     }
@@ -26,10 +23,7 @@ fn multiprocessing_recv(socket: usize, size: usize, vm: &VirtualMachine) -> PyRe
     let nread =
         unsafe { winsock2::recv(socket as SOCKET, buf.as_mut_ptr() as *mut _, size as i32, 0) };
     if nread < 0 {
-        Err(super::os::convert_io_error(
-            vm,
-            std::io::Error::last_os_error(),
-        ))
+        Err(super::os::errno_err(vm))
     } else {
         Ok(nread)
     }
@@ -41,20 +35,11 @@ fn multiprocessing_send(
     buf: PyBytesLike,
     vm: &VirtualMachine,
 ) -> PyResult<libc::c_int> {
-    let buf = buf.to_cow();
-    let ret = unsafe {
-        winsock2::send(
-            socket as SOCKET,
-            buf.as_ptr() as *const _,
-            buf.len() as i32,
-            0,
-        )
-    };
+    let ret = buf.with_ref(|b| unsafe {
+        winsock2::send(socket as SOCKET, b.as_ptr() as *const _, b.len() as i32, 0)
+    });
     if ret < 0 {
-        Err(super::os::convert_io_error(
-            vm,
-            std::io::Error::last_os_error(),
-        ))
+        Err(super::os::errno_err(vm))
     } else {
         Ok(ret)
     }
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
diff --git a/vm/src/stdlib/winapi.rs b/vm/src/stdlib/winapi.rs
