Skip to main content

iris_chat_protocol/protocol_engine/
engine_persistence.rs

1use super::*;
2
3impl ProtocolEngine {
4    pub(super) fn has_authoritative_local_roster(&self) -> bool {
5        if self.local_app_keys_observed {
6            return true;
7        }
8        self.session_manager
9            .roster(self.local_owner)
10            .is_some_and(|roster| {
11                let devices = roster.devices();
12                !devices.is_empty()
13                    && (devices.len() > 1 || devices[0].device_pubkey != self.local_device)
14            })
15    }
16
17    #[cfg_attr(all(test, target_os = "linux"), track_caller)]
18    pub(super) fn persist(&mut self) -> anyhow::Result<()> {
19        #[cfg(all(test, target_os = "linux"))]
20        tests::audit_retry_persist_requested(std::panic::Location::caller());
21        if self.batch_depth.get() > 0 {
22            self.batch_persist_dirty.set(true);
23            return Ok(());
24        }
25        self.persist_now()
26    }
27
28    #[cfg_attr(all(test, target_os = "linux"), track_caller)]
29    pub(super) fn persist_now(&mut self) -> anyhow::Result<()> {
30        let state = ProtocolEnginePersistedState {
31            version: PROTOCOL_ENGINE_STATE_VERSION,
32            group_invite_uses: self.group_invite_uses.clone(),
33            session_manager: self.session_manager.snapshot(),
34            group_manager: self.group_manager.snapshot(),
35            verified_app_keys_owners: self.verified_app_keys_owners.clone(),
36            app_keys_provenance_version: PROTOCOL_APP_KEYS_PROVENANCE_VERSION,
37            group_ingress_provenance_version: PROTOCOL_GROUP_INGRESS_PROVENANCE_VERSION,
38            invite_owner_app_keys_evidence: self.invite_owner_app_keys_evidence.clone(),
39            processed_private_invite_response_ids: self
40                .processed_private_invite_response_ids
41                .clone(),
42            pending_inbound: Vec::new(),
43            pending_group_fanouts: Vec::new(),
44            ready_group_publications: self.ready_group_publications.clone(),
45            pending_local_sibling_sends: self.pending_local_sibling_sends.clone(),
46            pending_remote_sends: self.pending_remote_sends.clone(),
47            pending_group_pairwise_payloads: self.pending_group_pairwise_payloads.clone(),
48            pending_group_sender_key_messages: PendingGroupMessages::default(),
49            pending_group_sender_key_repairs: self.pending_group_sender_key_repairs.clone(),
50            processed_group_sender_key_messages: self.processed_group_sender_key_messages.clone(),
51            answered_group_sender_key_repairs: self.answered_group_sender_key_repairs.clone(),
52            pending_decrypted_deliveries: self
53                .pending_decrypted_deliveries
54                .iter()
55                .filter(|delivery| !delivery.discarded)
56                .cloned()
57                .collect(),
58            group_roster_fact_histories: self.group_roster_fact_histories.clone(),
59            subscription_generation: self.subscription_generation,
60        };
61        let (checkpoint, layout) = record_protocol_checkpoint(
62            state,
63            &self.checkpoint_layout.borrow(),
64            &self.pending_group_sender_key_messages,
65            &self.pending_group_fanouts,
66            &self.pending_inbound,
67        )?;
68        self.storage
69            .put_checkpoint(PROTOCOL_ENGINE_STATE_KEY, &checkpoint)?;
70        #[cfg(all(test, target_os = "linux"))]
71        tests::audit_retry_persisted(std::panic::Location::caller(), self.batch_depth.get());
72        self.batch_persist_dirty.set(false);
73        self.pending_decrypted_deliveries
74            .retain(|delivery| !delivery.discarded);
75        *self.checkpoint_layout.borrow_mut() = layout;
76        Ok(())
77    }
78
79    pub fn enter_batch(&self) {
80        if self.batch_depth.get() == 0 {
81            self.group_sender_key_retry.borrow_mut().reset_budget();
82        }
83        self.batch_depth
84            .set(self.batch_depth.get().saturating_add(1));
85    }
86
87    pub fn exit_batch(&mut self) -> anyhow::Result<()> {
88        let depth = self.batch_depth.get();
89        if depth == 0 {
90            return Ok(());
91        }
92        self.batch_depth.set(depth - 1);
93        if self.batch_depth.get() == 0 && self.batch_persist_dirty.get() {
94            self.persist_now()?;
95        }
96        Ok(())
97    }
98}
99
100fn record_protocol_checkpoint(
101    core: ProtocolEnginePersistedState,
102    previous: &ProtocolCheckpointLayout,
103    pending: &PendingGroupMessages,
104    fanouts: &[ProtocolPendingGroupFanout],
105    inbound: &[ProtocolPendingInbound],
106) -> anyhow::Result<(crate::StorageCheckpoint, ProtocolCheckpointLayout)> {
107    let core_json = serde_json::to_string(&core)?;
108    let empty_prefix = format!(
109        "{{\"version\":{},\"pending_group_sender_key_messages\":[],\"pending_group_fanouts\":[],\"pending_inbound\":[]",
110        PROTOCOL_ENGINE_STATE_VERSION
111    );
112    // These are the first four fields of our own serializer, not untrusted
113    // substring matches. A schema/order change must update this contract.
114    anyhow::ensure!(
115        core_json.starts_with(&empty_prefix),
116        "unexpected checkpoint core prefix"
117    );
118    let core_previous = previous.core.as_deref();
119    let empty_layout = ProtocolCheckpointLayout::default();
120    let (core_json, core_layout) = layout_protocol_checkpoint_with_cached_fields(
121        core_json,
122        core_previous.unwrap_or(&empty_layout),
123        None,
124        &[],
125    );
126    let tail = core_json
127        .strip_prefix(&empty_prefix)
128        .ok_or_else(|| anyhow::anyhow!("unexpected checkpoint layout prefix"))?;
129    let mut value = crate::StorageCheckpoint::new();
130    value.push_inline(
131        "prefix",
132        format!(
133            "{{\"version\":{},\"pending_group_sender_key_messages\":[",
134            PROTOCOL_ENGINE_STATE_VERSION
135        ),
136        0,
137    );
138    for (index, body) in pending.record_bodies()?.into_iter().enumerate() {
139        value.push_record(
140            0,
141            body,
142            if index == 0 { "" } else { "," }.into(),
143            String::new(),
144            0,
145        );
146    }
147    value.push_inline("between", "],\"pending_group_fanouts\":[".into(), 0);
148    #[derive(Serialize)]
149    struct FanoutTail<'a> {
150        remaining_devices: &'a Option<BTreeSet<NdrDevicePubkey>>,
151        inner_event_id: &'a Option<String>,
152        created_at_secs: u64,
153        next_retry_at_secs: u64,
154    }
155    for (index, pending) in fanouts.iter().enumerate() {
156        let suffix = serde_json::to_string(&FanoutTail {
157            remaining_devices: &pending.remaining_devices,
158            inner_event_id: &pending.inner_event_id,
159            created_at_secs: pending.created_at_secs,
160            next_retry_at_secs: pending.next_retry_at_secs,
161        })?;
162        let prefix = format!(
163            "{}{{\"group_id\":{},\"fanout\":",
164            if index == 0 { "" } else { "," },
165            serde_json::to_string(&pending.group_id)?
166        );
167        value.push_record(
168            1,
169            pending.fanout.checkpoint_body()?,
170            prefix,
171            format!(",{}", &suffix[1..]),
172            0,
173        );
174    }
175    value.push_inline("before_inbound", "],\"pending_inbound\":[".into(), 0);
176    append_checkpoint_inbound(&mut value, inbound)?;
177    value.push_inline("core", format!("]{tail}"), 0);
178    let compact_bytes = value.len();
179    if compact_bytes < CHECKPOINT_RESERVE {
180        return Ok((value, ProtocolCheckpointLayout::default()));
181    }
182    // Padding belongs to the record, not its old ordinal. Append/removal
183    // cannot dirty the descriptors of every surviving backlog record.
184    value.pad_records(CHECKPOINT_PAGE, CHECKPOINT_PAGE / 2);
185    let reuse = previous.bytes <= compact_bytes.saturating_mul(3) + CHECKPOINT_RESERVE;
186    let target = if reuse && value.len() <= previous.bytes {
187        previous.bytes
188    } else {
189        value.len().div_ceil(CHECKPOINT_RESERVE) * CHECKPOINT_RESERVE
190    };
191    value.push_inline("reserve", String::new(), target - value.len());
192    let layout = ProtocolCheckpointLayout {
193        bytes: value.len(),
194        record_ends: Vec::new(),
195        core: Some(Box::new(core_layout)),
196    };
197    Ok((value, layout))
198}