use matrix_sdk_common::events::ToDeviceEvent;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
collections::BTreeMap,
convert::{TryFrom, TryInto},
fmt,
ops::Deref,
sync::{
atomic::{AtomicBool, AtomicI64, Ordering},
Arc,
},
};
use tracing::{debug, trace, warn};
#[cfg(test)]
use matrix_sdk_common::events::EventType;
use matrix_sdk_common::{
api::r0::keys::{
upload_keys, upload_signatures::Request as SignatureUploadRequest, OneTimeKey, SignedKey,
},
encryption::DeviceKeys,
events::{room::encrypted::EncryptedEventContent, AnyToDeviceEvent},
identifiers::{
DeviceId, DeviceIdBox, DeviceKeyAlgorithm, DeviceKeyId, EventEncryptionAlgorithm, RoomId,
UserId,
},
instant::Instant,
locks::Mutex,
CanonicalJsonValue, Raw, UInt,
};
use olm_rs::{
account::{IdentityKeys, OlmAccount, OneTimeKeys},
errors::{OlmAccountError, OlmSessionError},
session::{OlmMessage, PreKeyMessage},
PicklingMode,
};
use crate::{
error::{EventError, OlmResult, SessionCreationError},
identities::ReadOnlyDevice,
requests::UploadSigningKeysRequest,
store::{Changes, Store},
utilities::encode,
OlmError,
};
use super::{
EncryptionSettings, InboundGroupSession, OutboundGroupSession, PrivateCrossSigningIdentity,
Session,
};
#[derive(Debug, Clone)]
pub struct Account {
pub(crate) inner: ReadOnlyAccount,
pub(crate) store: Store,
}
#[derive(Debug, Clone)]
pub enum SessionType {
New(Session),
Existing(Session),
}
impl SessionType {
#[cfg(test)]
pub fn session(self) -> Session {
match self {
SessionType::New(s) => s,
SessionType::Existing(s) => s,
}
}
}
#[derive(Debug, Clone)]
pub struct OlmDecryptionInfo {
pub session: SessionType,
pub message_hash: OlmMessageHash,
pub event: Raw<AnyToDeviceEvent>,
pub signing_key: String,
pub sender_key: String,
pub inbound_group_session: Option<InboundGroupSession>,
}
#[derive(Debug, Clone)]
pub struct OlmMessageHash {
pub sender_key: String,
pub hash: String,
}
impl Deref for Account {
type Target = ReadOnlyAccount;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl Account {
pub async fn decrypt_to_device_event(
&self,
event: &ToDeviceEvent<EncryptedEventContent>,
) -> OlmResult<OlmDecryptionInfo> {
debug!("Decrypting to-device event");
let content = if let EncryptedEventContent::OlmV1Curve25519AesSha2(c) = &event.content {
c
} else {
warn!("Error, unsupported encryption algorithm");
return Err(EventError::UnsupportedAlgorithm.into());
};
let identity_keys = self.inner.identity_keys();
let own_key = identity_keys.curve25519();
let own_ciphertext = content.ciphertext.get(own_key);
if let Some(ciphertext) = own_ciphertext {
let message_type: u8 = ciphertext
.message_type
.try_into()
.map_err(|_| EventError::UnsupportedOlmType)?;
let sha = Sha256::new()
.chain(&content.sender_key)
.chain(&[message_type])
.chain(&ciphertext.body);
let message_hash = OlmMessageHash {
sender_key: content.sender_key.clone(),
hash: encode(sha.finalize().as_slice()),
};
let message =
OlmMessage::from_type_and_ciphertext(message_type.into(), ciphertext.body.clone())
.map_err(|_| EventError::UnsupportedOlmType)?;
let (session, event, signing_key) = match self
.decrypt_olm_message(&event.sender, &content.sender_key, message)
.await
{
Ok(d) => d,
Err(OlmError::SessionWedged(user_id, sender_key)) => {
if self.store.is_message_known(&message_hash).await? {
return Err(OlmError::ReplayedMessage(user_id, sender_key));
} else {
return Err(OlmError::SessionWedged(user_id, sender_key));
}
}
Err(e) => return Err(e),
};
debug!("Decrypted a to-device event {:?}", event);
Ok(OlmDecryptionInfo {
session,
message_hash,
event,
signing_key,
sender_key: content.sender_key.clone(),
inbound_group_session: None,
})
} else {
warn!("Olm event doesn't contain a ciphertext for our key");
Err(EventError::MissingCiphertext.into())
}
}
pub async fn update_uploaded_key_count(&self, key_count: &BTreeMap<DeviceKeyAlgorithm, UInt>) {
let one_time_key_count = key_count.get(&DeviceKeyAlgorithm::SignedCurve25519);
let count: u64 = one_time_key_count.map_or(0, |c| (*c).into());
self.inner.update_uploaded_key_count(count);
}
pub async fn receive_keys_upload_response(
&self,
response: &upload_keys::Response,
) -> OlmResult<()> {
if !self.inner.shared() {
debug!("Marking account as shared");
}
self.inner.mark_as_shared();
let one_time_key_count = response
.one_time_key_counts
.get(&DeviceKeyAlgorithm::SignedCurve25519);
let count: u64 = one_time_key_count.map_or(0, |c| (*c).into());
debug!(
"Updated uploaded one-time key count {} -> {}, marking keys as published",
self.inner.uploaded_key_count(),
count
);
self.inner.update_uploaded_key_count(count);
self.inner.mark_keys_as_published().await;
self.store.save_account(self.inner.clone()).await?;
Ok(())
}
async fn try_decrypt_olm_message(
&self,
sender: &UserId,
sender_key: &str,
message: &OlmMessage,
) -> OlmResult<Option<(Session, String)>> {
let s = self.store.get_sessions(sender_key).await?;
let sessions = if let Some(s) = s {
s
} else {
return Ok(None);
};
let mut decrypted: Option<(Session, String)> = None;
for session in &mut *sessions.lock().await {
let mut matches = false;
if let OlmMessage::PreKey(m) = &message {
matches = session.matches(sender_key, m.clone()).await?;
if !matches {
continue;
}
}
let ret = session.decrypt(message.clone()).await;
match ret {
Ok(p) => {
decrypted = Some((session.clone(), p));
break;
}
Err(e) => {
if matches {
warn!(
"Found a matching Olm session yet decryption failed
for sender {} and sender_key {} {:?}",
sender, sender_key, e
);
return Err(OlmError::SessionWedged(
sender.to_owned(),
sender_key.to_owned(),
));
}
}
}
}
Ok(decrypted)
}
async fn decrypt_olm_message(
&self,
sender: &UserId,
sender_key: &str,
message: OlmMessage,
) -> OlmResult<(SessionType, Raw<AnyToDeviceEvent>, String)> {
let (session, plaintext) = if let Some(d) = self
.try_decrypt_olm_message(sender, sender_key, &message)
.await?
{
(SessionType::Existing(d.0), d.1)
} else {
let mut session = match &message {
OlmMessage::Message(_) => {
warn!(
"Failed to decrypt a non-pre-key message with all \
available sessions {} {}",
sender, sender_key
);
return Err(OlmError::SessionWedged(
sender.to_owned(),
sender_key.to_owned(),
));
}
OlmMessage::PreKey(m) => {
let session = match self
.inner
.create_inbound_session(sender_key, m.clone())
.await
{
Ok(s) => s,
Err(e) => {
warn!(
"Failed to create a new Olm session for {} {}
from a prekey message: {}",
sender, sender_key, e
);
return Err(OlmError::SessionWedged(
sender.to_owned(),
sender_key.to_owned(),
));
}
};
session
}
};
let plaintext = session.decrypt(message).await?;
(SessionType::New(session), plaintext)
};
trace!("Successfully decrypted a Olm message: {}", plaintext);
let (event, signing_key) = match self.parse_decrypted_to_device_event(sender, &plaintext) {
Ok(r) => r,
Err(e) => {
match session {
SessionType::New(s) => {
let changes = Changes {
account: Some(self.inner.clone()),
sessions: vec![s],
..Default::default()
};
self.store.save_changes(changes).await?;
}
SessionType::Existing(s) => {
self.store.save_sessions(&[s]).await?;
}
}
return Err(e);
}
};
Ok((session, event, signing_key))
}
fn parse_decrypted_to_device_event(
&self,
sender: &UserId,
plaintext: &str,
) -> OlmResult<(Raw<AnyToDeviceEvent>, String)> {
let decrypted_json: Value = serde_json::from_str(&plaintext)?;
let encrytped_sender = decrypted_json
.get("sender")
.cloned()
.ok_or_else(|| EventError::MissingField("sender".to_string()))?;
let encrytped_sender: UserId = serde_json::from_value(encrytped_sender)?;
let recipient = decrypted_json
.get("recipient")
.cloned()
.ok_or_else(|| EventError::MissingField("recipient".to_string()))?;
let recipient: UserId = serde_json::from_value(recipient)?;
let recipient_keys: BTreeMap<DeviceKeyAlgorithm, String> = serde_json::from_value(
decrypted_json
.get("recipient_keys")
.cloned()
.ok_or_else(|| EventError::MissingField("recipient_keys".to_string()))?,
)?;
let keys: BTreeMap<DeviceKeyAlgorithm, String> = serde_json::from_value(
decrypted_json
.get("keys")
.cloned()
.ok_or_else(|| EventError::MissingField("keys".to_string()))?,
)?;
if &recipient != self.user_id() || sender != &encrytped_sender {
return Err(EventError::MissmatchedSender.into());
}
if self.inner.identity_keys().ed25519()
!= recipient_keys
.get(&DeviceKeyAlgorithm::Ed25519)
.ok_or(EventError::MissingSigningKey)?
{
return Err(EventError::MissmatchedKeys.into());
}
let signing_key = keys
.get(&DeviceKeyAlgorithm::Ed25519)
.ok_or(EventError::MissingSigningKey)?;
Ok((
Raw::from(serde_json::from_value::<AnyToDeviceEvent>(decrypted_json)?),
signing_key.to_owned(),
))
}
}
#[derive(Clone)]
pub struct ReadOnlyAccount {
pub(crate) user_id: Arc<UserId>,
pub(crate) device_id: Arc<Box<DeviceId>>,
inner: Arc<Mutex<OlmAccount>>,
pub(crate) identity_keys: Arc<IdentityKeys>,
shared: Arc<AtomicBool>,
uploaded_signed_key_count: Arc<AtomicI64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccountPickle(String);
impl AccountPickle {
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<String> for AccountPickle {
fn from(value: String) -> Self {
Self(value)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PickledAccount {
pub user_id: UserId,
pub device_id: DeviceIdBox,
pub pickle: AccountPickle,
pub shared: bool,
pub uploaded_signed_key_count: i64,
}
#[cfg(not(tarpaulin_include))]
impl fmt::Debug for ReadOnlyAccount {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Account")
.field("identity_keys", self.identity_keys())
.field("shared", &self.shared())
.finish()
}
}
impl ReadOnlyAccount {
const ALGORITHMS: &'static [&'static EventEncryptionAlgorithm] = &[
&EventEncryptionAlgorithm::OlmV1Curve25519AesSha2,
&EventEncryptionAlgorithm::MegolmV1AesSha2,
];
#[allow(clippy::ptr_arg)]
pub fn new(user_id: &UserId, device_id: &DeviceId) -> Self {
let account = OlmAccount::new();
let identity_keys = account.parsed_identity_keys();
Self {
user_id: Arc::new(user_id.to_owned()),
device_id: Arc::new(device_id.into()),
inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::new(false)),
uploaded_signed_key_count: Arc::new(AtomicI64::new(0)),
}
}
pub fn user_id(&self) -> &UserId {
&self.user_id
}
pub fn device_id(&self) -> &DeviceId {
&self.device_id
}
pub fn identity_keys(&self) -> &IdentityKeys {
&self.identity_keys
}
pub(crate) fn update_uploaded_key_count(&self, new_count: u64) {
let key_count = i64::try_from(new_count).unwrap_or(i64::MAX);
self.uploaded_signed_key_count
.store(key_count, Ordering::Relaxed);
}
pub fn uploaded_key_count(&self) -> i64 {
self.uploaded_signed_key_count.load(Ordering::Relaxed)
}
pub fn shared(&self) -> bool {
self.shared.load(Ordering::Relaxed)
}
pub(crate) fn mark_as_shared(&self) {
self.shared.store(true, Ordering::Relaxed);
}
pub(crate) async fn one_time_keys(&self) -> OneTimeKeys {
self.inner.lock().await.parsed_one_time_keys()
}
pub(crate) async fn generate_one_time_keys_helper(&self, count: usize) {
self.inner.lock().await.generate_one_time_keys(count);
}
pub(crate) async fn max_one_time_keys(&self) -> usize {
self.inner.lock().await.max_number_of_one_time_keys()
}
pub(crate) async fn generate_one_time_keys(&self) -> Result<u64, ()> {
let count = self.uploaded_key_count() as u64;
let max_keys = self.max_one_time_keys().await;
let max_on_server = (max_keys as u64) / 2;
if count >= (max_on_server) {
return Err(());
}
let key_count = (max_on_server) - count;
let key_count: usize = key_count.try_into().unwrap_or(max_keys);
self.generate_one_time_keys_helper(key_count).await;
Ok(key_count as u64)
}
pub(crate) async fn should_upload_keys(&self) -> bool {
if !self.shared() {
return true;
}
let count = self.uploaded_key_count() as u64;
let max_keys = self.max_one_time_keys().await as u64;
if count > (max_keys / 2) {
return false;
}
let key_count = (max_keys / 2) - count;
key_count > 0
}
pub(crate) async fn keys_for_upload(
&self,
) -> Option<(
Option<DeviceKeys>,
Option<BTreeMap<DeviceKeyId, OneTimeKey>>,
)> {
if !self.should_upload_keys().await {
return None;
}
let device_keys = if !self.shared() {
Some(self.device_keys().await)
} else {
None
};
let one_time_keys = self.signed_one_time_keys().await.ok();
Some((device_keys, one_time_keys))
}
pub(crate) async fn mark_keys_as_published(&self) {
self.inner.lock().await.mark_keys_as_published();
}
pub async fn sign(&self, string: &str) -> String {
self.inner.lock().await.sign(string)
}
pub async fn pickle(&self, pickle_mode: PicklingMode) -> PickledAccount {
let pickle = AccountPickle(self.inner.lock().await.pickle(pickle_mode));
PickledAccount {
user_id: self.user_id().to_owned(),
device_id: self.device_id().to_owned(),
pickle,
shared: self.shared(),
uploaded_signed_key_count: self.uploaded_key_count(),
}
}
pub fn from_pickle(
pickle: PickledAccount,
pickle_mode: PicklingMode,
) -> Result<Self, OlmAccountError> {
let account = OlmAccount::unpickle(pickle.pickle.0, pickle_mode)?;
let identity_keys = account.parsed_identity_keys();
Ok(Self {
user_id: Arc::new(pickle.user_id),
device_id: Arc::new(pickle.device_id),
inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::from(pickle.shared)),
uploaded_signed_key_count: Arc::new(AtomicI64::new(pickle.uploaded_signed_key_count)),
})
}
pub(crate) fn unsigned_device_keys(&self) -> DeviceKeys {
let identity_keys = self.identity_keys();
let mut keys = BTreeMap::new();
keys.insert(
DeviceKeyId::from_parts(DeviceKeyAlgorithm::Curve25519, &self.device_id),
identity_keys.curve25519().to_owned(),
);
keys.insert(
DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, &self.device_id),
identity_keys.ed25519().to_owned(),
);
DeviceKeys::new(
(*self.user_id).clone(),
(*self.device_id).clone(),
Self::ALGORITHMS.iter().map(|a| (&**a).clone()).collect(),
keys,
BTreeMap::new(),
)
}
pub(crate) async fn device_keys(&self) -> DeviceKeys {
let mut device_keys = self.unsigned_device_keys();
let json_device_keys = json!({
"user_id": device_keys.user_id,
"device_id": device_keys.device_id,
"algorithms": device_keys.algorithms,
"keys": device_keys.keys,
});
device_keys
.signatures
.entry(self.user_id().clone())
.or_insert_with(BTreeMap::new)
.insert(
DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, &self.device_id),
self.sign_json(json_device_keys).await,
);
device_keys
}
pub(crate) async fn bootstrap_cross_signing(
&self,
) -> (
PrivateCrossSigningIdentity,
UploadSigningKeysRequest,
SignatureUploadRequest,
) {
PrivateCrossSigningIdentity::new_with_account(self).await
}
pub async fn sign_json(&self, json: Value) -> String {
let canonical_json: CanonicalJsonValue =
json.try_into().expect("Can't canonicalize the json value");
self.sign(&canonical_json.to_string()).await
}
pub(crate) async fn signed_one_time_keys_helper(
&self,
) -> Result<BTreeMap<DeviceKeyId, OneTimeKey>, ()> {
let one_time_keys = self.one_time_keys().await;
let mut one_time_key_map = BTreeMap::new();
for (key_id, key) in one_time_keys.curve25519().iter() {
let key_json = json!({
"key": key,
});
let signature = self.sign_json(key_json).await;
let mut signature_map = BTreeMap::new();
signature_map.insert(
DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, &self.device_id),
signature,
);
let mut signatures = BTreeMap::new();
signatures.insert((*self.user_id).clone(), signature_map);
let signed_key = SignedKey {
key: key.to_owned(),
signatures,
};
one_time_key_map.insert(
DeviceKeyId::from_parts(
DeviceKeyAlgorithm::SignedCurve25519,
key_id.as_str().into(),
),
OneTimeKey::SignedKey(signed_key),
);
}
Ok(one_time_key_map)
}
pub(crate) async fn signed_one_time_keys(
&self,
) -> Result<BTreeMap<DeviceKeyId, OneTimeKey>, ()> {
let _ = self.generate_one_time_keys().await?;
self.signed_one_time_keys_helper().await
}
pub(crate) async fn create_outbound_session_helper(
&self,
their_identity_key: &str,
their_one_time_key: &SignedKey,
) -> Result<Session, OlmSessionError> {
let session = self
.inner
.lock()
.await
.create_outbound_session(their_identity_key, &their_one_time_key.key)?;
let now = Instant::now();
let session_id = session.session_id();
Ok(Session {
user_id: self.user_id.clone(),
device_id: self.device_id.clone(),
our_identity_keys: self.identity_keys.clone(),
inner: Arc::new(Mutex::new(session)),
session_id: session_id.into(),
sender_key: their_identity_key.into(),
creation_time: Arc::new(now),
last_use_time: Arc::new(now),
})
}
pub(crate) async fn create_outbound_session(
&self,
device: ReadOnlyDevice,
key_map: &BTreeMap<DeviceKeyId, OneTimeKey>,
) -> Result<Session, SessionCreationError> {
let one_time_key = key_map.values().next().ok_or_else(|| {
SessionCreationError::OneTimeKeyMissing(
device.user_id().to_owned(),
device.device_id().into(),
)
})?;
let one_time_key = match one_time_key {
OneTimeKey::SignedKey(k) => k,
OneTimeKey::Key(_) => {
return Err(SessionCreationError::OneTimeKeyNotSigned(
device.user_id().to_owned(),
device.device_id().into(),
));
}
};
device.verify_one_time_key(&one_time_key).map_err(|e| {
SessionCreationError::InvalidSignature(
device.user_id().to_owned(),
device.device_id().into(),
e,
)
})?;
let curve_key = device
.get_key(DeviceKeyAlgorithm::Curve25519)
.ok_or_else(|| {
SessionCreationError::DeviceMissingCurveKey(
device.user_id().to_owned(),
device.device_id().into(),
)
})?;
self.create_outbound_session_helper(curve_key, &one_time_key)
.await
.map_err(|e| {
SessionCreationError::OlmError(
device.user_id().to_owned(),
device.device_id().into(),
e,
)
})
}
pub(crate) async fn create_inbound_session(
&self,
their_identity_key: &str,
message: PreKeyMessage,
) -> Result<Session, OlmSessionError> {
let session = self
.inner
.lock()
.await
.create_inbound_session_from(their_identity_key, message)?;
self.inner
.lock()
.await
.remove_one_time_keys(&session)
.expect(
"Session was successfully created but the account doesn't hold a matching one-time key",
);
let now = Instant::now();
let session_id = session.session_id();
Ok(Session {
user_id: self.user_id.clone(),
device_id: self.device_id.clone(),
our_identity_keys: self.identity_keys.clone(),
inner: Arc::new(Mutex::new(session)),
session_id: session_id.into(),
sender_key: their_identity_key.into(),
creation_time: Arc::new(now),
last_use_time: Arc::new(now),
})
}
pub(crate) async fn create_group_session_pair(
&self,
room_id: &RoomId,
settings: EncryptionSettings,
) -> Result<(OutboundGroupSession, InboundGroupSession), ()> {
if settings.algorithm != EventEncryptionAlgorithm::MegolmV1AesSha2 {
return Err(());
}
let outbound = OutboundGroupSession::new(
self.device_id.clone(),
self.identity_keys.clone(),
room_id,
settings,
);
let identity_keys = self.identity_keys();
let sender_key = identity_keys.curve25519();
let signing_key = identity_keys.ed25519();
let inbound = InboundGroupSession::new(
sender_key,
signing_key,
&room_id,
outbound.session_key().await,
)
.expect("Can't create inbound group session from a newly created outbound group session");
Ok((outbound, inbound))
}
#[cfg(test)]
pub(crate) async fn create_group_session_pair_with_defaults(
&self,
room_id: &RoomId,
) -> Result<(OutboundGroupSession, InboundGroupSession), ()> {
self.create_group_session_pair(room_id, EncryptionSettings::default())
.await
}
#[cfg(test)]
pub(crate) async fn create_session_for(&self, other: &ReadOnlyAccount) -> (Session, Session) {
other.generate_one_time_keys_helper(1).await;
let one_time = other.signed_one_time_keys().await.unwrap();
let device = ReadOnlyDevice::from_account(other).await;
let mut our_session = self
.create_outbound_session(device.clone(), &one_time)
.await
.unwrap();
other.mark_keys_as_published().await;
let message = our_session
.encrypt(&device, EventType::Dummy, json!({}))
.await
.unwrap();
let content = if let EncryptedEventContent::OlmV1Curve25519AesSha2(c) = message {
c
} else {
panic!("Invalid encrypted event algorithm");
};
let own_ciphertext = content
.ciphertext
.get(other.identity_keys.curve25519())
.unwrap();
let message_type: u8 = own_ciphertext.message_type.try_into().unwrap();
let message =
OlmMessage::from_type_and_ciphertext(message_type.into(), own_ciphertext.body.clone())
.unwrap();
let prekey = if let OlmMessage::PreKey(m) = message.clone() {
m
} else {
panic!("Wrong Olm message type");
};
let our_device = ReadOnlyDevice::from_account(self).await;
let mut other_session = other
.create_inbound_session(
our_device
.keys()
.get(&DeviceKeyId::from_parts(
DeviceKeyAlgorithm::Curve25519,
our_device.device_id(),
))
.unwrap(),
prekey,
)
.await
.unwrap();
other_session.decrypt(message).await.unwrap();
(our_session, other_session)
}
}
impl PartialEq for ReadOnlyAccount {
fn eq(&self, other: &Self) -> bool {
self.identity_keys() == other.identity_keys() && self.shared() == other.shared()
}
}