diff --git a/rust/bridge/shared/src/net/keytrans.rs b/rust/bridge/shared/src/net/keytrans.rs index 45be05e43..02d8decde 100644 --- a/rust/bridge/shared/src/net/keytrans.rs +++ b/rust/bridge/shared/src/net/keytrans.rs @@ -16,8 +16,8 @@ use libsignal_keytrans::{ }; use libsignal_net_chat::api::RequestError; use libsignal_net_chat::api::keytrans::{ - CheckMode, Error, KeyTransparencyClient, MaybePartial, SearchKey, UnauthenticatedChatApi as _, - UsernameHash, check, + CheckMode, Error, KeyTransparencyClient, MaybePartial, SearchKey, TreeHeadWithTimestamp, + UnauthenticatedChatApi as _, UsernameHash, check, }; use libsignal_protocol::PublicKey; use prost::{DecodeError, Message}; @@ -63,8 +63,6 @@ async fn KeyTransparency_Check( let account_data = account_data.map(try_decode_account_data).transpose()?; - let last_distinguished_tree_head = try_decode_distinguished(last_distinguished_tree_head)?; - let mode = if is_self_check { CheckMode::SelfCheck { is_e164_discoverable, @@ -75,10 +73,13 @@ async fn KeyTransparency_Check( let e164_pair = make_e164_pair(e164, unidentified_access_key)?; - let maybe_partial_result = chat_connection + let last_distinguished_tree_head = try_decode_distinguished(last_distinguished_tree_head); + + let (maybe_partial_result, _updated_distinguished) = chat_connection .as_typed(|chat| { Box::pin(async move { let kt = KeyTransparencyClient::new(*chat, config); + check( &kt, &aci, @@ -86,7 +87,7 @@ async fn KeyTransparency_Check( e164_pair, username_hash, account_data, - &last_distinguished_tree_head, + last_distinguished_tree_head, mode, ) .await @@ -94,12 +95,15 @@ async fn KeyTransparency_Check( }) .await?; - maybe_partial_result.into_serialized_account_data( + let now = SystemTime::now(); + let serialized_account_data = maybe_partial_result.into_serialized_account_data( aci.as_search_key(), e164.map(|x| x.as_search_key()), maybe_hash_search_key, - SystemTime::now(), - ) + now, + )?; + // let serialized_distinguished = updated_distinguished.into_stored(now).encode_to_vec(); + Ok(serialized_account_data) } #[bridge_io(TokioAsyncContext)] @@ -169,11 +173,10 @@ fn try_decode_account_data(bytes: Box<[u8]>) -> Result) -> Result> { +fn try_decode_distinguished(bytes: Box<[u8]>) -> Option { try_decode(bytes) - .map(|stored: StoredTreeHead| stored.into_last_tree_head()) - .map_err(|_| invalid_request("could not decode last distinguished tree head"))? - .ok_or(invalid_request("last distinguished tree is required")) + .ok() + .and_then(TreeHeadWithTimestamp::from_stored) } trait MaybePartialExt { diff --git a/rust/net/chat/src/api/keytrans.rs b/rust/net/chat/src/api/keytrans.rs index cc1075bb5..ec90133f9 100644 --- a/rust/net/chat/src/api/keytrans.rs +++ b/rust/net/chat/src/api/keytrans.rs @@ -24,7 +24,7 @@ use libsignal_keytrans::{ use libsignal_net::env::KeyTransConfig; use libsignal_protocol::PublicKey; pub use maybe_partial::{AccountDataField, MaybePartial}; -pub use monitor_and_search::check; +pub use monitor_and_search::{TreeHeadWithTimestamp, check}; use verify_ext::KeyTransparencyVerifyExt as _; use super::RequestError; @@ -654,24 +654,11 @@ pub(crate) mod test_support { pub username_hash_bytes: Option>, } - #[derive(Default)] - pub struct SearchStub { - pub result: Option, RequestError>>, - pub invocations: Vec, - } - - impl SearchStub { - fn new(res: Option, RequestError>>) -> Self { - Self { - result: res, - invocations: vec![], - } - } - } - pub struct TestKt { pub monitor: Cell>>>, - pub search: Cell, + pub search: Cell, RequestError>>>, + pub distinguished: Cell>>>, + search_invocation: Cell>, } impl TestKt { @@ -683,13 +670,21 @@ pub(crate) mod test_support { Self::new(None, Some(search)) } + #[must_use] + pub fn with_distinguished(self, result: Result>) -> Self { + self.distinguished.set(Some(result)); + self + } + pub fn new( monitor: Option>>, search: Option, RequestError>>, ) -> Self { Self { monitor: Cell::new(monitor), - search: Cell::new(SearchStub::new(search)), + search: Cell::new(search), + distinguished: Cell::new(Some(Ok(test_distinguished_tree()))), + search_invocation: Cell::new(None), } } @@ -703,8 +698,8 @@ pub(crate) mod test_support { assert_matches!(result, Err(RequestError::Unexpected { log_safe }) if log_safe == "test error") } - pub fn take_searches(&self) -> Vec { - self.search.take().invocations + pub fn search_invocation(self) -> Option { + self.search_invocation.into_inner() } } @@ -719,31 +714,30 @@ pub(crate) mod test_support { _distinguished_tree_head: &LastTreeHead, ) -> impl Future, RequestError>> + Send { - let mut search_stub = self.search.take(); - - search_stub.invocations.push(OwnedParameters { + self.search_invocation.set(Some(OwnedParameters { aci: *aci, e164, username_hash_bytes: username_hash.map(|x| x.as_ref().to_vec()), - }); + })); - let result = search_stub - .result - .as_ref() - .expect("unexpected call to search") - .clone(); - self.search.set(search_stub); - std::future::ready(result.clone()) + let result = self.search.take().expect("unexpected call to search"); + std::future::ready(result) } fn distinguished( &self, _: Option, ) -> impl Future>> { - // not used in the tests - unreachable!(); - #[allow(unreachable_code)] // without this, `impl Future` gets confused - std::future::pending() + let tree_head = self + .distinguished + .take() + .expect("unexpected call to distinguished"); + let state_update = tree_head.map(|tree_head| SearchStateUpdate { + tree_head: tree_head.0, + tree_root: tree_head.1, + monitoring_data: None, + }); + std::future::ready(state_update) } fn monitor( diff --git a/rust/net/chat/src/api/keytrans/monitor_and_search.rs b/rust/net/chat/src/api/keytrans/monitor_and_search.rs index 2a29cbbef..a5007e8da 100644 --- a/rust/net/chat/src/api/keytrans/monitor_and_search.rs +++ b/rust/net/chat/src/api/keytrans/monitor_and_search.rs @@ -3,9 +3,14 @@ // SPDX-License-Identifier: AGPL-3.0-only // +use std::ops::ControlFlow; +use std::time::{Duration, SystemTime}; + use libsignal_core::curve::PublicKey; use libsignal_core::{Aci, E164}; -use libsignal_keytrans::{AccountData, LastTreeHead, MonitoringData}; +use libsignal_keytrans::{ + AccountData, LastTreeHead, LocalStateUpdate, MonitoringData, StoredTreeHead, +}; use crate::api::RequestError; use crate::api::keytrans::maybe_partial::MaybePartial; @@ -13,6 +18,9 @@ use crate::api::keytrans::{ AccountDataField, CheckMode, Error, SearchKey, UnauthenticatedChatApi, UsernameHash, }; +const MAX_DISTINGUISHED_TREE_AGE: Duration = + Duration::from_secs(7 * 24 * 60 * 60 /* one week */); + /// The main entry point to the module. pub async fn check( kt: &impl UnauthenticatedChatApi, @@ -21,9 +29,11 @@ pub async fn check( e164: Option<(E164, Vec)>, username_hash: Option>, stored_account_data: Option, - distinguished_tree_head: &LastTreeHead, + distinguished_tree_head: Option, mode: CheckMode, -) -> Result, RequestError> { +) -> Result<(MaybePartial, LastTreeHead), RequestError> { + let distinguished_tree_head = + update_distinguished_if_needed(kt, distinguished_tree_head).await?; let action = Action::plan( aci, e164.as_ref(), @@ -31,9 +41,65 @@ pub async fn check( stored_account_data, ); - action - .execute(kt, aci_identity_key, distinguished_tree_head, mode) - .await + let result = action + .execute(kt, aci_identity_key, &distinguished_tree_head, mode) + .await?; + + Ok((result, distinguished_tree_head)) +} + +fn is_too_old(stored_at_ms: u64) -> bool { + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("valid SystemTime") + .saturating_sub(Duration::from_millis(stored_at_ms)) + > MAX_DISTINGUISHED_TREE_AGE +} + +fn select_baseline_tree_head( + tree_head: Option, +) -> ControlFlow> { + match tree_head { + None => ControlFlow::Continue(None), + // Not recent enough to be used for search/monitor but a fine + // baseline for refresh. + Some(timed_head) if is_too_old(timed_head.stored_at_ms) => { + ControlFlow::Continue(Some(timed_head.tree_head)) + } + Some(timed_head) => ControlFlow::Break(timed_head.tree_head), + } +} + +async fn update_distinguished_if_needed( + kt: &impl UnauthenticatedChatApi, + tree_head: Option, +) -> Result> { + let tree_head = match select_baseline_tree_head(tree_head) { + ControlFlow::Break(tree_head) => return Ok(tree_head), + ControlFlow::Continue(baseline) => baseline, + }; + let LocalStateUpdate { + tree_head, + tree_root, + monitoring_data: _, + } = kt.distinguished(tree_head).await?; + + Ok(LastTreeHead(tree_head, tree_root)) +} + +#[cfg_attr(test, derive(Clone, PartialEq))] +pub struct TreeHeadWithTimestamp { + pub tree_head: LastTreeHead, + pub stored_at_ms: u64, +} + +impl TreeHeadWithTimestamp { + pub fn from_stored(stored: StoredTreeHead) -> Option { + Some(Self { + stored_at_ms: stored.stored_at_ms, + tree_head: stored.into_last_tree_head()?, + }) + } } /// A more ergonomic search for the clients. @@ -600,15 +666,20 @@ async fn monitor_then_search<'a>( #[cfg(test)] mod test { + use std::ops::ControlFlow; + use std::time::SystemTime; + use assert_matches::assert_matches; + use futures_util::FutureExt; use libsignal_core::E164; use libsignal_keytrans::AccountData; use nonzero_ext::nonzero; use test_case::{test_case, test_matrix}; use super::{ - Action, Parameters, PostMonitorAction, VersionChanged, check, merge_account_data, - modal_search, monitor_then_search, + Action, Parameters, PostMonitorAction, TreeHeadWithTimestamp, VersionChanged, check, + is_too_old, merge_account_data, modal_search, monitor_then_search, + select_baseline_tree_head, }; use crate::api::RequestError; use crate::api::keytrans::test_support::{ @@ -629,6 +700,19 @@ mod test { } } + fn recent_distinguished_tree() -> TreeHeadWithTimestamp { + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("valid time") + .as_millis() + .try_into() + .expect("fits in u64"); + TreeHeadWithTimestamp { + tree_head: test_distinguished_tree(), + stored_at_ms: now, + } + } + #[test] fn parameters_merge_with_prefers_other_optionals_and_self_aci() { let aci = test_account::aci(); @@ -894,7 +978,7 @@ mod test { None, None, Some(test_account_data()), - &test_distinguished_tree(), + Some(recent_distinguished_tree()), CheckMode::ContactCheck, ) .await; @@ -923,7 +1007,7 @@ mod test { None, Some(test_account::username_hash()), Some(stored_account_data), - &test_distinguished_tree(), + Some(recent_distinguished_tree()), CheckMode::ContactCheck, ) .await; @@ -955,7 +1039,7 @@ mod test { None, None, Some(test_account_data()), - &test_distinguished_tree(), + Some(recent_distinguished_tree()), CheckMode::ContactCheck, ) .await; @@ -975,12 +1059,12 @@ mod test { None, None, Some(test_account_data()), - &test_distinguished_tree(), + Some(recent_distinguished_tree()), CheckMode::ContactCheck, ) .await .expect("monitor should succeed"); - assert_eq!(actual, monitor_result.into()); + assert_eq!(actual.0, monitor_result.into()); } #[test] @@ -1366,10 +1450,10 @@ mod test { assert_eq!(result.inner, test_account_data()); - let invocation = &kt.take_searches()[0]; + let invocation = kt.search_invocation().expect("search is invoked"); // In case of version change we perform a single search using monitor parameters - assert_eq!(*invocation, monitor_parameters); + assert_eq!(invocation, monitor_parameters); } #[tokio::test] @@ -1460,9 +1544,9 @@ mod test { .await .expect("succeeds"); - let invocation = &kt.take_searches()[0]; + let invocation = kt.search_invocation().expect("search is invoked"); - assert_eq!(*invocation, search_parameters); + assert_eq!(invocation, search_parameters); // Importantly, despite having done a new search for ACI as well // we should have retained the ACI result from an earlier monitor. @@ -1513,9 +1597,9 @@ mod test { .await .expect("succeeds"); - let invocation = &kt.take_searches()[0]; + let invocation = kt.search_invocation().expect("search is invoked"); - assert_eq!(*invocation, search_parameters); + assert_eq!(invocation, search_parameters); assert!(result.inner.e164.is_some()); assert!(result.missing_fields.is_empty()); @@ -1623,9 +1707,9 @@ mod test { .await .expect("succeeds"); - let invocation = &kt.take_searches()[0]; + let invocation = kt.search_invocation().expect("search is invoked"); - assert_eq!(*invocation, search_parameters); + assert_eq!(invocation, search_parameters); assert!(result.missing_fields.is_empty()); assert_eq!(&result.inner.e164, &monitor_result.e164); @@ -1710,10 +1794,10 @@ mod test { .await .expect("succeeds"); - let invocation = &kt.take_searches()[0]; + let invocation = kt.search_invocation().expect("search is invoked"); assert_eq!( - *invocation, + invocation, Parameters { aci: &aci, e164: Some(&e164), @@ -1735,4 +1819,133 @@ mod test { } } } + + #[test] + fn check_updates_distinguished_when_its_unknown() { + let mut expected = test_distinguished_tree(); + expected.1 = [42; 32]; + let kt = TestKt::for_search(Ok(test_account_data().into())) + .with_distinguished(Ok(expected.clone())); + + let result = check( + &kt, + &test_account::aci(), + &test_account::aci_identity_key(), + None, + None, + None, + None, + CheckMode::ContactCheck, + ) + .now_or_never() + .expect("sync") + .expect("check succeeds"); + + // Distinguished stub has been consumed + assert!(kt.distinguished.take().is_none()); + // Should return the most recent distinguished tree + assert_eq!(expected, result.1); + } + + #[test] + fn check_does_not_update_distinguished_when_its_recent() { + let stored_distinguished = TreeHeadWithTimestamp { + tree_head: test_distinguished_tree(), + stored_at_ms: now_ms() - 60 * 1000, + }; + + let mut updated_distinguished = test_distinguished_tree(); + updated_distinguished.1 = [42; 32]; + let kt = TestKt::for_search(Ok(test_account_data().into())) + .with_distinguished(Ok(updated_distinguished)); + + let result = check( + &kt, + &test_account::aci(), + &test_account::aci_identity_key(), + None, + None, + None, + Some(stored_distinguished.clone()), + CheckMode::ContactCheck, + ) + .now_or_never() + .expect("sync") + .expect("check succeeds"); + + // Distinguished stub has not been consumed + assert!(kt.distinguished.take().is_some()); + // Should return the most recent distinguished tree + assert_eq!(stored_distinguished.tree_head, result.1); + } + + #[test] + fn check_propagates_distinguished_update_failure() { + let kt = TestKt::for_search(Ok(test_account_data().into())) + .with_distinguished(Err(TestKt::expected_error())); + + let result = check( + &kt, + &test_account::aci(), + &test_account::aci_identity_key(), + None, + None, + None, + None, + CheckMode::ContactCheck, + ) + .now_or_never() + .expect("sync"); + + TestKt::assert_expected_error(result); + } + + fn now_ms() -> u64 { + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("valid time") + .as_millis() + .try_into() + .expect("fits in u64") + } + + #[test_case(0 => false; "now")] + #[test_case(6 * 24 * 60 * 60 * 1000 => false; "not quite a week ago")] + #[test_case(2 * 7 * 24 * 60 * 60 * 1000 => true; "two weeks ago")] + fn is_too_old_for_values_in_the_past(ms_in_the_past: u64) -> bool { + is_too_old(now_ms() - ms_in_the_past) + } + + #[test] + fn is_too_old_value_in_the_future() { + assert!(!is_too_old(now_ms() + 1)); + } + + #[test] + fn select_baseline_tree_head_no_stored_head_continues_with_none() { + assert_matches!(select_baseline_tree_head(None), ControlFlow::Continue(None)); + } + + #[test] + fn select_baseline_tree_head_too_old_continues_with_tree_head() { + let tree_head = test_distinguished_tree(); + let stored = TreeHeadWithTimestamp { + tree_head: tree_head.clone(), + stored_at_ms: 0, + }; + assert_eq!( + select_baseline_tree_head(Some(stored)), + ControlFlow::Continue(Some(tree_head)), + ); + } + + #[test] + fn select_baseline_tree_head_recent_breaks_with_tree_head() { + let recent = recent_distinguished_tree(); + let expected = recent.tree_head.clone(); + assert_eq!( + select_baseline_tree_head(Some(recent)), + ControlFlow::Break(expected), + ); + } }