Skip to main content

max / makenotwork

10.1 KB · 343 lines History Blame Raw
1 //! SyncKit end-to-end key rotation state machine: begin a rotation, stream the
2 //! re-encrypted entries in batches, and atomically complete it. `begin` and
3 //! `complete` guard against concurrent rotations for the same key.
4
5 use serde_json::Value as JsonValue;
6 use sqlx::PgPool;
7
8 use crate::db::{SyncAppId, SyncDeviceId, UserId};
9 use crate::error::Result;
10
11 // ── Key Rotation ──
12
13 /// Begin a key rotation. Returns the rotation row if created, or the existing
14 /// row if this device already has an active rotation (resume support).
15 ///
16 /// Returns `None` if the key_version doesn't match (caller should 409).
17 /// Returns `Err` if a different device has an active rotation.
18 #[tracing::instrument(skip_all)]
19 pub async fn begin_key_rotation(
20 pool: &PgPool,
21 app_id: SyncAppId,
22 user_id: UserId,
23 device_id: SyncDeviceId,
24 new_encrypted_key: &str,
25 expected_key_version: i32,
26 ) -> Result<std::result::Result<crate::db::models::DbSyncKeyRotation, &'static str>> {
27 // The whole begin is a check-then-insert: without serialization two devices
28 // can both observe "no existing rotation" and both INSERT, racing on the
29 // unique constraint (one 500s, or two rotations exist). Run it in a
30 // transaction and take a row lock on the `sync_keys` row up front,
31 // `FOR UPDATE` serializes every concurrent begin for the same (app, user),
32 // so the existing-rotation check and the insert are atomic. (Compare
33 // `complete_key_rotation`, which already locks.)
34 let mut tx = pool.begin().await?;
35
36 // Verify key_version matches, holding the key row for the rest of the tx.
37 let key_row: Option<(i32, i32)> = sqlx::query_as(
38 "SELECT key_version, key_id FROM sync_keys WHERE app_id = $1 AND user_id = $2 FOR UPDATE",
39 )
40 .bind(app_id)
41 .bind(user_id)
42 .fetch_optional(&mut *tx)
43 .await?;
44
45 let Some((current_version, current_key_id)) = key_row else {
46 return Ok(Err("no encryption key exists"));
47 };
48
49 if current_version != expected_key_version {
50 return Ok(Err("key version mismatch"));
51 }
52
53 // Check for existing rotation (serialized by the lock above).
54 let existing = sqlx::query_as::<_, crate::db::models::DbSyncKeyRotation>(
55 "SELECT * FROM sync_key_rotations WHERE app_id = $1 AND user_id = $2",
56 )
57 .bind(app_id)
58 .bind(user_id)
59 .fetch_optional(&mut *tx)
60 .await?;
61
62 if let Some(rotation) = existing {
63 if rotation.device_id == device_id {
64 // Same device resuming, return existing rotation
65 return Ok(Ok(rotation));
66 }
67 return Ok(Err("rotation already in progress by another device"));
68 }
69
70 // Get target_seq (max seq for this user)
71 let target_seq: i64 = sqlx::query_scalar(
72 "SELECT COALESCE(MAX(seq), 0) FROM sync_log WHERE app_id = $1 AND user_id = $2",
73 )
74 .bind(app_id)
75 .bind(user_id)
76 .fetch_one(&mut *tx)
77 .await?;
78
79 let new_key_id = current_key_id + 1;
80
81 let rotation = sqlx::query_as::<_, crate::db::models::DbSyncKeyRotation>(
82 r"
83 INSERT INTO sync_key_rotations (app_id, user_id, device_id, new_encrypted_key, old_key_version, new_key_id, target_seq)
84 VALUES ($1, $2, $3, $4, $5, $6, $7)
85 RETURNING *
86 ",
87 )
88 .bind(app_id)
89 .bind(user_id)
90 .bind(device_id)
91 .bind(new_encrypted_key)
92 .bind(current_version)
93 .bind(new_key_id)
94 .bind(target_seq)
95 .fetch_one(&mut *tx)
96 .await?;
97
98 tx.commit().await?;
99 Ok(Ok(rotation))
100 }
101
102 /// Get the active rotation for a user, if any.
103 #[tracing::instrument(skip_all)]
104 pub async fn get_key_rotation(
105 pool: &PgPool,
106 app_id: SyncAppId,
107 user_id: UserId,
108 ) -> Result<Option<crate::db::models::DbSyncKeyRotation>> {
109 let rotation = sqlx::query_as::<_, crate::db::models::DbSyncKeyRotation>(
110 "SELECT * FROM sync_key_rotations WHERE app_id = $1 AND user_id = $2",
111 )
112 .bind(app_id)
113 .bind(user_id)
114 .fetch_optional(pool)
115 .await?;
116
117 Ok(rotation)
118 }
119
120 /// One sync-log entry awaiting re-encryption: (seq, table_name, row_id, data).
121 pub type RotationEntry = (i64, String, String, Option<JsonValue>);
122
123 /// Pull sync log entries that need re-encryption (key_id != new_key_id).
124 /// Returns entries ordered by seq, paginated by after_seq.
125 #[tracing::instrument(skip_all)]
126 pub async fn get_rotation_entries(
127 pool: &PgPool,
128 app_id: SyncAppId,
129 user_id: UserId,
130 new_key_id: i32,
131 after_seq: i64,
132 limit: i64,
133 ) -> Result<Vec<RotationEntry>> {
134 let entries: Vec<RotationEntry> = sqlx::query_as(
135 r"
136 SELECT seq, table_name, row_id, data FROM sync_log
137 WHERE app_id = $1 AND user_id = $2
138 AND seq > $3
139 AND (key_id IS NULL OR key_id != $5)
140 ORDER BY seq ASC
141 LIMIT $4
142 ",
143 )
144 .bind(app_id)
145 .bind(user_id)
146 .bind(after_seq)
147 .bind(limit)
148 .bind(new_key_id)
149 .fetch_all(pool)
150 .await?;
151
152 Ok(entries)
153 }
154
155 /// Batch-update re-encrypted sync log entries during key rotation.
156 /// Sets data and key_id for each (seq) in the batch.
157 /// Returns the number of rows updated.
158 #[tracing::instrument(skip_all)]
159 pub async fn submit_rotation_batch(
160 pool: &PgPool,
161 app_id: SyncAppId,
162 user_id: UserId,
163 rotation_id: uuid::Uuid,
164 new_key_id: i32,
165 entries: &[(i64, Option<JsonValue>)],
166 ) -> Result<u64> {
167 if entries.is_empty() {
168 return Ok(0);
169 }
170
171 let mut seqs: Vec<i64> = Vec::with_capacity(entries.len());
172 let mut data_values: Vec<JsonValue> = Vec::with_capacity(entries.len());
173 for (seq, data) in entries {
174 seqs.push(*seq);
175 data_values.push(data.clone().unwrap_or(JsonValue::Null));
176 }
177
178 let mut tx = pool.begin().await?;
179
180 let updated = sqlx::query(
181 r"
182 UPDATE sync_log AS sl
183 SET data = CASE WHEN batch.new_data = 'null'::jsonb THEN NULL ELSE batch.new_data END,
184 key_id = $3
185 FROM UNNEST($4::bigint[], $5::jsonb[]) AS batch(seq, new_data)
186 WHERE sl.seq = batch.seq AND sl.app_id = $1 AND sl.user_id = $2
187 ",
188 )
189 .bind(app_id)
190 .bind(user_id)
191 .bind(new_key_id)
192 .bind(&seqs)
193 .bind(&data_values)
194 .execute(&mut *tx)
195 .await?;
196
197 // Update progress marker
198 if let Some(&max_seq) = seqs.iter().max() {
199 sqlx::query(
200 r"
201 UPDATE sync_key_rotations
202 SET re_encrypted_through_seq = GREATEST(re_encrypted_through_seq, $1),
203 updated_at = NOW()
204 WHERE id = $2
205 ",
206 )
207 .bind(max_seq)
208 .bind(rotation_id)
209 .execute(&mut *tx)
210 .await?;
211 }
212
213 tx.commit().await?;
214
215 Ok(updated.rows_affected())
216 }
217
218 /// Complete a key rotation: swap the new key into sync_keys and delete the rotation.
219 /// Returns `Err("entries remain")` if un-rotated entries still exist, with the count.
220 #[tracing::instrument(skip_all)]
221 pub async fn complete_key_rotation(
222 pool: &PgPool,
223 app_id: SyncAppId,
224 user_id: UserId,
225 rotation_id: uuid::Uuid,
226 ) -> Result<std::result::Result<i32, i64>> {
227 // Lock the rotation row and run the remaining-entries check, the key swap,
228 // and the rotation delete in ONE transaction. The `FOR UPDATE` serializes
229 // this against a concurrent `rotation_batch` for the same rotation, so the
230 // count that authorizes the swap can't go stale between the read and the
231 // write (the prior code counted on the pool, outside the swap tx).
232 let mut tx = pool.begin().await?;
233
234 let rotation = sqlx::query_as::<_, crate::db::models::DbSyncKeyRotation>(
235 "SELECT * FROM sync_key_rotations WHERE id = $1 AND app_id = $2 AND user_id = $3 FOR UPDATE",
236 )
237 .bind(rotation_id)
238 .bind(app_id)
239 .bind(user_id)
240 .fetch_optional(&mut *tx)
241 .await?;
242
243 let Some(rotation) = rotation else {
244 return Ok(Err(0)); // No rotation found
245 };
246
247 // Check for remaining un-rotated entries up to the target_seq captured at
248 // rotation start. Entries arriving after rotation began are excluded, they
249 // will use the new key once rotation completes.
250 let remaining: i64 = sqlx::query_scalar(
251 r"
252 SELECT COUNT(*) FROM sync_log
253 WHERE app_id = $1 AND user_id = $2
254 AND seq <= $4
255 AND (key_id IS NULL OR key_id != $3)
256 ",
257 )
258 .bind(app_id)
259 .bind(user_id)
260 .bind(rotation.new_key_id)
261 .bind(rotation.target_seq)
262 .fetch_one(&mut *tx)
263 .await?;
264
265 if remaining > 0 {
266 return Ok(Err(remaining));
267 }
268
269 // Swap key and delete rotation (same transaction as the authorizing count).
270 sqlx::query(
271 r"
272 UPDATE sync_keys
273 SET encrypted_key = $3,
274 key_version = key_version + 1,
275 key_id = $4,
276 updated_at = NOW()
277 WHERE app_id = $1 AND user_id = $2
278 ",
279 )
280 .bind(app_id)
281 .bind(user_id)
282 .bind(&rotation.new_encrypted_key)
283 .bind(rotation.new_key_id)
284 .execute(&mut *tx)
285 .await?;
286
287 sqlx::query("DELETE FROM sync_key_rotations WHERE id = $1")
288 .bind(rotation_id)
289 .execute(&mut *tx)
290 .await?;
291
292 tx.commit().await?;
293
294 Ok(Ok(rotation.new_key_id))
295 }
296
297 /// Cancel a stale rotation (only if older than the stale threshold).
298 /// Returns true if cancelled, false if not found or not stale.
299 #[tracing::instrument(skip_all)]
300 pub async fn cancel_stale_rotation(
301 pool: &PgPool,
302 app_id: SyncAppId,
303 user_id: UserId,
304 stale_hours: i64,
305 ) -> Result<bool> {
306 let result = sqlx::query(
307 r"
308 DELETE FROM sync_key_rotations
309 WHERE app_id = $1 AND user_id = $2
310 AND updated_at < NOW() - make_interval(hours => $3::int)
311 ",
312 )
313 .bind(app_id)
314 .bind(user_id)
315 .bind(stale_hours)
316 .execute(pool)
317 .await?;
318
319 Ok(result.rows_affected() > 0)
320 }
321
322 /// Get sync status: total changes and latest seq for a user within an app.
323 #[tracing::instrument(skip_all)]
324 pub async fn get_sync_status(
325 pool: &PgPool,
326 app_id: SyncAppId,
327 user_id: UserId,
328 ) -> Result<(i64, Option<i64>)> {
329 let row: (i64, Option<i64>) = sqlx::query_as(
330 r"
331 SELECT COUNT(*), MAX(seq)
332 FROM sync_log
333 WHERE app_id = $1 AND user_id = $2
334 ",
335 )
336 .bind(app_id)
337 .bind(user_id)
338 .fetch_one(pool)
339 .await?;
340
341 Ok(row)
342 }
343