mirror/libsignal
0
0
Fork 0
mirror of https://github.com/signalapp/libsignal.git synced 2024-09-20 12:02:18 +02:00
Code

Add a simple benchmark

Browse Source 
Shuffle some of the test support code into support/mod.rs so it can
be used in benchmarks also
This commit is contained in:
Jack Lloyd 2020-08-03 19:34:14 -04:00
parent 2d8098ad84
commit 993f845301
6 changed files with 206 additions and 136 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
Cargo.toml
View File

@ -22,6 +22,11 @@ block-modes = "0.5"
[dev-dependencies]
hex = "0.4"
criterion = "0.3"
[build-dependencies]
prost-build = "0.6"
[[bench]]
name = "session"
harness = false

53
benches/session.rs Normal file
View File

@ -0,0 +1,53 @@
use criterion::{criterion_group, criterion_main, Criterion};
use libsignal_protocol_rust::*;
#[path = "../tests/support/mod.rs"]
mod support;
pub fn session_encrypt_decrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolError> {
let (alice_session, bob_session) = support::initialize_sessions_v3()?;
let alice_session_record = SessionRecord::new(alice_session);
let bob_session_record = SessionRecord::new(bob_session);
let alice_address = ProtocolAddress::new("+14159999999".to_owned(), 1);
let bob_address = ProtocolAddress::new("+14158888888".to_owned(), 1);
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
alice_store.store_session(&bob_address, &alice_session_record)?;
bob_store.store_session(&alice_address, &bob_session_record)?;
let message_to_decrypt = support::encrypt(
&mut alice_store,
&bob_address,
"a short message")?;
c.bench_function("session encrypt", |b| {
b.iter(|| {
let mut alice_store = alice_store.clone();
support::encrypt(
&mut alice_store,
&bob_address,
"a short message").expect("success");
})});
c.bench_function("session decrypt", |b| {
b.iter(|| {
let mut bob_store = bob_store.clone();
support::decrypt(
&mut bob_store,
&alice_address,
&message_to_decrypt).expect("success");
})});
Ok(())
}
pub fn session_encrypt_decrypt(mut c: &mut Criterion) {
session_encrypt_decrypt_result(&mut c).expect("success");
}
criterion_group!(benches, session_encrypt_decrypt);
criterion_main!(benches);

6
src/storage/inmem.rs
View File

