Update binascii with base64, add with_ref method to PyBytesLike

coolreader18 · coolreader18 · commit 69658bd215fd · 2019-12-14T15:48:48.000-06:00
diff --git a/Cargo.lock b/Cargo.lock
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/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
@@ -41,15 +41,9 @@ 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,
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
@@ -204,25 +204,23 @@ 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())
+        bytes
+            .with_ref(|b| self.sock().send(b))
             .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())
+        bytes
+            .with_ref(|b| self.sock.borrow_mut().write_all(b))
             .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)?;
-        self.sock()
-            .send_to(bytes.to_cow().as_ref(), &addr)
+        bytes
+            .with_ref(|b| self.sock().send_to(b, &addr))
             .map_err(|err| convert_sock_error(vm, err))?;
         Ok(())
     }
