poc for PQC

rust/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [lib]
-crate-type = ["cdylib", "staticlib"]
+crate-type = ["cdylib", "staticlib", "rlib"]
 
 [dependencies]
 flutter_rust_bridge = "=2.12.0"
@@ -18,24 +18,11 @@ sqlx = { version = "0.9.0-alpha.1", default-features = false, features = [
     "derive",
     "json",
 ] }
-# mdk-core = { version = "0.8.0", git = "https://github.com/marmot-protocol/mdk", rev = "7f809f8549458a0d7f7d885bcdd694023abf299c", features = [
-#     "mip04",
-#     "mip05",
-# ] }
-# mdk-sqlite-storage = { version = "0.8.0", git = "https://github.com/marmot-protocol/mdk", rev = "7f809f8549458a0d7f7d885bcdd694023abf299c" }
-# mdk-storage-traits = { version = "0.8.0", git = "https://github.com/marmot-protocol/mdk", rev = "7f809f8549458a0d7f7d885bcdd694023abf299c" }
-# nostr-sdk = { version = "0.44", features = [
-#     "nip04",
-#     "nip44",
-#     "nip47",
-#     "nip59",
-# ] }
 libsqlite3-sys = { version = "0.35.0", features = [
     "bundled-sqlcipher-vendored-openssl",
 ] }
 tokio = { version = "1.44", features = ["full"] }
 tracing = "0.1.44"
-rand = "0.10.1"
 prost = "0.14.1"
 blahaj = "0.6.0"
 serde_json = "1.0"
@@ -56,6 +43,10 @@ chrono = { version = "0.4", features = ["serde"] }
 zip = { version = "2.2.2", default-features = false, features = ["deflate"] }
 scrypt = { version = "0.12", default-features = false }
 walkdir = "2.5.0"
+libsignal-protocol = { git = "https://github.com/signalapp/libsignal", version = "0.1.0" }
+rand08 = { version = "0.8.5", package = "rand" }
+rand = "0.9.4"
+async-trait = "0.1.91"
 [target.'cfg(target_os = "ios")'.dependencies]
 # iOS backend: Requires the 'protected' feature for Data Protection Keychain
 apple-native-keyring-store = { version = "1", features = ["protected"] }

rust/src/bridge/wrapper/key_manager.rs

@@ -36,6 +36,7 @@ impl RustKeyManager {
             identity_key_pair_structure,
             registration_id,
             pre_key_store: signed_pre_key_store,
+            pqc_identity_key: None,
         });
         key_manager.store_to_keychain(&ctx.secure_storage)?;
         Ok(())
@@ -132,13 +133,19 @@ impl RustKeyManager {
         }
         let mut key_array = [0u8; 32];
         key_array.copy_from_slice(&media_key);
-        Ok(key_manager.main_key.encrypt_cloud_media_key(&key_array, &addition))
+        Ok(key_manager
+            .main_key
+            .encrypt_cloud_media_key(&key_array, &addition))
     }
 
-    pub async fn decrypt_cloud_media_key(encrypted_media_key: Vec<u8>, addition: String) -> Result<Vec<u8>> {
+    pub async fn decrypt_cloud_media_key(
+        encrypted_media_key: Vec<u8>,
+        addition: String,
+    ) -> Result<Vec<u8>> {
         let key_manager = get_twonly_flutter()?.key_manager.lock().await;
-        let decrypted = key_manager.main_key.decrypt_cloud_media_key(&encrypted_media_key, &addition)?;
+        let decrypted = key_manager
+            .main_key
+            .decrypt_cloud_media_key(&encrypted_media_key, &addition)?;
         Ok(decrypted.to_vec())
     }
 }
-

rust/src/database/migrations/0002_create_signal_tables.sql

@@ -0,0 +1,34 @@
+-- Signal Identities (Stores remote IdentityKeys)
+CREATE TABLE IF NOT EXISTS signal_identities (
+    name TEXT NOT NULL,
+    device_id INTEGER NOT NULL,
+    identity_key BLOB NOT NULL,
+    timestamp INTEGER NOT NULL,
+    PRIMARY KEY (name, device_id)
+);
+
+-- Signal PreKeys (Stores local one-time prekeys)
+CREATE TABLE IF NOT EXISTS signal_pre_keys (
+    pre_key_id INTEGER PRIMARY KEY,
+    record_bytes BLOB NOT NULL
+);
+
+-- Signal Signed PreKeys (Stores local signed prekeys)
+CREATE TABLE IF NOT EXISTS signal_signed_pre_keys (
+    signed_pre_key_id INTEGER PRIMARY KEY,
+    record_bytes BLOB NOT NULL
+);
+
+-- Signal Kyber PreKeys (Stores local PQC Kyber prekeys)
+CREATE TABLE IF NOT EXISTS signal_kyber_pre_keys (
+    kyber_pre_key_id INTEGER PRIMARY KEY,
+    record_bytes BLOB NOT NULL
+);
+
+-- Signal Sessions (Stores active sessions with remote devices)
+CREATE TABLE IF NOT EXISTS signal_sessions (
+    name TEXT NOT NULL,
+    device_id INTEGER NOT NULL,
+    record_bytes BLOB NOT NULL,
+    PRIMARY KEY (name, device_id)
+);

