Skip to main content

iris_chat_protocol/protocol_engine/
engine_persistence.rs

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