Skip to main content

aven_core/sync/encrypted_tail/
server.rs

1use super::*;
2use crate::db::{Database, begin_immediate};
3use crate::sync::persistence::parent_liveness::ParentState;
4use anyhow::Context as _;
5use sqlx::SqliteConnection;
6
7async fn found(conn: &mut SqliteConnection, id: &str) -> Result<Option<Accepted>> {
8    let row: Option<(i64, Vec<u8>, Vec<u8>)> = sqlx::query_as(
9        "SELECT sequence,commitment,record FROM server_e2ee_tail WHERE operation_id=?",
10    )
11    .bind(id)
12    .fetch_optional(conn)
13    .await?;
14    row.map(|(sequence, commitment, record)| {
15        Ok(Accepted {
16            mapping: Mapping {
17                operation_id: id.into(),
18                sequence,
19                commitment: commitment
20                    .try_into()
21                    .map_err(|_| anyhow::anyhow!("error encrypted-tail-storage"))?,
22            },
23            record,
24        })
25    })
26    .transpose()
27}
28async fn mapping(conn: &mut SqliteConnection, id: &str) -> Result<Option<Mapping>> {
29    let row: Option<(i64, Vec<u8>)> =
30        sqlx::query_as("SELECT sequence,commitment FROM server_e2ee_tail WHERE operation_id=?")
31            .bind(id)
32            .fetch_optional(conn)
33            .await?;
34    row.map(|(sequence, commitment)| {
35        Ok(Mapping {
36            operation_id: id.into(),
37            sequence,
38            commitment: commitment
39                .try_into()
40                .map_err(|_| anyhow::anyhow!("error encrypted-tail-storage"))?,
41        })
42    })
43    .transpose()
44}
45fn serialized_len(value: &impl Serialize) -> Result<usize> {
46    struct Count(usize);
47    impl std::io::Write for Count {
48        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
49            self.0 += bytes.len();
50            Ok(bytes.len())
51        }
52        fn flush(&mut self) -> std::io::Result<()> {
53            Ok(())
54        }
55    }
56    let mut count = Count(0);
57    serde_json::to_writer(&mut count, value)?;
58    Ok(count.0)
59}
60async fn prefix(conn: &mut SqliteConnection, id: &str) -> Result<bool> {
61    Ok(sqlx::query_scalar(
62        "SELECT EXISTS(SELECT 1 FROM server_bootstrap_prefix WHERE operation_id=?)",
63    )
64    .bind(id)
65    .fetch_one(conn)
66    .await?)
67}
68impl Database {
69    pub async fn encrypted_tail_exchange(
70        &self,
71        context: &Context,
72        bearer: &super::super::seed_claim::Secret,
73        op: Operation,
74    ) -> Result<Reply> {
75        let mut conn = self.acquire_writer().await?;
76        let mut tx = begin_immediate(&mut conn).await?;
77        let current =
78            crate::sync::seed_claim::membership::persistence::current(self, &mut tx).await?;
79        aven_protocol::tail::authenticate(&current.membership, context, bearer)?;
80        let binding = current.membership.publication().binding();
81        let n = i64::try_from(binding.prefix_count)?;
82        let high = i64::try_from(
83            crate::sync::seed_claim::membership::persistence::allocator(
84                &mut tx,
85                &current.membership,
86            )
87            .await?,
88        )?;
89        let reply = match op {
90            Operation::Features => Reply::Features(BatchFeatures {
91                count: BATCH_COUNT,
92                bytes: BATCH_BYTES,
93            }),
94            Operation::Append { record, ticket } => {
95                let e = codec::parse(&record)?;
96                valid(e.vault == context.vault && e.stream == context.stream)?;
97                ensure!(
98                    !prefix(&mut tx, &e.id).await?,
99                    super::PrefixIdentityCollision
100                );
101                if let Some(old) = found(&mut tx, &e.id).await? {
102                    Reply::Appended(old.mapping)
103                } else {
104                    aven_protocol::tail::admit_generations(&current.membership, [e.generation])?;
105                    let sequence = high
106                        .checked_add(1)
107                        .context("error encrypted-tail-sequence-exhausted")?;
108                    let mapping = Mapping {
109                        operation_id: e.id.clone(),
110                        sequence,
111                        commitment: hash(&record),
112                    };
113                    sqlx::query("INSERT INTO server_e2ee_tail(operation_id,sequence,commitment,record) VALUES(?,?,?,?)")
114                        .bind(&e.id).bind(sequence).bind(mapping.commitment.as_slice()).bind(&record).execute(&mut *tx).await?;
115                    super::attachments::server::admit(
116                        &mut tx,
117                        context,
118                        &current.membership,
119                        &e.id,
120                        &e.projection,
121                        ticket.as_ref(),
122                    )
123                    .await?;
124                    apply_parent(&mut tx, &e.id, &e.projection).await?;
125                    sqlx::query("UPDATE server_e2ee_allocator SET high_water=? WHERE singleton=1")
126                        .bind(sequence)
127                        .execute(&mut *tx)
128                        .await?;
129                    Reply::Appended(mapping)
130                }
131            }
132            Operation::Lookup {
133                operation_id,
134                expected,
135            } => {
136                valid(!operation_id.is_empty() && operation_id.len() <= 256)?;
137                if let Some(want) = &expected {
138                    valid(want.operation_id == operation_id)?;
139                }
140                if prefix(&mut tx, &operation_id).await? {
141                    valid(expected.is_none())?;
142                    Reply::Bootstrap
143                } else if let Some(record) = found(&mut tx, &operation_id).await? {
144                    valid(expected.as_ref().is_none_or(|m| m == &record.mapping))?;
145                    Reply::Found(record)
146                } else {
147                    valid(expected.is_none())?;
148                    Reply::Absent
149                }
150            }
151            Operation::Pull {
152                after,
153                limit,
154                watermark,
155            } => {
156                let watermark = watermark.unwrap_or(high);
157                valid(
158                    after >= n
159                        && after <= watermark
160                        && watermark <= high
161                        && (1..=PAGE_COUNT).contains(&limit),
162                )?;
163                let rows:Vec<(String,i64,Vec<u8>,Vec<u8>)>=sqlx::query_as("SELECT operation_id,sequence,commitment,record FROM server_e2ee_tail WHERE sequence>? AND sequence<=? ORDER BY sequence LIMIT ?")
164                    .bind(after).bind(watermark).bind(limit as i64).fetch_all(&mut *tx).await?;
165                let mut records = Vec::new();
166                let mut size = 0;
167                let mut cursor = after;
168                for (id, sequence, commitment, record) in rows {
169                    let accepted = Accepted {
170                        mapping: Mapping {
171                            operation_id: id,
172                            sequence,
173                            commitment: commitment
174                                .try_into()
175                                .map_err(|_| anyhow::anyhow!("error encrypted-tail-storage"))?,
176                        },
177                        record,
178                    };
179                    // Measure the JSON the transport sends, plus one separator.
180                    let len = serialized_len(&accepted)? + 1;
181                    if size + len > PAGE_BYTES {
182                        break;
183                    }
184                    size += len;
185                    cursor = sequence;
186                    records.push(accepted);
187                }
188                let has_more:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM server_e2ee_tail WHERE sequence>? AND sequence<=?)").bind(cursor).bind(watermark).fetch_one(&mut *tx).await?;
189                valid(!has_more || cursor > after)?;
190                Reply::Page(Page {
191                    after,
192                    watermark,
193                    cursor,
194                    has_more,
195                    records,
196                })
197            }
198        };
199        tx.commit().await?;
200        Ok(reply)
201    }
202
203    pub async fn encrypted_tail_batch_exchange(
204        &self,
205        context: &Context,
206        bearer: &super::super::seed_claim::Secret,
207        op: BatchOperation,
208    ) -> Result<BatchReply> {
209        let mut conn = self.acquire_writer().await?;
210        let mut tx = begin_immediate(&mut conn).await?;
211        let current =
212            crate::sync::seed_claim::membership::persistence::current(self, &mut tx).await?;
213        aven_protocol::tail::authenticate(&current.membership, context, bearer)?;
214        let reply = match op {
215            BatchOperation::Append { records } => {
216                valid((1..=BATCH_COUNT).contains(&records.len()))?;
217                valid(
218                    records
219                        .iter()
220                        .try_fold(0usize, |total, record| total.checked_add(record.0.len()))
221                        .is_some_and(|total| total <= BATCH_BYTES),
222                )?;
223                let mut ids = std::collections::HashSet::new();
224                let mut parsed = Vec::with_capacity(records.len());
225                for record in &records {
226                    let envelope = codec::parse(&record.0)?;
227                    valid(
228                        envelope.vault == context.vault
229                            && envelope.stream == context.stream
230                            && !matches!(envelope.projection, domain::Projection::Ref { .. })
231                            && ids.insert(envelope.id.clone()),
232                    )?;
233                    ensure!(
234                        !prefix(&mut tx, &envelope.id).await?,
235                        super::PrefixIdentityCollision
236                    );
237                    ensure!(
238                        mapping(&mut tx, &envelope.id).await?.is_none(),
239                        "error encrypted-tail-batch-known"
240                    );
241                    parsed.push(envelope);
242                }
243                aven_protocol::tail::admit_generations(
244                    &current.membership,
245                    parsed.iter().map(|envelope| envelope.generation),
246                )?;
247                let high = i64::try_from(
248                    crate::sync::seed_claim::membership::persistence::allocator(
249                        &mut tx,
250                        &current.membership,
251                    )
252                    .await?,
253                )?;
254                let end = high
255                    .checked_add(i64::try_from(records.len())?)
256                    .context("error encrypted-tail-sequence-exhausted")?;
257                let mut mappings = Vec::with_capacity(records.len());
258                for (offset, (record, envelope)) in records.iter().zip(&parsed).enumerate() {
259                    let sequence = high
260                        .checked_add(i64::try_from(offset)? + 1)
261                        .context("error encrypted-tail-sequence-exhausted")?;
262                    let mapping = Mapping {
263                        operation_id: envelope.id.clone(),
264                        sequence,
265                        commitment: hash(&record.0),
266                    };
267                    sqlx::query("INSERT INTO server_e2ee_tail(operation_id,sequence,commitment,record) VALUES(?,?,?,?)")
268                        .bind(&envelope.id)
269                        .bind(sequence)
270                        .bind(mapping.commitment.as_slice())
271                        .bind(&record.0)
272                        .execute(&mut *tx)
273                        .await?;
274                    super::attachments::server::admit(
275                        &mut tx,
276                        context,
277                        &current.membership,
278                        &envelope.id,
279                        &envelope.projection,
280                        None,
281                    )
282                    .await?;
283                    apply_parent(&mut tx, &envelope.id, &envelope.projection).await?;
284                    mappings.push(mapping.into());
285                }
286                sqlx::query("UPDATE server_e2ee_allocator SET high_water=? WHERE singleton=1")
287                    .bind(end)
288                    .execute(&mut *tx)
289                    .await?;
290                BatchReply::Appended(mappings)
291            }
292            BatchOperation::Resolve { operation_ids } => {
293                valid((1..=BATCH_COUNT).contains(&operation_ids.len()))?;
294                let mut ids = std::collections::HashSet::new();
295                let mut resolutions = Vec::with_capacity(operation_ids.len());
296                for operation_id in operation_ids {
297                    valid(
298                        !operation_id.is_empty()
299                            && operation_id.len() <= 256
300                            && ids.insert(operation_id.clone()),
301                    )?;
302                    resolutions.push(if prefix(&mut tx, &operation_id).await? {
303                        Resolution::Bootstrap { operation_id }
304                    } else if let Some(mapping) = mapping(&mut tx, &operation_id).await? {
305                        Resolution::Found(mapping.into())
306                    } else {
307                        Resolution::Absent { operation_id }
308                    });
309                }
310                BatchReply::Resolved(resolutions)
311            }
312        };
313        tx.commit().await?;
314        Ok(reply)
315    }
316}
317async fn apply_parent(conn: &mut SqliteConnection, id: &str, p: &domain::Projection) -> Result<()> {
318    let domain::Projection::Parent {
319        action,
320        workspace,
321        task,
322        deleted,
323        version,
324    } = p
325    else {
326        return Ok(());
327    };
328    let workspace = super::moves::server_workspace(conn, task, workspace).await?;
329    let workspace = &workspace;
330    let existing:Option<(Option<String>,bool,bool)>=sqlx::query_as("SELECT version,deleted,protected FROM server_e2ee_image_parents WHERE workspace=? AND parent=?").bind(workspace).bind(task).fetch_optional(&mut *conn).await?;
331    if existing.is_none() {
332        let count: i64 = sqlx::query_scalar("SELECT count(*) FROM server_e2ee_image_parents")
333            .fetch_one(&mut *conn)
334            .await?;
335        valid(count < 262144)?;
336    }
337    let mut state = existing
338        .map(|(version, deleted, protected)| ParentState {
339            version,
340            deleted,
341            protected,
342        })
343        .unwrap_or_default();
344    state.apply(*action, id, *deleted, version.as_deref());
345    sqlx::query("INSERT INTO server_e2ee_image_parents(workspace,parent,version,deleted,protected) VALUES(?,?,?,?,?) ON CONFLICT(workspace,parent) DO UPDATE SET version=excluded.version,deleted=excluded.deleted,protected=excluded.protected")
346        .bind(workspace).bind(task).bind(state.version).bind(state.deleted).bind(state.protected).execute(&mut *conn).await?;
347    super::attachments::server::retain_parent(conn, workspace, task).await?;
348    Ok(())
349}
350
351#[cfg(test)]
352mod rotation_tests;