Skip to content

Commit

Permalink
Move generate_pre_keys into pre_keys.rs
Browse files Browse the repository at this point in the history
  • Loading branch information
@rubdos
rubdos committed Jan 30, 2024
1 parent 1c3d6dd commit c2ceb80
Show file tree
Hide file tree
Showing 3 changed files with 194 additions and 202 deletions.
201 changes: 26 additions & 175 deletions libsignal-service/src/account_manager.rs
View file
Original file line number Diff line number Diff line change
@@ -1,23 +1,23 @@
use base64::prelude::*;
use std::collections::HashMap;
use std::convert::{TryFrom, TryInto};
use std::time::SystemTime;
use std::convert::TryInto;

use aes::cipher::{KeyIvInit, StreamCipher as _};
use hmac::digest::Output;
use hmac::{Hmac, Mac};
use libsignal_protocol::{
kem, GenericSignedPreKey, IdentityKeyStore, KeyPair, KyberPreKeyRecord,
PreKeyRecord, PrivateKey, PublicKey, SignalProtocolError,
SignedPreKeyRecord,
IdentityKeyStore, KeyPair, PrivateKey, PublicKey, SignalProtocolError,
};
use prost::Message;
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use tracing_futures::Instrument;
use zkgroup::profiles::ProfileKey;

