Do not assume base64 input and use bytes

Cargo.toml

@@ -10,7 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 paste = "1.0.15"
 thiserror = "1.0.63"
-vodozemac = "0.7.0"
+vodozemac = { git = "https://github.com/matrix-org/vodozemac.git" }
 
 [package.metadata.maturin]
 name = "vodozemac"

src/account.rs

@@ -70,7 +70,7 @@ impl Account {
         self.inner.curve25519_key().into()
     }
 
-    fn sign(&self, message: &str) -> Ed25519Signature {
+    fn sign(&self, message: &[u8]) -> Ed25519Signature {
         self.inner.sign(message).into()
     }
 
@@ -127,7 +127,7 @@ impl Account {
         &mut self,
         identity_key: &Curve25519PublicKey,
         message: &OlmMessage,
-    ) -> Result<(Session, String), SessionError> {
+    ) -> Result<(Session, Vec<u8>), SessionError> {
         let message =
             vodozemac::olm::OlmMessage::from_parts(message.message_type, &message.ciphertext)?;
 
@@ -140,7 +140,7 @@ impl Account {
                 Session {
                     inner: result.session,
                 },
-                String::from_utf8(result.plaintext)?,
+                result.plaintext,
             ))
         } else {
             Err(SessionError::InvalidMessageType)

src/group_sessions.rs

@@ -28,12 +28,12 @@ impl GroupSession {
     }
 
     #[getter]
-    fn session_key(&self) -> String {
-        self.inner.session_key().to_base64()
+    fn session_key(&self) -> Vec<u8> {
+        self.inner.session_key().to_bytes()
     }
 
-    fn encrypt(&mut self, plaintext: &str) -> String {
-        self.inner.encrypt(plaintext).to_base64()
+    fn encrypt(&mut self, plaintext: &[u8]) -> Vec<u8> {
+        self.inner.encrypt(plaintext).to_bytes()
     }
 
     fn pickle(&self, pickle_key: &[u8]) -> Result<String, PickleError> {
@@ -64,7 +64,7 @@ impl GroupSession {
 #[pyclass]
 pub struct DecryptedMessage {
     #[pyo3(get)]
-    plaintext: String,
+    plaintext: Vec<u8>,
     #[pyo3(get)]
     message_index: u32,
 }
@@ -77,8 +77,8 @@ pub struct InboundGroupSession {
 #[pymethods]
 impl InboundGroupSession {
     #[new]
-    fn new(session_key: &str) -> Result<Self, SessionKeyDecodeError> {
-        let key = SessionKey::from_base64(session_key)?;
+    fn new(session_key: &[u8]) -> Result<Self, SessionKeyDecodeError> {
+        let key = SessionKey::from_bytes(session_key)?;
 
         Ok(Self {
             inner: vodozemac::megolm::InboundGroupSession::new(&key, SessionConfig::version_1()),
@@ -88,9 +88,9 @@ impl InboundGroupSession {
     #[classmethod]
     fn import_session(
         _cls: &Bound<'_, PyType>,
-        session_key: &str,
+        session_key: &[u8],
     ) -> Result<Self, SessionKeyDecodeError> {
-        let key = ExportedSessionKey::from_base64(session_key)?;
+        let key = ExportedSessionKey::from_bytes(session_key)?;
 
         Ok(Self {
             inner: vodozemac::megolm::InboundGroupSession::import(&key, SessionConfig::version_1()),
@@ -107,16 +107,16 @@ impl InboundGroupSession {
         self.inner.first_known_index()
     }
 
-    fn export_at(&mut self, index: u32) -> Option<String> {
-        self.inner.export_at(index).map(|k| k.to_base64())
+    fn export_at(&mut self, index: u32) -> Option<Vec<u8>> {
+        self.inner.export_at(index).map(|k| k.to_bytes())
     }
 
-    fn decrypt(&mut self, ciphertext: &str) -> Result<DecryptedMessage, MegolmDecryptionError> {
-        let message = MegolmMessage::from_base64(ciphertext)?;
+    fn decrypt(&mut self, ciphertext: &[u8]) -> Result<DecryptedMessage, MegolmDecryptionError> {
+        let message = MegolmMessage::from_bytes(ciphertext)?;
         let ret = self.inner.decrypt(&message)?;
 
         Ok(DecryptedMessage {
-            plaintext: String::from_utf8(ret.plaintext)?,
+            plaintext: ret.plaintext,
             message_index: ret.message_index,
         })
     }

src/lib.rs

@@ -11,17 +11,17 @@ use pyo3::prelude::*;
 #[pyclass]
 pub struct OlmMessage {
     #[pyo3(get)]
-    ciphertext: String,
+    ciphertext: Vec<u8>,
     #[pyo3(get)]
     message_type: usize,
 }
 
 #[pymethods]
 impl OlmMessage {
     #[new]
-    pub fn new(message_type: usize, ciphertext: &str) -> Self {
+    pub fn new(message_type: usize, ciphertext: &[u8]) -> Self {
         Self {
-            ciphertext: ciphertext.to_owned(),
+            ciphertext: ciphertext.to_vec(),
             message_type,
         }
     }

src/sas.rs

@@ -1,20 +1,19 @@
 use pyo3::prelude::*;
-use vodozemac::Curve25519PublicKey;
 
-use crate::error::SasError;
+use crate::{error::SasError, types::Curve25519PublicKey};
 
 #[pyclass]
 pub struct Sas {
     inner: Option<vodozemac::sas::Sas>,
-    public_key: String,
+    public_key: vodozemac::Curve25519PublicKey,
 }
 
 #[pymethods]
 impl Sas {
     #[new]
     fn new() -> Self {
         let sas = vodozemac::sas::Sas::new();
-        let public_key = sas.public_key().to_base64();
+        let public_key = sas.public_key();
 
         Self {
             inner: Some(sas),
@@ -23,15 +22,13 @@ impl Sas {
     }
 
     #[getter]
-    fn public_key(&self) -> &str {
-        &self.public_key
+    fn public_key(&self) -> Curve25519PublicKey {
+        Curve25519PublicKey::from(self.public_key)
     }
 
-    fn diffie_hellman(&mut self, key: &str) -> Result<EstablishedSas, SasError> {
+    fn diffie_hellman(&mut self, key: Curve25519PublicKey) -> Result<EstablishedSas, SasError> {
         if let Some(sas) = self.inner.take() {
-            let key = Curve25519PublicKey::from_base64(key)?;
-            let sas = sas.diffie_hellman(key)?;
-
+            let sas = sas.diffie_hellman(key.inner)?;
             Ok(EstablishedSas { inner: sas })
         } else {
             Err(SasError::Used)

src/types/curve25519.rs

@@ -1,5 +1,5 @@
 use crate::error::*;
-use pyo3::{prelude::*, types::PyType};
+use pyo3::{prelude::*, types::{PyBytes, PyType}};
 
 #[pyclass]
 #[derive(Clone)]
@@ -22,10 +22,29 @@ impl Curve25519PublicKey {
         })
     }
 
+    #[classmethod]
+    pub fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, KeyError> {
+        let key: [u8; 32] = bytes.try_into().map_err(|_| {
+            KeyError::from(vodozemac::KeyError::InvalidKeyLength {
+                key_type: "Curve25519PublicKey",
+                expected_length: 32,
+                length: bytes.len(),
+            })
+        })?;
+
+        Ok(Self {
+            inner: vodozemac::Curve25519PublicKey::from_bytes(key),
+        })
+    }
+
     pub fn to_base64(&self) -> String {
         self.inner.to_base64()
     }
 
+    pub fn to_bytes(&self) -> Py<PyBytes> {
+        Python::with_gil(|py| PyBytes::new_bound(py, self.inner.to_bytes().as_slice()).into())
+    }
+
     #[classattr]
     const __hash__: Option<PyObject> = None;
 

src/types/ed25519.rs

@@ -15,6 +15,21 @@ impl Ed25519PublicKey {
         })
     }
 
+    #[classmethod]
+    pub fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, KeyError> {
+        let key: &[u8; 32] = bytes.try_into().map_err(|_| {
+            KeyError::from(vodozemac::KeyError::InvalidKeyLength {
+                key_type: "Ed25519PublicKey",
+                expected_length: 32,
+                length: bytes.len(),
+            })
+        })?;
+
+        Ok(Self {
+            inner: vodozemac::Ed25519PublicKey::from_slice(key)?,
+        })
+    }
+
     pub fn to_base64(&self) -> String {
         self.inner.to_base64()
     }