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

iris_chat_protocol/protocol_engine/
engine_persistence.rs