use crate::pre_keys::{KyberPreKeyEntity, PreKeysStore, SignedPreKeyEntity};
use crate::pre_keys::{
KyberPreKeyEntity, PreKeysStore, SignedPreKeyEntity, PRE_KEY_BATCH_SIZE,
PRE_KEY_MINIMUM,
};
use crate::proto::DeviceName;
use crate::provisioning::generate_registration_id;
use crate::push_service::{AvatarWrite, RecaptchaAttributes, ServiceIdType};
Expand All @@ -26,7 +26,7 @@ use crate::session_store::SessionStoreExt;
use crate::ServiceAddress;
use crate::{
configuration::{Endpoint, ServiceCredentials},
pre_keys::{PreKeyEntity, PreKeyState},
pre_keys::PreKeyState,
profile_cipher::{ProfileCipher, ProfileCipherError},
profile_name::ProfileName,
proto::{ProvisionEnvelope, ProvisionMessage, ProvisioningVersion},
Expand Down Expand Up @@ -73,10 +73,6 @@ pub struct Profile {
pub about_emoji: Option<String>,
}

const PRE_KEY_MINIMUM: u32 = 10;
const PRE_KEY_BATCH_SIZE: u32 = 100;
const PRE_KEY_MEDIUM_MAX_VALUE: u32 = 0xFFFFFF;

impl<Service: PushService> AccountManager<Service> {
pub fn new(service: Service, profile_key: Option<ProfileKey>) -> Self {
Self {
Expand All @@ -85,150 +81,6 @@ impl<Service: PushService> AccountManager<Service> {
}
}

#[tracing::instrument(skip(self, protocol_store, csprng))]
// XXX: Maybe refactor this and move into pre_keys.rs
pub(crate) async fn generate_pre_keys<
R: rand::Rng + rand::CryptoRng,
P: PreKeysStore,
>(
&mut self,
protocol_store: &mut P,
service_id_type: ServiceIdType,
csprng: &mut R,
use_last_resort_key: bool,
pre_key_count: u32,
kyber_pre_key_count: u32,
) -> Result<
(
Vec<PreKeyEntity>,
SignedPreKeyEntity,
Vec<KyberPreKeyEntity>,
Option<KyberPreKeyEntity>,
),
ServiceError,
> {
let pre_keys_offset_id = protocol_store.next_pre_key_id().await?;
let next_signed_pre_key_id =
protocol_store.next_signed_pre_key_id().await?;
let pq_pre_keys_offset_id = protocol_store.next_pq_pre_key_id().await?;

let span = tracing::span!(tracing::Level::DEBUG, "Generating pre keys");

let identity_key_pair = protocol_store
.get_identity_key_pair()
.instrument(
tracing::trace_span!(parent: &span, "get identity key pair"),
)
.await?;

let mut pre_key_entities = vec![];
let mut pq_pre_key_entities = vec![];

// EC keys
for i in 0..pre_key_count {
let key_pair = KeyPair::generate(csprng);
let pre_key_id = (((pre_keys_offset_id + i)
% (PRE_KEY_MEDIUM_MAX_VALUE - 1))
+ 1)
.into();
let pre_key_record = PreKeyRecord::new(pre_key_id, &key_pair);
protocol_store
.save_pre_key(pre_key_id, &pre_key_record)
.instrument(tracing::trace_span!(parent: &span, "save pre key", ?pre_key_id)).await?;
// TODO: Shouldn't this also remove the previous pre-keys from storage?
// I think we might want to update the storage, and then sync the storage to the
// server.

pre_key_entities.push(PreKeyEntity::try_from(pre_key_record)?);
}

// Kyber keys
for i in 0..kyber_pre_key_count {
let pre_key_id = (((pq_pre_keys_offset_id + i)
% (PRE_KEY_MEDIUM_MAX_VALUE - 1))
+ 1)
.into();
let pre_key_record = KyberPreKeyRecord::generate(
kem::KeyType::Kyber1024,
pre_key_id,
identity_key_pair.private_key(),
)?;
protocol_store
.save_kyber_pre_key(pre_key_id, &pre_key_record)
.instrument(tracing::trace_span!(parent: &span, "save kyber pre key", ?pre_key_id)).await?;
// TODO: Shouldn't this also remove the previous pre-keys from storage?
// I think we might want to update the storage, and then sync the storage to the
// server.

pq_pre_key_entities
.push(KyberPreKeyEntity::try_from(pre_key_record)?);
}

// Generate and store the next signed prekey
let signed_pre_key_pair = KeyPair::generate(csprng);
let signed_pre_key_public = signed_pre_key_pair.public_key;
let signed_pre_key_signature = identity_key_pair
.private_key()
.calculate_signature(&signed_pre_key_public.serialize(), csprng)?;

let unix_time = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();

let signed_prekey_record = SignedPreKeyRecord::new(
next_signed_pre_key_id.into(),
unix_time.as_millis() as u64,
&signed_pre_key_pair,
&signed_pre_key_signature,
);

protocol_store
.save_signed_pre_key(
next_signed_pre_key_id.into(),
&signed_prekey_record,
)
.instrument(tracing::trace_span!(parent: &span, "save signed pre key", signed_pre_key_id = ?next_signed_pre_key_id)).await?;
let signed_prekey_entity =
SignedPreKeyEntity::try_from(signed_prekey_record)?;

let pq_last_resort_key = if use_last_resort_key {
let pre_key_id = (((pq_pre_keys_offset_id + kyber_pre_key_count)
% (PRE_KEY_MEDIUM_MAX_VALUE - 1))
+ 1)
.into();

if !pq_pre_key_entities.is_empty() {
assert_eq!(
pq_pre_key_entities.last().unwrap().key_id + 1,
u32::from(pre_key_id)
);
}

let pre_key_record = KyberPreKeyRecord::generate(
kem::KeyType::Kyber1024,
pre_key_id,
identity_key_pair.private_key(),
)?;
protocol_store
.store_last_resort_kyber_pre_key(pre_key_id, &pre_key_record)
.instrument(tracing::trace_span!(parent: &span, "save last resort kyber pre key", ?pre_key_id)).await?;
// TODO: Shouldn't this also remove the previous pre-keys from storage?
// I think we might want to update the storage, and then sync the storage to the
// server.

Some(KyberPreKeyEntity::try_from(pre_key_record)?)
} else {
None
};

Ok((
pre_key_entities,
signed_prekey_entity,
pq_pre_key_entities,
pq_last_resort_key,
))
}

/// Checks the availability of pre-keys, and updates them as necessary.
///
/// Parameters are the protocol's `StoreContext`, and the offsets for the next pre-key and
Expand Down Expand Up @@ -279,24 +131,24 @@ impl<Service: PushService> AccountManager<Service> {
return Ok(());
}

let (pre_keys, signed_pre_key, pq_pre_keys, pq_last_resort_key) = self
.generate_pre_keys(
let identity_key_pair = protocol_store
.get_identity_key_pair()
.instrument(tracing::trace_span!("get identity key pair"))
.await?;

let (pre_keys, signed_pre_key, pq_pre_keys, pq_last_resort_key) =
crate::pre_keys::generate_pre_keys(
protocol_store,
service_id_type,
&identity_key_pair,
csprng,
use_last_resort_key,
PRE_KEY_BATCH_SIZE,
PRE_KEY_BATCH_SIZE,
)
.await?;

let identity_key = protocol_store
.get_identity_key_pair()
.instrument(tracing::trace_span!("get identity key pair"))
.await?
.identity_key()
.public_key()
.clone();
let identity_key =
identity_key_pair.identity_key().public_key().clone();

Check warning on line 151 in libsignal-service/src/account_manager.rs

View workflow job for this annotation

GitHub Actions / clippy

using `clone` on type `PublicKey` which implements the `Copy` trait 

warning: using `clone` on type `PublicKey` which implements the `Copy` trait
   --> libsignal-service/src/account_manager.rs:151:13
    |
151 |             identity_key_pair.identity_key().public_key().clone();
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: try dereferencing it: `*identity_key_pair.identity_key().public_key()`
    |
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#clone_on_copy
    = note: `#[warn(clippy::clone_on_copy)]` on by default

let pre_key_state = PreKeyState {
pre_keys,
Expand Down Expand Up @@ -680,16 +532,15 @@ impl<Service: PushService> AccountManager<Service> {
signed_pre_key_entity,
_kyber_pre_key_entities,
last_resort_kyber_prekey,
) = self
.generate_pre_keys(
pni_protocol_store,
ServiceIdType::PhoneNumberIdentity,
csprng,
true,
0,
0,
)
.await?;
) = crate::pre_keys::generate_pre_keys(
pni_protocol_store,
&pni_identity_key_pair,
csprng,
true,
0,
0,
)
.await?;
let registration_id = generate_registration_id(csprng);

let local_device_id_s = local_device_id.to_string();
Expand Down
Loading

0 comments on commit c2ceb80

Please sign in to comment.