rust/src/database/mod.rs

@@ -5,12 +5,12 @@ use std::time::Duration;
 
 pub(crate) mod tables;
 
-pub(crate) struct Database {
+pub struct Database {
-    pub(crate) pool: SqlitePool,
+    pub pool: SqlitePool,
 }
 
 impl Database {
-    pub(crate) async fn new(
+    pub async fn new(
         db_path: &String,
         encryption_key: Option<&str>,
         read_only: bool,
@@ -47,7 +47,7 @@ impl Database {
         Ok(Self { pool })
     }
 
-    pub(crate) async fn run_migrations(&self) -> Result<()> {
+    pub async fn run_migrations(&self) -> Result<()> {
         sqlx::migrate!("./src/database/migrations")
             .run(&self.pool)
             .await

rust/src/error.rs

@@ -67,6 +67,9 @@ pub enum TwonlyError {
     InvalidOutputLen(#[from] InvalidOutputLen),
     #[error("AES-GCM error")]
     AesGcm,
+
+    #[error("Signal protocol error: {0}")]
+    Signal(String),
 }
 
 impl From<String> for TwonlyError {

rust/src/keys/identity_key/signal_identity_key.rs

@@ -5,10 +5,10 @@ use zeroize::{Zeroize, ZeroizeOnDrop};
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub(crate) struct SignalIdentityKey {
-    // https://github.com/MixinNetwork/libsignal_protocol_dart/blob/c95a1586057022acdbb9c76b1692d94cc549bcc7/protobuf/LocalStorageProtocol.proto#L85
     pub(crate) identity_key_pair_structure: Vec<u8>,
     pub(crate) registration_id: i64,
     pub(crate) pre_key_store: HashMap<i64, Vec<u8>>,
+    pub(crate) pqc_identity_key: Option<Vec<u8>>,
 }
 
 impl SignalIdentityKey {}
@@ -21,6 +21,9 @@ impl Zeroize for SignalIdentityKey {
             value.zeroize();
         }
         self.pre_key_store.clear();
+        if let Some(pqc) = &mut self.pqc_identity_key {
+            pqc.zeroize();
+        }
     }
 }
 

rust/src/lib.rs

@@ -1,12 +1,13 @@
 mod backup;
 pub mod bridge;
 mod context;
-mod database;
+pub mod database;
 mod error;
 mod frb_generated;
 mod keys;
 mod log;
 mod secure_storage;
+pub mod signal;
 mod standalone;
 mod user_discovery;
 mod utils;

rust/src/signal/engine.rs

@@ -0,0 +1,336 @@
+use crate::error::{Result, TwonlyError};
+use libsignal_protocol::{
+    message_encrypt, process_prekey_bundle, DeviceId, GenericSignedPreKey, IdentityKey,
+    IdentityKeyPair, IdentityKeyStore, KyberPreKeyId, KyberPreKeyStore, PreKeyBundle, PreKeyId,
+    PreKeySignalMessage, PreKeyStore, ProtocolAddress, PublicKey, SignalMessage, SignedPreKeyId,
+    SignedPreKeyStore, Timestamp,
+};
+use std::sync::Arc;
+use std::time::{SystemTime, UNIX_EPOCH};
+use tokio::sync::Mutex;
+
+use crate::signal::store::DbSignalProtocolStore;
+
+pub struct RustSignalEngine {
+    store: Arc<Mutex<DbSignalProtocolStore>>,
+}
+
+pub struct FrbPreKeyBundle {
+    pub registration_id: u32,
+    pub device_id: u32,
+    pub pre_key_id: Option<u32>,
+    pub pre_key_public: Option<Vec<u8>>,
+    pub signed_pre_key_id: u32,
+    pub signed_pre_key_public: Vec<u8>,
+    pub signed_pre_key_signature: Vec<u8>,
+    pub kyber_pre_key_id: u32,
+    pub kyber_pre_key_public: Vec<u8>,
+    pub kyber_pre_key_signature: Vec<u8>,
+    pub identity_key: Vec<u8>,
+}
+
+impl RustSignalEngine {
+    pub async fn new() -> Result<Self> {
+        let twonly = crate::bridge::get_twonly_flutter()?;
+        let pool = twonly.rust_db.pool.clone();
+
+        let km = twonly.key_manager.lock().await;
+        let signal_identity = km
+            .signal_identity
+            .as_ref()
+            .ok_or_else(|| TwonlyError::Generic("No signal identity found".to_string()))?;
+
+        let identity_key_pair =
+            IdentityKeyPair::try_from(&signal_identity.identity_key_pair_structure[..])
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let store = DbSignalProtocolStore::new(
+            pool,
+            identity_key_pair,
+            signal_identity.registration_id as u32,
+        );
+
+        Ok(Self {
+            store: Arc::new(Mutex::new(store)),
+        })
+    }
+
+    pub fn new_with_pool(
+        pool: sqlx::SqlitePool,
+        identity_key_pair_bytes: Vec<u8>,
+        local_registration_id: u32,
+    ) -> Result<Self> {
+        let identity_key_pair = IdentityKeyPair::try_from(&identity_key_pair_bytes[..])
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let store = DbSignalProtocolStore::new(pool, identity_key_pair, local_registration_id);
+
+        Ok(Self {
+            store: Arc::new(Mutex::new(store)),
+        })
+    }
+
+    pub fn generate_identity_key_pair() -> Result<Vec<u8>> {
+        let mut csprng = rand::rng();
+        let key_pair = IdentityKeyPair::generate(&mut csprng);
+        Ok(key_pair.serialize().to_vec())
+    }
+
+    pub async fn generate_bundle(
+        &self,
+        pre_key_id: u32,
+        signed_pre_key_id: u32,
+    ) -> Result<FrbPreKeyBundle> {
+        let mut store_guard = self.store.lock().await;
+        let store = &mut *store_guard;
+        let mut csprng = rand::rng();
+
+        let pre_key_pair = libsignal_protocol::KeyPair::generate(&mut csprng);
+        store
+            .pre_key_store
+            .save_pre_key(
+                pre_key_id.into(),
+                &libsignal_protocol::PreKeyRecord::new(pre_key_id.into(), &pre_key_pair),
+            )
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let signed_pre_key_pair = libsignal_protocol::KeyPair::generate(&mut csprng);
+        let signature = store
+            .identity_store
+            .get_identity_key_pair()
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?
+            .private_key()
+            .calculate_signature_for_multipart_message(
+                &[&signed_pre_key_pair.public_key.serialize()],
+                &mut csprng,
+            )
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+        let timestamp = Timestamp::from_epoch_millis(
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_millis() as u64,
+        );
+        store
+            .signed_pre_key_store
+            .save_signed_pre_key(
+                signed_pre_key_id.into(),
+                &libsignal_protocol::SignedPreKeyRecord::new(
+                    signed_pre_key_id.into(),
+                    timestamp,
+                    &signed_pre_key_pair,
+                    &signature,
+                ),
+            )
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let kyber_pre_key_id = 1;
+        let kyber_key_pair = libsignal_protocol::kem::KeyPair::generate(
+            libsignal_protocol::kem::KeyType::Kyber1024,
+            &mut csprng,
+        );
+        let kyber_signature = store
+            .identity_store
+            .get_identity_key_pair()
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?
+            .private_key()
+            .calculate_signature_for_multipart_message(
+                &[&kyber_key_pair.public_key.serialize()],
+                &mut csprng,
+            )
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+        let kyber_sig_arr: [u8; 64] = kyber_signature[..].try_into().unwrap();
+        store
+            .kyber_pre_key_store
+            .save_kyber_pre_key(
+                kyber_pre_key_id.into(),
+                &libsignal_protocol::KyberPreKeyRecord::new(
+                    kyber_pre_key_id.into(),
+                    timestamp,
+                    &kyber_key_pair,
+                    &kyber_sig_arr,
+                ),
+            )
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        Ok(FrbPreKeyBundle {
+            registration_id: store
+                .identity_store
+                .get_local_registration_id()
+                .await
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?,
+            device_id: 1,
+            pre_key_id: Some(pre_key_id),
+            pre_key_public: Some(pre_key_pair.public_key.serialize().to_vec()),
+            signed_pre_key_id: signed_pre_key_id,
+            signed_pre_key_public: signed_pre_key_pair.public_key.serialize().to_vec(),
+            signed_pre_key_signature: signature.to_vec(),
+            kyber_pre_key_id: kyber_pre_key_id,
+            kyber_pre_key_public: kyber_key_pair.public_key.serialize().to_vec(),
+            kyber_pre_key_signature: kyber_signature.to_vec(),
+            identity_key: store
+                .identity_store
+                .get_identity_key_pair()
+                .await
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?
+                .identity_key()
+                .serialize()
+                .to_vec(),
+        })
+    }
+
+    pub async fn process_prekey_bundle(
+        &self,
+        name: String,
+        device_id: u32,
+        bundle: FrbPreKeyBundle,
+    ) -> Result<()> {
+        let mut store_guard = self.store.lock().await;
+        let store = &mut *store_guard;
+
+        let d_id = DeviceId::try_from(device_id)
+            .map_err(|_| TwonlyError::Generic(format!("Invalid device id: {}", device_id)))?;
+        let remote_address = ProtocolAddress::new(name, d_id);
+        let local_address =
+            ProtocolAddress::new("local".to_string(), DeviceId::try_from(1).unwrap());
+
+        let identity_key = IdentityKey::decode(&bundle.identity_key)
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let signed_pre_key_public = PublicKey::deserialize(&bundle.signed_pre_key_public)
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        let pre_key_public = match bundle.pre_key_public {
+            Some(pk) => {
+                Some(PublicKey::deserialize(&pk).map_err(|e| TwonlyError::Signal(e.to_string()))?)
+            }
+            None => None,
+        };
+
+        let pre_key = match (bundle.pre_key_id, pre_key_public) {
+            (Some(id), Some(pk)) => Some((PreKeyId::from(id), pk)),
+            _ => None,
+        };
+
+        let mut csprng = rand::rng();
+
+        let pre_key_bundle = PreKeyBundle::new(
+            bundle.registration_id,
+            DeviceId::try_from(bundle.device_id).unwrap_or(DeviceId::try_from(1).unwrap()),
+            pre_key,
+            SignedPreKeyId::from(bundle.signed_pre_key_id),
+            signed_pre_key_public,
+            bundle.signed_pre_key_signature,
+            KyberPreKeyId::from(bundle.kyber_pre_key_id),
+            libsignal_protocol::kem::PublicKey::deserialize(&bundle.kyber_pre_key_public)
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?,
+            bundle.kyber_pre_key_signature,
+            identity_key,
+        )
+        .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        process_prekey_bundle(
+            &remote_address,
+            &local_address,
+            &mut store.session_store,
+            &mut store.identity_store,
+            &pre_key_bundle,
+            SystemTime::now(),
+            &mut csprng,
+        )
+        .await
+        .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        Ok(())
+    }
+
+    pub async fn encrypt_message(
+        &self,
+        name: String,
+        device_id: u32,
+        plaintext: Vec<u8>,
+    ) -> Result<Vec<u8>> {
+        let mut store_guard = self.store.lock().await;
+        let store = &mut *store_guard;
+
+        let d_id = DeviceId::try_from(device_id)
+            .map_err(|_| TwonlyError::Generic(format!("Invalid device id: {}", device_id)))?;
+        let remote_address = ProtocolAddress::new(name, d_id);
+        let local_address =
+            ProtocolAddress::new("local".to_string(), DeviceId::try_from(1).unwrap());
+
+        let mut csprng = rand::rng();
+        let now = SystemTime::now();
+
+        let ciphertext = message_encrypt(
+            &plaintext,
+            &remote_address,
+            &local_address,
+            &mut store.session_store,
+            &mut store.identity_store,
+            now,
+            &mut csprng,
+        )
+        .await
+        .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+
+        Ok(ciphertext.serialize().to_vec())
+    }
+
+    pub async fn decrypt_message(
+        &self,
+        name: String,
+        device_id: u32,
+        ciphertext_bytes: Vec<u8>,
+        is_prekey_message: bool,
+    ) -> Result<Vec<u8>> {
+        let mut store_guard = self.store.lock().await;
+        let store = &mut *store_guard;
+
+        let d_id = DeviceId::try_from(device_id)
+            .map_err(|_| TwonlyError::Generic(format!("Invalid device id: {}", device_id)))?;
+        let remote_address = ProtocolAddress::new(name, d_id);
+        let local_address =
+            ProtocolAddress::new("local".to_string(), DeviceId::try_from(1).unwrap());
+
+        let plaintext = if is_prekey_message {
+            let message = PreKeySignalMessage::try_from(&ciphertext_bytes[..])
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+            let mut csprng = rand::rng();
+            libsignal_protocol::message_decrypt_prekey(
+                &message,
+                &remote_address,
+                &local_address,
+                &mut store.session_store,
+                &mut store.identity_store,
+                &mut store.pre_key_store,
+                &store.signed_pre_key_store,
+                &mut store.kyber_pre_key_store,
+                &mut csprng,
+            )
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?
+        } else {
+            let message = SignalMessage::try_from(&ciphertext_bytes[..])
+                .map_err(|e| TwonlyError::Signal(e.to_string()))?;
+            let mut csprng = rand::rng();
+            libsignal_protocol::message_decrypt_signal(
+                &message,
+                &remote_address,
+                &local_address,
+                &mut store.session_store,
+                &mut store.identity_store,
+                &mut csprng,
+            )
+            .await
+            .map_err(|e| TwonlyError::Signal(e.to_string()))?
+        };
+
+        Ok(plaintext.to_vec())
+    }
+}

rust/src/signal/mod.rs

@@ -0,0 +1,2 @@
+pub mod engine;
+pub mod store;

rust/src/signal/store.rs

@@ -0,0 +1,352 @@
+use async_trait::async_trait;
+use libsignal_protocol::{
+    Direction, GenericSignedPreKey, IdentityChange, IdentityKey, IdentityKeyPair, IdentityKeyStore,
+    KyberPreKeyId, KyberPreKeyRecord, KyberPreKeyStore, PreKeyId, PreKeyRecord, PreKeyStore,
+    ProtocolAddress, PublicKey, SessionRecord, SessionStore, SignalProtocolError, SignedPreKeyId,
+    SignedPreKeyRecord, SignedPreKeyStore,
+};
+use sqlx::SqlitePool;
+
+pub struct DbSignalProtocolStore {
+    pub identity_store: DbIdentityKeyStore,
+    pub pre_key_store: DbPreKeyStore,
+    pub signed_pre_key_store: DbSignedPreKeyStore,
+    pub kyber_pre_key_store: DbKyberPreKeyStore,
+    pub session_store: DbSessionStore,
+}
+
+impl DbSignalProtocolStore {
+    pub fn new(
+        pool: SqlitePool,
+        identity_key_pair: IdentityKeyPair,
+        local_registration_id: u32,
+    ) -> Self {
+        Self {
+            identity_store: DbIdentityKeyStore {
+                pool: pool.clone(),
+                identity_key_pair,
+                local_registration_id,
+            },
+            pre_key_store: DbPreKeyStore { pool: pool.clone() },
+            signed_pre_key_store: DbSignedPreKeyStore { pool: pool.clone() },
+            kyber_pre_key_store: DbKyberPreKeyStore { pool: pool.clone() },
+            session_store: DbSessionStore { pool },
+        }
+    }
+}
+
+pub struct DbIdentityKeyStore {
+    pool: SqlitePool,
+    identity_key_pair: IdentityKeyPair,
+    local_registration_id: u32,
+}
+
+#[async_trait(?Send)]
+impl IdentityKeyStore for DbIdentityKeyStore {
+    async fn get_identity_key_pair(&self) -> Result<IdentityKeyPair, SignalProtocolError> {
+        Ok(self.identity_key_pair.clone())
+    }
+
+    async fn get_local_registration_id(&self) -> Result<u32, SignalProtocolError> {
+        Ok(self.local_registration_id)
+    }
+
+    async fn save_identity(
+        &mut self,
+        address: &ProtocolAddress,
+        identity: &IdentityKey,
+    ) -> Result<IdentityChange, SignalProtocolError> {
+        let name = address.name();
+        let device_id: u32 = address.device_id().into();
+        let identity_bytes = identity.serialize();
+        let timestamp = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_millis() as i64;
+
+        let existing: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT identity_key FROM signal_identities WHERE name = ? AND device_id = ?",
+        )
+        .bind(name)
+        .bind(device_id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        let changed = if let Some(row) = existing {
+            row.0 != identity_bytes.as_ref()
+        } else {
+            false
+        };
+
+        sqlx::query("INSERT INTO signal_identities (name, device_id, identity_key, timestamp) VALUES (?, ?, ?, ?) ON CONFLICT(name, device_id) DO UPDATE SET identity_key = excluded.identity_key, timestamp = excluded.timestamp")
+            .bind(name)
+            .bind(device_id)
+            .bind(identity_bytes.as_ref())
+            .bind(timestamp)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        Ok(IdentityChange::from_changed(changed))
+    }
+
+    async fn is_trusted_identity(
+        &self,
+        address: &ProtocolAddress,
+        identity: &IdentityKey,
+        _direction: Direction,
+    ) -> Result<bool, SignalProtocolError> {
+        let name = address.name();
+        let device_id: u32 = address.device_id().into();
+        let identity_bytes = identity.serialize();
+
+        let row: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT identity_key FROM signal_identities WHERE name = ? AND device_id = ?",
+        )
+        .bind(name)
+        .bind(device_id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        if let Some((stored_key,)) = row {
+            Ok(stored_key == identity_bytes.as_ref())
+        } else {
+            Ok(true)
+        }
+    }
+
+    async fn get_identity(
+        &self,
+        address: &ProtocolAddress,
+    ) -> Result<Option<IdentityKey>, SignalProtocolError> {
+        let name = address.name();
+        let device_id: u32 = address.device_id().into();
+
+        let row: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT identity_key FROM signal_identities WHERE name = ? AND device_id = ?",
+        )
+        .bind(name)
+        .bind(device_id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        if let Some((bytes,)) = row {
+            let key = IdentityKey::decode(&bytes)
+                .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+            Ok(Some(key))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+pub struct DbPreKeyStore {
+    pool: SqlitePool,
+}
+
+#[async_trait(?Send)]
+impl PreKeyStore for DbPreKeyStore {
+    async fn get_pre_key(&self, prekey_id: PreKeyId) -> Result<PreKeyRecord, SignalProtocolError> {
+        let id: u32 = prekey_id.into();
+        let row: Option<(Vec<u8>,)> =
+            sqlx::query_as("SELECT record_bytes FROM signal_pre_keys WHERE pre_key_id = ?")
+                .bind(id)
+                .fetch_optional(&self.pool)
+                .await
+                .map_err(|_| SignalProtocolError::InvalidPreKeyId)?;
+
+        if let Some((bytes,)) = row {
+            PreKeyRecord::deserialize(&bytes).map_err(|_| SignalProtocolError::InvalidPreKeyId)
+        } else {
+            Err(SignalProtocolError::InvalidPreKeyId)
+        }
+    }
+
+    async fn save_pre_key(
+        &mut self,
+        prekey_id: PreKeyId,
+        record: &PreKeyRecord,
+    ) -> Result<(), SignalProtocolError> {
+        let id: u32 = prekey_id.into();
+        let bytes = record
+            .serialize()
+            .map_err(|_| SignalProtocolError::InvalidPreKeyId)?;
+
+        sqlx::query("INSERT INTO signal_pre_keys (pre_key_id, record_bytes) VALUES (?, ?) ON CONFLICT(pre_key_id) DO UPDATE SET record_bytes = excluded.record_bytes")
+            .bind(id)
+            .bind(bytes)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::InvalidPreKeyId)?;
+        Ok(())
+    }
+
+    async fn remove_pre_key(&mut self, prekey_id: PreKeyId) -> Result<(), SignalProtocolError> {
+        let id: u32 = prekey_id.into();
+        sqlx::query("DELETE FROM signal_pre_keys WHERE pre_key_id = ?")
+            .bind(id)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::InvalidPreKeyId)?;
+        Ok(())
+    }
+}
+
+pub struct DbSignedPreKeyStore {
+    pool: SqlitePool,
+}
+
+#[async_trait(?Send)]
+impl SignedPreKeyStore for DbSignedPreKeyStore {
+    async fn get_signed_pre_key(
+        &self,
+        signed_prekey_id: SignedPreKeyId,
+    ) -> Result<SignedPreKeyRecord, SignalProtocolError> {
+        let id: u32 = signed_prekey_id.into();
+        let row: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT record_bytes FROM signal_signed_pre_keys WHERE signed_pre_key_id = ?",
+        )
+        .bind(id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::InvalidSignedPreKeyId)?;
+
+        if let Some((bytes,)) = row {
+            SignedPreKeyRecord::deserialize(&bytes)
+                .map_err(|_| SignalProtocolError::InvalidSignedPreKeyId)
+        } else {
+            Err(SignalProtocolError::InvalidSignedPreKeyId)
+        }
+    }
+
+    async fn save_signed_pre_key(
+        &mut self,
+        signed_prekey_id: SignedPreKeyId,
+        record: &SignedPreKeyRecord,
+    ) -> Result<(), SignalProtocolError> {
+        let id: u32 = signed_prekey_id.into();
+        let bytes = record
+            .serialize()
+            .map_err(|_| SignalProtocolError::InvalidSignedPreKeyId)?;
+
+        sqlx::query("INSERT INTO signal_signed_pre_keys (signed_pre_key_id, record_bytes) VALUES (?, ?) ON CONFLICT(signed_pre_key_id) DO UPDATE SET record_bytes = excluded.record_bytes")
+            .bind(id)
+            .bind(bytes)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::InvalidSignedPreKeyId)?;
+        Ok(())
+    }
+}
+
+pub struct DbKyberPreKeyStore {
+    pool: SqlitePool,
+}
+
+#[async_trait(?Send)]
+impl KyberPreKeyStore for DbKyberPreKeyStore {
+    async fn get_kyber_pre_key(
+        &self,
+        kyber_prekey_id: KyberPreKeyId,
+    ) -> Result<KyberPreKeyRecord, SignalProtocolError> {
+        let id: u32 = kyber_prekey_id.into();
+        let row: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT record_bytes FROM signal_kyber_pre_keys WHERE kyber_pre_key_id = ?",
+        )
+        .bind(id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::InvalidKyberPreKeyId)?;
+
+        if let Some((bytes,)) = row {
+            KyberPreKeyRecord::deserialize(&bytes)
+                .map_err(|_| SignalProtocolError::InvalidKyberPreKeyId)
+        } else {
+            Err(SignalProtocolError::InvalidKyberPreKeyId)
+        }
+    }
+
+    async fn save_kyber_pre_key(
+        &mut self,
+        kyber_prekey_id: KyberPreKeyId,
+        record: &KyberPreKeyRecord,
+    ) -> Result<(), SignalProtocolError> {
+        let id: u32 = kyber_prekey_id.into();
+        let bytes = record
+            .serialize()
+            .map_err(|_| SignalProtocolError::InvalidKyberPreKeyId)?;
+
+        sqlx::query("INSERT INTO signal_kyber_pre_keys (kyber_pre_key_id, record_bytes) VALUES (?, ?) ON CONFLICT(kyber_pre_key_id) DO UPDATE SET record_bytes = excluded.record_bytes")
+            .bind(id)
+            .bind(bytes)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::InvalidKyberPreKeyId)?;
+        Ok(())
+    }
+
+    async fn mark_kyber_pre_key_used(
+        &mut self,
+        _kyber_prekey_id: KyberPreKeyId,
+        _ec_prekey_id: SignedPreKeyId,
+        _base_key: &PublicKey,
+    ) -> Result<(), SignalProtocolError> {
+        Ok(())
+    }
+}
+
+pub struct DbSessionStore {
+    pool: SqlitePool,
+}
+
+#[async_trait(?Send)]
+impl SessionStore for DbSessionStore {
+    async fn load_session(
+        &self,
+        address: &ProtocolAddress,
+    ) -> Result<Option<SessionRecord>, SignalProtocolError> {
+        let name = address.name();
+        let device_id: u32 = address.device_id().into();
+
+        let row: Option<(Vec<u8>,)> = sqlx::query_as(
+            "SELECT record_bytes FROM signal_sessions WHERE name = ? AND device_id = ?",
+        )
+        .bind(name)
+        .bind(device_id)
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        if let Some((bytes,)) = row {
+            let record = SessionRecord::deserialize(&bytes)
+                .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+            Ok(Some(record))
+        } else {
+            Ok(None)
+        }
+    }
+
+    async fn store_session(
+        &mut self,
+        address: &ProtocolAddress,
+        record: &SessionRecord,
+    ) -> Result<(), SignalProtocolError> {
+        let name = address.name();
+        let device_id: u32 = address.device_id().into();
+        let bytes = record
+            .serialize()
+            .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+
+        sqlx::query("INSERT INTO signal_sessions (name, device_id, record_bytes) VALUES (?, ?, ?) ON CONFLICT(name, device_id) DO UPDATE SET record_bytes = excluded.record_bytes")
+            .bind(name)
+            .bind(device_id)
+            .bind(bytes)
+            .execute(&self.pool)
+            .await
+            .map_err(|_| SignalProtocolError::UntrustedIdentity(address.clone()))?;
+        Ok(())
+    }
+}

rust/tests/libsignal_compatibility_test.rs

@@ -0,0 +1,194 @@
+use libsignal_protocol::{
+    message_decrypt, message_encrypt, process_prekey_bundle, CiphertextMessageType, DeviceId,
+    GenericSignedPreKey, IdentityKeyPair, InMemSignalProtocolStore, KeyPair, KyberPreKeyStore,
+    PreKeyBundle, PreKeyRecord, PreKeyStore, ProtocolAddress, SignedPreKeyStore, Timestamp,
+};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+#[tokio::test]
+async fn libsignal_compatibility_test() -> Result<(), Box<dyn std::error::Error>> {
+    let mut csprng = rand::rng();
+
+    // ==========================================
+    // 1. SETUP BOB (Receiver)
+    // ==========================================
+    let bob_identity = IdentityKeyPair::generate(&mut csprng);
+    let bob_registration_id = 5678;
+    let mut bob_store =
+        InMemSignalProtocolStore::new(bob_identity.clone(), bob_registration_id).unwrap();
+    let bob_address = ProtocolAddress::new("bob".to_string(), DeviceId::try_from(1).unwrap());
+
+    // Bob generates PreKey
+    let bob_pre_key_id = 1.into();
+    let bob_pre_key_pair = KeyPair::generate(&mut csprng);
+    bob_store
+        .pre_key_store
+        .save_pre_key(
+            bob_pre_key_id,
+            &PreKeyRecord::new(bob_pre_key_id, &bob_pre_key_pair),
+        )
+        .await?;
+
+    // Bob generates SignedPreKey
+    let bob_signed_pre_key_id = 1.into();
+    let bob_signed_pre_key_pair = KeyPair::generate(&mut csprng);
+    let bob_signed_pre_key_sig = bob_identity
+        .private_key()
+        .calculate_signature_for_multipart_message(
+            &[&bob_signed_pre_key_pair.public_key.serialize()],
+            &mut csprng,
+        )?;
+    let timestamp = Timestamp::from_epoch_millis(
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_millis() as u64,
+    );
+    bob_store
+        .signed_pre_key_store
+        .save_signed_pre_key(
+            bob_signed_pre_key_id,
+            &libsignal_protocol::SignedPreKeyRecord::new(
+                bob_signed_pre_key_id,
+                timestamp,
+                &bob_signed_pre_key_pair,
+                &bob_signed_pre_key_sig,
+            ),
+        )
+        .await?;
+
+    // Bob generates KyberPreKey (Post-Quantum!)
+    let bob_kyber_pre_key_id = 1.into();
+    let bob_kyber_key_pair = libsignal_protocol::kem::KeyPair::generate(
+        libsignal_protocol::kem::KeyType::Kyber1024,
+        &mut csprng,
+    );
+    let bob_kyber_sig = bob_identity
+        .private_key()
+        .calculate_signature_for_multipart_message(
+            &[&bob_kyber_key_pair.public_key.serialize()],
+            &mut csprng,
+        )?;
+
+    // Kyber signature must be exactly 64 bytes
+    let kyber_sig_arr: [u8; 64] = bob_kyber_sig[..].try_into().unwrap();
+
+    bob_store
+        .kyber_pre_key_store
+        .save_kyber_pre_key(
+            bob_kyber_pre_key_id,
+            &libsignal_protocol::KyberPreKeyRecord::new(
+                bob_kyber_pre_key_id,
+                timestamp,
+                &bob_kyber_key_pair,
+                &kyber_sig_arr,
+            ),
+        )
+        .await?;
+
+    let bob_bundle = PreKeyBundle::new(
+        bob_registration_id,
+        DeviceId::try_from(1).unwrap(),
+        Some((bob_pre_key_id, bob_pre_key_pair.public_key)),
+        bob_signed_pre_key_id,
+        bob_signed_pre_key_pair.public_key,
+        bob_signed_pre_key_sig.to_vec(),
+        bob_kyber_pre_key_id,
+        bob_kyber_key_pair.public_key,
+        kyber_sig_arr.to_vec(),
+        *bob_identity.identity_key(),
+    )?;
+
+    // ==========================================
+    // 2. SETUP ALICE (Sender)
+    // ==========================================
+    let alice_identity = IdentityKeyPair::generate(&mut csprng);
+    let alice_registration_id = 1234;
+    let mut alice_store =
+        InMemSignalProtocolStore::new(alice_identity, alice_registration_id).unwrap();
+    let alice_address = ProtocolAddress::new("alice".to_string(), DeviceId::try_from(1).unwrap());
+
+    // Alice processes Bob's bundle
+    process_prekey_bundle(
+        &bob_address,
+        &alice_address,
+        &mut alice_store.session_store,
+        &mut alice_store.identity_store,
+        &bob_bundle,
+        SystemTime::now(),
+        &mut csprng,
+    )
+    .await?;
+
+    // ==========================================
+    // 3. EXCHANGE 100 MESSAGES
+    // ==========================================
+    let mut alice_to_bob = true;
+    for i in 1..=100 {
+        if alice_to_bob {
+            let plaintext = format!("Message {} from Alice to Bob", i).into_bytes();
+            let ciphertext = message_encrypt(
+                &plaintext,
+                &bob_address,
+                &alice_address,
+                &mut alice_store.session_store,
+                &mut alice_store.identity_store,
+                SystemTime::now(),
+                &mut csprng,
+            )
+            .await?;
+
+            if i == 1 {
+                assert_eq!(ciphertext.message_type(), CiphertextMessageType::PreKey);
+            } else {
+                assert_eq!(ciphertext.message_type(), CiphertextMessageType::Whisper);
+            }
+
+            let decrypted = message_decrypt(
+                &ciphertext,
+                &alice_address,
+                &bob_address,
+                &mut bob_store.session_store,
+                &mut bob_store.identity_store,
+                &mut bob_store.pre_key_store,
+                &bob_store.signed_pre_key_store,
+                &mut bob_store.kyber_pre_key_store,
+                &mut csprng,
+            )
+            .await?;
+            assert_eq!(plaintext, decrypted);
+        } else {
+            let plaintext = format!("Message {} from Bob to Alice", i).into_bytes();
+            let ciphertext = message_encrypt(
+                &plaintext,
+                &alice_address,
+                &bob_address,
+                &mut bob_store.session_store,
+                &mut bob_store.identity_store,
+                SystemTime::now(),
+                &mut csprng,
+            )
+            .await?;
+
+            assert_eq!(ciphertext.message_type(), CiphertextMessageType::Whisper);
+
+            let decrypted = message_decrypt(
+                &ciphertext,
+                &bob_address,
+                &alice_address,
+                &mut alice_store.session_store,
+                &mut alice_store.identity_store,
+                &mut alice_store.pre_key_store,
+                &alice_store.signed_pre_key_store,
+                &mut alice_store.kyber_pre_key_store,
+                &mut csprng,
+            )
+            .await?;
+            assert_eq!(plaintext, decrypted);
+        }
+
+        alice_to_bob = !alice_to_bob;
+    }
+
+    Ok(())
+}

rust/tests/pqxdh_tests.rs

@@ -0,0 +1,66 @@
+#[tokio::test]
+async fn test_twonly_api_100_messages() -> Result<(), Box<dyn std::error::Error>> {
+    use rust_lib_twonly::database::Database;
+    use rust_lib_twonly::signal::engine::RustSignalEngine;
+
+    let _ = pretty_env_logger::try_init();
+
+    // 1. Setup in-memory databases
+    let alice_db = Database::new(&"sqlite::memory:".to_string(), None, false).await?;
+    alice_db.run_migrations().await?;
+
+    let bob_db = Database::new(&"sqlite::memory:".to_string(), None, false).await?;
+    bob_db.run_migrations().await?;
+
+    // 2. Setup Alice and Bob identity keys
+    let alice_identity_bytes = RustSignalEngine::generate_identity_key_pair()?;
+    let bob_identity_bytes = RustSignalEngine::generate_identity_key_pair()?;
+
+    // 3. Initialize engines with the DB pools
+    let alice_engine =
+        RustSignalEngine::new_with_pool(alice_db.pool.clone(), alice_identity_bytes, 1234)?;
+    let bob_engine =
+        RustSignalEngine::new_with_pool(bob_db.pool.clone(), bob_identity_bytes, 5678)?;
+
+    // 4. Bob generates a bundle
+    let bob_bundle = bob_engine.generate_bundle(1, 1).await?;
+
+    // 5. Alice processes Bob's bundle
+    alice_engine
+        .process_prekey_bundle("bob".to_string(), 1, bob_bundle)
+        .await?;
+
+    // 6. Exchange 100 messages
+    let mut alice_to_bob = true;
+    for i in 1..=100 {
+        if alice_to_bob {
+            let plaintext = format!("Message {} from Alice to Bob", i).into_bytes();
+
+            let ciphertext = alice_engine
+                .encrypt_message("bob".to_string(), 1, plaintext.clone())
+                .await?;
+
+            let is_prekey = i == 1; // Only the first message is a PreKeySignalMessage
+            let decrypted = bob_engine
+                .decrypt_message("alice".to_string(), 1, ciphertext, is_prekey)
+                .await?;
+
+            assert_eq!(plaintext, decrypted);
+        } else {
+            let plaintext = format!("Message {} from Bob to Alice", i).into_bytes();
+
+            let ciphertext = bob_engine
+                .encrypt_message("alice".to_string(), 1, plaintext.clone())
+                .await?;
+
+            let decrypted = alice_engine
+                .decrypt_message("bob".to_string(), 1, ciphertext, false)
+                .await?;
+
+            assert_eq!(plaintext, decrypted);
+        }
+        alice_to_bob = !alice_to_bob;
+    }
+
+    Ok(())
+}