Do not use Vec<u8> as argument type in the binding

See PyO3/pyo3#2888
Scille · Jan 20, 2023 · 4f02a4b · 4f02a4b
1 parent 3de9817
commit 4f02a4b
Show file tree

Hide file tree

Showing 15 changed files with 251 additions and 83 deletions.
diff --git a/newsfragments/3925.misc.rst b/newsfragments/3925.misc.rst
@@ -0,0 +1 @@
+Fix an underlying issue that caused an overall performance degradation
diff --git a/src/addrs.rs b/src/addrs.rs
@@ -10,6 +10,7 @@ use std::str::FromStr;
 
 use crate::{
     api_crypto::VerifyKey,
+    binding_utils::BytesWrapper,
     enumerate::InvitationType,
     ids::{EntryID, InvitationToken, OrganizationID},
 };
@@ -447,10 +448,11 @@ impl BackendOrganizationFileLinkAddr {
     fn new(
         organization_id: OrganizationID,
         workspace_id: EntryID,
-        encrypted_path: Vec<u8>,
-        encrypted_timestamp: Option<Vec<u8>>,
+        encrypted_path: BytesWrapper,
+        encrypted_timestamp: Option<BytesWrapper>,
         py_kwargs: Option<&PyDict>,
     ) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(encrypted_path, encrypted_timestamp);
         let addr = match py_kwargs {
             Some(dict) => BackendAddr::new(
                 match dict.get_item("hostname") {
@@ -568,9 +570,10 @@ impl BackendOrganizationFileLinkAddr {
         _cls: &PyType,
         organization_addr: BackendOrganizationAddr,
         workspace_id: EntryID,
-        encrypted_path: Vec<u8>,
-        encrypted_timestamp: Option<Vec<u8>>,
+        encrypted_path: BytesWrapper,
+        encrypted_timestamp: Option<BytesWrapper>,
     ) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(encrypted_path, encrypted_timestamp);
         Ok(Self(
             libparsec::types::BackendOrganizationFileLinkAddr::new(
                 organization_addr.get_backend_addr().0,

diff --git a/src/binding_utils.rs b/src/binding_utils.rs
@@ -1,14 +1,80 @@
 // Parsec Cloud (https://parsec.cloud) Copyright (c) BUSL-1.1 (eventually AGPL-3.0) 2016-present Scille SAS
 
 use pyo3::{
-    exceptions::PyNotImplementedError, pyclass::CompareOp, types::PyFrozenSet, FromPyObject, PyAny,
-    PyResult,
+    exceptions::PyNotImplementedError, pyclass::CompareOp, types::PyByteArray, types::PyBytes,
+    types::PyFrozenSet, FromPyObject, PyAny, PyResult,
 };
 use std::{
     collections::{hash_map::DefaultHasher, HashSet},
     hash::{Hash, Hasher},
 };
 
+#[derive(FromPyObject)]
+pub enum BytesWrapper<'py> {
+    Bytes(&'py PyBytes),
+    ByteArray(&'py PyByteArray),
+}
+
+impl From<BytesWrapper<'_>> for Vec<u8> {
+    fn from(wrapper: BytesWrapper) -> Self {
+        match wrapper {
+            BytesWrapper::Bytes(bytes) => bytes.as_bytes().to_vec(),
+            BytesWrapper::ByteArray(byte_array) => byte_array.to_vec(),
+        }
+    }
+}
+
+pub trait UnwrapBytesWrapper {
+    type ResultType;
+    fn unwrap_bytes(self) -> Self::ResultType;
+}
+
+impl UnwrapBytesWrapper for BytesWrapper<'_> {
+    type ResultType = Vec<u8>;
+    fn unwrap_bytes(self) -> Self::ResultType {
+        self.into()
+    }
+}
+
+impl<T> UnwrapBytesWrapper for Option<T>
+where
+    T: UnwrapBytesWrapper,
+{
+    type ResultType = Option<T::ResultType>;
+    fn unwrap_bytes(self) -> Self::ResultType {
+        self.map(|x| x.unwrap_bytes())
+    }
+}
+
+impl<T> UnwrapBytesWrapper for Vec<T>
+where
+    T: UnwrapBytesWrapper,
+{
+    type ResultType = Vec<T::ResultType>;
+    fn unwrap_bytes(self) -> Self::ResultType {
+        self.into_iter().map(|x| x.unwrap_bytes()).collect()
+    }
+}
+
+macro_rules! _unwrap_bytes {
+    ($name:ident) => {
+        let $name = $name.unwrap_bytes();
+    };
+    ($x:ident, $($y:ident),+) => {
+        crate::binding_utils::_unwrap_bytes!($x);
+        crate::binding_utils::_unwrap_bytes!($($y),+);
+
+    }
+}
+
+macro_rules! unwrap_bytes {
+    ($($name:ident),+) => {
+        use crate::binding_utils::UnwrapBytesWrapper;
+        crate::binding_utils::_unwrap_bytes!($($name),+);
+
+    }
+}
+
 pub fn comp_eq<T: std::cmp::PartialEq>(op: CompareOp, h1: T, h2: T) -> PyResult<bool> {
     Ok(match op {
         CompareOp::Eq => h1 == h2,
@@ -233,9 +299,11 @@ macro_rules! impl_enum_field {
     };
 }
 
+pub(crate) use _unwrap_bytes;
 pub(crate) use create_exception;
 pub(crate) use gen_proto;
 pub(crate) use impl_enum_field;
 pub(crate) use parse_kwargs;
 pub(crate) use parse_kwargs_optional;
 pub(crate) use py_object;
+pub(crate) use unwrap_bytes;
diff --git a/src/data/organization.rs b/src/data/organization.rs
@@ -3,7 +3,9 @@ use pyo3::{
     types::{PyBytes, PyTuple},
 };
 
-use crate::{data::UsersPerProfileDetailItem, protocol::ActiveUsersLimit};
+use crate::{
+    binding_utils::BytesWrapper, data::UsersPerProfileDetailItem, protocol::ActiveUsersLimit,
+};
 
 #[pyclass]
 pub(crate) struct OrganizationStats(pub libparsec::types::OrganizationStats);
@@ -81,9 +83,10 @@ impl OrganizationConfig {
     fn new(
         user_profile_outsider_allowed: bool,
         active_users_limit: ActiveUsersLimit,
-        sequester_authority: Option<Vec<u8>>,
-        sequester_services: Option<Vec<Vec<u8>>>,
+        sequester_authority: Option<BytesWrapper>,
+        sequester_services: Option<Vec<BytesWrapper>>,
     ) -> Self {
+        crate::binding_utils::unwrap_bytes!(sequester_authority, sequester_services);
         Self(libparsec::types::OrganizationConfig {
             user_profile_outsider_allowed,
             active_users_limit: active_users_limit.0,

diff --git a/src/data/pki.rs b/src/data/pki.rs
@@ -8,6 +8,7 @@ use std::{collections::HashMap, path::Path};
 use crate::{
     addrs::BackendPkiEnrollmentAddr,
     api_crypto::{PublicKey, VerifyKey},
+    binding_utils::BytesWrapper,
     data::{DataResult, PkiEnrollmentLocalPendingResult},
     enumerate::UserProfile,
     ids::{DeviceID, DeviceLabel, EnrollmentID, HumanHandle},
@@ -135,10 +136,11 @@ impl X509Certificate {
     fn new(
         issuer: HashMap<String, String>,
         subject: HashMap<String, String>,
-        der_x509_certificate: Vec<u8>,
-        certificate_sha1: Vec<u8>,
+        der_x509_certificate: BytesWrapper,
+        certificate_sha1: BytesWrapper,
         certificate_id: Option<String>,
     ) -> Self {
+        crate::binding_utils::unwrap_bytes!(der_x509_certificate, certificate_sha1);
         Self(libparsec::types::X509Certificate {
             issuer,
             subject,
@@ -154,10 +156,11 @@ impl X509Certificate {
             py_kwargs,
             [issuer: HashMap<String, String>, "issuer"],
             [subject: HashMap<String, String>, "subject"],
-            [der_x509_certificate: Vec<u8>, "der_x509_certificate"],
-            [certificate_sha1: Vec<u8>, "certificate_sha1"],
+            [der_x509_certificate: BytesWrapper, "der_x509_certificate"],
+            [certificate_sha1: BytesWrapper, "certificate_sha1"],
             [certificate_id: Option<String>, "certificate_id"],
         );
+        crate::binding_utils::unwrap_bytes!(der_x509_certificate, certificate_sha1);
 
         let mut r = self.0.clone();
 
@@ -241,9 +244,10 @@ impl LocalPendingEnrollment {
         submitted_on: DateTime,
         enrollment_id: EnrollmentID,
         submit_payload: PkiEnrollmentSubmitPayload,
-        encrypted_key: Vec<u8>,
-        ciphertext: Vec<u8>,
+        encrypted_key: BytesWrapper,
+        ciphertext: BytesWrapper,
     ) -> Self {
+        crate::binding_utils::unwrap_bytes!(encrypted_key, ciphertext);
         Self(libparsec::types::LocalPendingEnrollment {
             x509_certificate: x509_certificate.0,
             addr: addr.0,

diff --git a/src/protocol/block.rs b/src/protocol/block.rs
@@ -8,6 +8,7 @@ use libparsec::protocol::{
 };
 
 use crate::{
+    binding_utils::BytesWrapper,
     ids::{BlockID, RealmID},
     protocol::{
         error::{ProtocolError, ProtocolErrorFields, ProtocolResult},
@@ -25,7 +26,8 @@ crate::binding_utils::gen_proto!(BlockCreateReq, __richcmp__, eq);
 #[pymethods]
 impl BlockCreateReq {
     #[new]
-    fn new(block_id: BlockID, realm_id: RealmID, block: Vec<u8>) -> PyResult<Self> {
+    fn new(block_id: BlockID, realm_id: RealmID, block: BytesWrapper) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(block);
         Ok(Self(block_create::Req {
             block_id: block_id.0,
             realm_id: realm_id.0,
@@ -130,7 +132,8 @@ pub(crate) struct BlockReadRepOk;
 #[pymethods]
 impl BlockReadRepOk {
     #[new]
-    fn new(block: Vec<u8>) -> PyResult<(Self, BlockReadRep)> {
+    fn new(block: BytesWrapper) -> PyResult<(Self, BlockReadRep)> {
+        crate::binding_utils::unwrap_bytes!(block);
         Ok((Self, BlockReadRep(block_read::Rep::Ok { block })))
     }
 

diff --git a/src/protocol/invite.rs b/src/protocol/invite.rs
@@ -24,6 +24,7 @@ use libparsec::protocol::{
 use crate::{
     api_crypto,
     api_crypto::{HashDigest, PublicKey},
+    binding_utils::BytesWrapper,
     enumerate::{
         InvitationDeletedReason, InvitationEmailSentStatus, InvitationStatus, InvitationType,
     },
@@ -677,7 +678,8 @@ pub(crate) struct Invite2aClaimerSendHashedNonceRepOk;
 #[pymethods]
 impl Invite2aClaimerSendHashedNonceRepOk {
     #[new]
-    fn new(greeter_nonce: Vec<u8>) -> PyResult<(Self, Invite2aClaimerSendHashedNonceRep)> {
+    fn new(greeter_nonce: BytesWrapper) -> PyResult<(Self, Invite2aClaimerSendHashedNonceRep)> {
+        crate::binding_utils::unwrap_bytes!(greeter_nonce);
         Ok((
             Self,
             Invite2aClaimerSendHashedNonceRep(invite_2a_claimer_send_hashed_nonce::Rep::Ok {
@@ -775,7 +777,8 @@ crate::binding_utils::gen_proto!(Invite2bClaimerSendNonceReq, __richcmp__, eq);
 #[pymethods]
 impl Invite2bClaimerSendNonceReq {
     #[new]
-    fn new(claimer_nonce: Vec<u8>) -> PyResult<Self> {
+    fn new(claimer_nonce: BytesWrapper) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(claimer_nonce);
         Ok(Self(invite_2b_claimer_send_nonce::Req { claimer_nonce }))
     }
 
@@ -829,7 +832,8 @@ crate::binding_utils::gen_proto!(Invite2bGreeterSendNonceReq, __richcmp__, eq);
 #[pymethods]
 impl Invite2bGreeterSendNonceReq {
     #[new]
-    fn new(token: InvitationToken, greeter_nonce: Vec<u8>) -> PyResult<Self> {
+    fn new(token: InvitationToken, greeter_nonce: BytesWrapper) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(greeter_nonce);
         let token = token.0;
         Ok(Self(invite_2b_greeter_send_nonce::Req {
             token,
@@ -874,7 +878,8 @@ pub(crate) struct Invite2bGreeterSendNonceRepOk;
 #[pymethods]
 impl Invite2bGreeterSendNonceRepOk {
     #[new]
-    fn new(claimer_nonce: Vec<u8>) -> PyResult<(Self, Invite2bGreeterSendNonceRep)> {
+    fn new(claimer_nonce: BytesWrapper) -> PyResult<(Self, Invite2bGreeterSendNonceRep)> {
+        crate::binding_utils::unwrap_bytes!(claimer_nonce);
         Ok((
             Self,
             Invite2bGreeterSendNonceRep(invite_2b_greeter_send_nonce::Rep::Ok { claimer_nonce }),
@@ -1110,7 +1115,8 @@ crate::binding_utils::gen_proto!(Invite4ClaimerCommunicateReq, __richcmp__, eq);
 #[pymethods]
 impl Invite4ClaimerCommunicateReq {
     #[new]
-    fn new(payload: Vec<u8>) -> PyResult<Self> {
+    fn new(payload: BytesWrapper) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(payload);
         Ok(Self(invite_4_claimer_communicate::Req { payload }))
     }
 
@@ -1146,7 +1152,8 @@ pub(crate) struct Invite4ClaimerCommunicateRepOk;
 #[pymethods]
 impl Invite4ClaimerCommunicateRepOk {
     #[new]
-    fn new(payload: Vec<u8>) -> PyResult<(Self, Invite4ClaimerCommunicateRep)> {
+    fn new(payload: BytesWrapper) -> PyResult<(Self, Invite4ClaimerCommunicateRep)> {
+        crate::binding_utils::unwrap_bytes!(payload);
         Ok((
             Self,
             Invite4ClaimerCommunicateRep(invite_4_claimer_communicate::Rep::Ok { payload }),
@@ -1174,7 +1181,8 @@ crate::binding_utils::gen_proto!(Invite4GreeterCommunicateReq, __richcmp__, eq);
 #[pymethods]
 impl Invite4GreeterCommunicateReq {
     #[new]
-    fn new(token: InvitationToken, payload: Vec<u8>) -> PyResult<Self> {
+    fn new(token: InvitationToken, payload: BytesWrapper) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(payload);
         let token = token.0;
         Ok(Self(invite_4_greeter_communicate::Req { token, payload }))
     }
@@ -1216,7 +1224,8 @@ pub(crate) struct Invite4GreeterCommunicateRepOk;
 #[pymethods]
 impl Invite4GreeterCommunicateRepOk {
     #[new]
-    fn new(payload: Vec<u8>) -> PyResult<(Self, Invite4GreeterCommunicateRep)> {
+    fn new(payload: BytesWrapper) -> PyResult<(Self, Invite4GreeterCommunicateRep)> {
+        crate::binding_utils::unwrap_bytes!(payload);
         Ok((
             Self,
             Invite4GreeterCommunicateRep(invite_4_greeter_communicate::Rep::Ok { payload }),

diff --git a/src/protocol/message.rs b/src/protocol/message.rs
@@ -8,12 +8,12 @@ use pyo3::{
 
 use libparsec::protocol::{authenticated_cmds::v2::message_get, Request};
 
-use crate::ids::DeviceID;
 use crate::protocol::{
     error::{ProtocolError, ProtocolErrorFields, ProtocolResult},
     gen_rep,
 };
 use crate::time::DateTime;
+use crate::{binding_utils::BytesWrapper, ids::DeviceID};
 
 #[pyclass]
 #[derive(Clone)]
@@ -25,7 +25,13 @@ crate::binding_utils::gen_proto!(Message, __richcmp__, eq);
 #[pymethods]
 impl Message {
     #[new]
-    fn new(count: u64, sender: DeviceID, timestamp: DateTime, body: Vec<u8>) -> PyResult<Self> {
+    fn new(
+        count: u64,
+        sender: DeviceID,
+        timestamp: DateTime,
+        body: BytesWrapper,
+    ) -> PyResult<Self> {
+        crate::binding_utils::unwrap_bytes!(body);
         let sender = sender.0;
         let timestamp = timestamp.0;
         Ok(Self(message_get::Message {

diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
@@ -148,9 +148,11 @@ macro_rules! gen_rep {
                 }
 
                 #[classmethod]
-                fn load<'py>(_cls: &::pyo3::types::PyType, buf: Vec<u8>, py: Python<'py>) -> PyResult<PyObject> {
+                fn load<'py>(_cls: &::pyo3::types::PyType, buf: crate::binding_utils::BytesWrapper, py: Python<'py>) -> PyResult<PyObject> {
                     use pyo3::{pyclass_init::PyObjectInit, PyTypeInfo};
 
+                    crate::binding_utils::unwrap_bytes!(buf);
+
                     let rep = $mod::Rep::load(&buf)
                         .map_err(|e| ProtocolErrorFields(libparsec::protocol::ProtocolError::DecodingError { exc: e.to_string() }))
                         .map_err(|e| ProtocolError::new_err(e))?;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Fix an underlying issue that caused an overall performance degradation