@ -5,6 +5,7 @@ use crate::{IdentityKey, IdentityKeyPair, ProtocolAddress, SenderKeyName, Sender
use std::collections::HashMap;
#[derive(Clone)]
pub struct InMemIdentityKeyStore {
key_pair: IdentityKeyPair,
id: u32,
@ -68,6 +69,7 @@ impl traits::IdentityKeyStore for InMemIdentityKeyStore {
}
}
#[derive(Clone)]
pub struct InMemPreKeyStore {
pre_keys: HashMap<PreKeyId, PreKeyRecord>,
}
@ -112,6 +114,7 @@ impl traits::PreKeyStore for InMemPreKeyStore {
}
}
#[derive(Clone)]
pub struct InMemSignedPreKeyStore {
signed_pre_keys: HashMap<SignedPreKeyId, SignedPreKeyRecord>,
}
@ -168,6 +171,7 @@ impl traits::SignedPreKeyStore for InMemSignedPreKeyStore {
}
}
#[derive(Clone)]
pub struct InMemSessionStore {
sessions: HashMap<ProtocolAddress, SessionRecord>,
}
@ -226,6 +230,7 @@ impl traits::SessionStore for InMemSessionStore {
}
}
#[derive(Clone)]
pub struct InMemSenderKeyStore {
keys: HashMap<SenderKeyName, SenderKeyRecord>,
}
@ -262,6 +267,7 @@ impl traits::SenderKeyStore for InMemSenderKeyStore {
}
}
#[derive(Clone)]
pub struct InMemSignalProtocolStore {
pub session_store: InMemSessionStore,
pub pre_key_store: InMemPreKeyStore,

35
tests/groups.rs
View File

@ -5,6 +5,7 @@ use rand::rngs::OsRng;
use rand::seq::SliceRandom;
use rand::Rng;
use std::convert::TryFrom;
use support::test_in_memory_protocol_store;
#[test]
fn group_no_send_session() -> Result<(), SignalProtocolError> {
@ -14,7 +15,7 @@ fn group_no_send_session() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
assert!(group_encrypt(
&mut alice_store,
@ -35,8 +36,8 @@ fn group_no_recv_session() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -66,8 +67,8 @@ fn group_basic_encrypt_decrypt() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -103,8 +104,8 @@ fn group_large_messages() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -141,8 +142,8 @@ fn group_basic_ratchet() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -200,8 +201,8 @@ fn group_late_join() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -246,8 +247,8 @@ fn group_out_of_order() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -300,8 +301,8 @@ fn group_too_far_in_the_future() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;
@ -344,8 +345,8 @@ fn group_message_key_limit() -> Result<(), SignalProtocolError> {
let group_sender =
SenderKeyName::new("summer camp planning committee".to_owned(), sender_address)?;
let mut alice_store = support::test_in_memory_protocol_store();
let mut bob_store = support::test_in_memory_protocol_store();
let mut alice_store = test_in_memory_protocol_store();
let mut bob_store = test_in_memory_protocol_store();
let sent_distribution_message =
create_sender_key_distribution_message(&group_sender, &mut alice_store, &mut csprng)?;

126
tests/session.rs
View File

@ -1,84 +1,9 @@
mod support;
use libsignal_protocol_rust::*;
use rand::{rngs::OsRng, CryptoRng, Rng};
use rand::rngs::OsRng;
use std::convert::TryFrom;
fn encrypt(
store: &mut InMemSignalProtocolStore,
remote_address: &ProtocolAddress,
msg: &str,
) -> Result<CiphertextMessage, SignalProtocolError> {
let mut session_cipher = SessionCipher::new(
remote_address.clone(),
&mut store.session_store,
&mut store.identity_store,
&mut store.signed_pre_key_store,
&mut store.pre_key_store,
);
session_cipher.encrypt(msg.as_bytes())
}
fn decrypt(
store: &mut InMemSignalProtocolStore,
remote_address: &ProtocolAddress,
msg: &CiphertextMessage,
) -> Result<Vec<u8>, SignalProtocolError> {
let mut session_cipher = SessionCipher::new(
remote_address.clone(),
&mut store.session_store,
&mut store.identity_store,
&mut store.signed_pre_key_store,
&mut store.pre_key_store,
);
let mut csprng = OsRng;
session_cipher.decrypt(msg, &mut csprng)
}
fn create_pre_key_bundle<R: Rng + CryptoRng>(
store: &mut dyn ProtocolStore,
mut csprng: &mut R,
) -> Result<PreKeyBundle, SignalProtocolError> {
let pre_key_pair = KeyPair::new(&mut csprng);
let signed_pre_key_pair = KeyPair::new(&mut csprng);
let signed_pre_key_public = signed_pre_key_pair.public_key.serialize();
let signed_pre_key_signature = store
.get_identity_key_pair()?
.private_key()
.calculate_signature(&signed_pre_key_public, &mut csprng)?;
let device_id: u32 = csprng.gen();
let pre_key_id: u32 = csprng.gen();
let signed_pre_key_id: u32 = csprng.gen();
let pre_key_bundle = PreKeyBundle::new(
store.get_local_registration_id()?,
device_id,
Some(pre_key_id),
Some(pre_key_pair.public_key),
signed_pre_key_id,
signed_pre_key_pair.public_key,
signed_pre_key_signature.to_vec(),
*store.get_identity_key_pair()?.identity_key(),
)?;
store.save_pre_key(pre_key_id, &PreKeyRecord::new(pre_key_id, &pre_key_pair))?;
let timestamp = csprng.gen();
store.save_signed_pre_key(
signed_pre_key_id,
&SignedPreKeyRecord::new(
signed_pre_key_id,
timestamp,
&signed_pre_key_pair,
&signed_pre_key_signature,
),
)?;
Ok(pre_key_bundle)
}
use support::*;
#[test]
fn test_basic_prekey_v3() -> Result<(), SignalProtocolError> {
@ -637,41 +562,6 @@ fn optional_one_time_prekey() -> Result<(), SignalProtocolError> {
Ok(())
}
fn initialize_sessions_v3() -> Result<(SessionState, SessionState), SignalProtocolError> {
let mut csprng = OsRng;
let alice_identity = IdentityKeyPair::generate(&mut csprng);
let bob_identity = IdentityKeyPair::generate(&mut csprng);
let alice_base_key = KeyPair::new(&mut csprng);
let bob_base_key = KeyPair::new(&mut csprng);
let bob_ephemeral_key = bob_base_key;
let alice_params = AliceSignalProtocolParameters::new(
alice_identity,
alice_base_key,
*bob_identity.identity_key(),
bob_base_key.public_key,
None,
bob_ephemeral_key.public_key,
);
let alice_session = initialize_alice_session(&alice_params, &mut csprng)?;
let bob_params = BobSignalProtocolParameters::new(
bob_identity,
bob_base_key,
None,
bob_ephemeral_key,
*alice_identity.identity_key(),
alice_base_key.public_key,
);
let bob_session = initialize_bob_session(&bob_params)?;
Ok((alice_session, bob_session))
}
#[test]
fn basic_session_v3() -> Result<(), SignalProtocolError> {
let (alice_session, bob_session) = initialize_sessions_v3()?;
@ -893,12 +783,12 @@ fn is_session_id_equal(
.load_session(bob_address)?
.unwrap()
.session_state()?
.alice_base_key()
== bob_store
.load_session(alice_address)?
.unwrap()
.session_state()?
.alice_base_key())
.alice_base_key().clone() ==
bob_store
.load_session(alice_address)?
.unwrap()
.session_state()?
.alice_base_key().clone())
}
#[test]

117
tests/support/mod.rs
View File

@ -1,5 +1,5 @@
use libsignal_protocol_rust::*;
use rand::rngs::OsRng;
use rand::{rngs::OsRng, CryptoRng, Rng};
pub fn test_in_memory_protocol_store() -> InMemSignalProtocolStore {
let mut csprng = OsRng;
@ -8,3 +8,118 @@ pub fn test_in_memory_protocol_store() -> InMemSignalProtocolStore {
InMemSignalProtocolStore::new(identity_key, registration_id).unwrap()
}
#[allow(dead_code)]
pub fn encrypt(
store: &mut InMemSignalProtocolStore,
remote_address: &ProtocolAddress,
msg: &str,
) -> Result<CiphertextMessage, SignalProtocolError> {
let mut session_cipher = SessionCipher::new(
remote_address.clone(),
&mut store.session_store,
&mut store.identity_store,
&mut store.signed_pre_key_store,
&mut store.pre_key_store,
);
session_cipher.encrypt(msg.as_bytes())
}
#[allow(dead_code)]
pub fn decrypt(
store: &mut InMemSignalProtocolStore,
remote_address: &ProtocolAddress,
msg: &CiphertextMessage,
) -> Result<Vec<u8>, SignalProtocolError> {
let mut session_cipher = SessionCipher::new(
remote_address.clone(),
&mut store.session_store,
&mut store.identity_store,
&mut store.signed_pre_key_store,
&mut store.pre_key_store,
);
let mut csprng = OsRng;
session_cipher.decrypt(msg, &mut csprng)
}
#[allow(dead_code)]
pub fn create_pre_key_bundle<R: Rng + CryptoRng>(
store: &mut dyn ProtocolStore,
mut csprng: &mut R,
) -> Result<PreKeyBundle, SignalProtocolError> {
let pre_key_pair = KeyPair::new(&mut csprng);
let signed_pre_key_pair = KeyPair::new(&mut csprng);
let signed_pre_key_public = signed_pre_key_pair.public_key.serialize();
let signed_pre_key_signature = store
.get_identity_key_pair()?
.private_key()
.calculate_signature(&signed_pre_key_public, &mut csprng)?;
let device_id: u32 = csprng.gen();
let pre_key_id: u32 = csprng.gen();
let signed_pre_key_id: u32 = csprng.gen();
let pre_key_bundle = PreKeyBundle::new(
store.get_local_registration_id()?,
device_id,
Some(pre_key_id),
Some(pre_key_pair.public_key),
signed_pre_key_id,
signed_pre_key_pair.public_key,
signed_pre_key_signature.to_vec(),
*store.get_identity_key_pair()?.identity_key(),
)?;
store.save_pre_key(pre_key_id, &PreKeyRecord::new(pre_key_id, &pre_key_pair))?;
let timestamp = csprng.gen();
store.save_signed_pre_key(
signed_pre_key_id,
&SignedPreKeyRecord::new(
signed_pre_key_id,
timestamp,
&signed_pre_key_pair,
&signed_pre_key_signature,
),
)?;
Ok(pre_key_bundle)
}
#[allow(dead_code)]
pub fn initialize_sessions_v3() -> Result<(SessionState, SessionState), SignalProtocolError> {
let mut csprng = OsRng;
let alice_identity = IdentityKeyPair::generate(&mut csprng);
let bob_identity = IdentityKeyPair::generate(&mut csprng);
let alice_base_key = KeyPair::new(&mut csprng);
let bob_base_key = KeyPair::new(&mut csprng);
let bob_ephemeral_key = bob_base_key;
let alice_params = AliceSignalProtocolParameters::new(
alice_identity,
alice_base_key,
*bob_identity.identity_key(),
bob_base_key.public_key,
None,
bob_ephemeral_key.public_key,
);
let alice_session = initialize_alice_session(&alice_params, &mut csprng)?;
let bob_params = BobSignalProtocolParameters::new(
bob_identity,
bob_base_key,
None,
bob_ephemeral_key,
*alice_identity.identity_key(),
alice_base_key.public_key,
);
let bob_session = initialize_bob_session(&bob_params)?;
Ok((alice_session, bob_session))
}