1use std::{
2 collections::BTreeMap,
3 iter::from_fn,
4 path::{Path, PathBuf},
5 pin::pin,
6 sync::Arc,
7 time::Duration,
8};
9
10use bytes::Bytes;
11use futures::{StreamExt, TryStreamExt};
12use object_store::Error as ObjectStoreError;
13use ruma::{
14 CanonicalJsonObject, CanonicalJsonValue, EventId, Mxc, OwnedRoomId, OwnedUserId, RoomId,
15 ServerName, UserId,
16};
17use serde::{Deserialize, de::IgnoredAny};
18use tokio::time::sleep;
19use tuwunel_core::{
20 Err, Error, Result, debug_warn, err, error, info,
21 itertools::Itertools,
22 utils,
23 utils::{ReadyExt, content_disposition::make_content_disposition, stream::TryIgnore},
24 warn,
25};
26use tuwunel_database::{Map, SEP};
27
28use crate::{Services, storage::Provider};
29
30#[derive(Deserialize)]
32struct HasRoomId {
33 room_id: Option<IgnoredAny>,
34}
35
36enum MediaSource {
40 Filesystem(PathBuf),
41 Provider(Arc<Provider>),
42}
43
44struct ConduitMediaEntry<'a> {
46 server_name: &'a ServerName,
47 media_id: &'a str,
48 sha256: &'a [u8],
49 filename: Option<&'a str>,
50 content_type: Option<&'a str>,
51}
52
53fn media_source(services: &Services) -> Result<MediaSource> {
57 let config = &services.server.config;
58
59 match config.conduit_source_media_provider.as_deref() {
60 | Some(name) => services
61 .storage
62 .provider(name)
63 .map(|provider| MediaSource::Provider(provider.clone())),
64 | None => {
65 let media_dir = config
66 .conduit_source_media_path
67 .clone()
68 .unwrap_or_else(|| config.database_path.join("media"));
69
70 Ok(MediaSource::Filesystem(media_dir))
71 },
72 }
73}
74
75const PROVIDER_READ_ATTEMPTS: u32 = 2;
80
81const PROVIDER_READ_RETRY_DELAY: Duration = Duration::from_secs(2);
83
84pub(super) async fn migrate_conduit_media(services: &Services) -> Result {
93 let db = &services.db;
94 let config = &services.server.config;
95
96 let Some(metadata) = db.open_cf("servernamemediaid_metadata")? else {
97 warn!("Conduit database has no media metadata; nothing to import.");
98 return Ok(());
99 };
100
101 let owners = db.open_cf("servernamemediaid_userlocalpart")?;
102 let owners = owners.as_ref();
103
104 let blocklist = db.open_cf("blocked_servername_mediaid")?;
105 let blocklist = blocklist.as_ref();
106
107 let depth = config.conduit_media_directory_depth;
108 let length = config.conduit_media_directory_length;
109 let source = media_source(services)?;
110
111 warn!("Importing Conduit media originals into tuwunel's key-addressed store...");
112
113 let cork = db.cork_and_sync();
114 let (imported, skipped, blocked) = metadata
115 .raw_stream()
116 .ignore_err()
117 .map(Ok::<_, Error>)
118 .try_fold(
119 (0_usize, 0_usize, 0_usize),
120 async |(imported, skipped, blocked), (key, value)| {
121 if conduit_media_blocked(blocklist, key).await? {
122 return Ok((imported, skipped, blocked.saturating_add(1)));
123 }
124
125 let imported_entry =
126 import_conduit_original(services, owners, &source, depth, length, key, value)
127 .await?;
128
129 Ok(if imported_entry {
130 (imported.saturating_add(1), skipped, blocked)
131 } else {
132 (imported, skipped.saturating_add(1), blocked)
133 })
134 },
135 )
136 .await?;
137
138 drop(cork);
139
140 if blocked > 0 {
141 warn!(%blocked, "Skipped Conduit media blocked by a moderator; not imported");
142 }
143
144 if skipped > 0 {
145 warn!(%imported, %skipped, "Imported Conduit media originals; some files were skipped");
146 } else {
147 info!(%imported, "Imported Conduit media originals");
148 }
149
150 Ok(())
151}
152
153async fn import_conduit_original(
159 services: &Services,
160 owners: Option<&Arc<Map>>,
161 source: &MediaSource,
162 depth: u8,
163 length: u8,
164 key: &[u8],
165 value: &[u8],
166) -> Result<bool> {
167 let entry = match parse_conduit_media_entry(key, value) {
168 | Ok(entry) => entry,
169 | Err(e) => {
170 debug_warn!(error = %e, "skipping unimportable Conduit media entry");
171 return Ok(false);
172 },
173 };
174
175 let Some(file) = read_conduit_original(source, depth, length, entry.sha256).await? else {
176 return Ok(false);
177 };
178
179 let content_disposition = make_content_disposition(None, entry.content_type, entry.filename);
180 let owner = conduit_media_owner(owners, key, entry.server_name).await;
181 let mxc = Mxc {
182 server_name: entry.server_name,
183 media_id: entry.media_id,
184 };
185
186 services
187 .media
188 .create(&mxc, owner.as_deref(), Some(&content_disposition), entry.content_type, &file)
189 .await?;
190
191 Ok(true)
192}
193
194fn parse_conduit_media_entry<'a>(
198 key: &'a [u8],
199 value: &'a [u8],
200) -> Result<ConduitMediaEntry<'a>> {
201 let Some(sep) = key.iter().position(|&byte| byte == SEP) else {
202 return Err!(Database("Conduit media key has no server-name separator"));
203 };
204 let server_name = <&ServerName>::try_from(str::from_utf8(&key[..sep])?)
205 .map_err(|_| err!(Database("Conduit media key has an invalid server name")))?;
206
207 let media_id = str::from_utf8(&key[sep.saturating_add(1)..])?;
208
209 let (sha256, filename, content_type) = parse_conduit_media_value(value)?;
210
211 Ok(ConduitMediaEntry {
212 server_name,
213 media_id,
214 sha256,
215 filename,
216 content_type,
217 })
218}
219
220async fn read_conduit_original(
226 source: &MediaSource,
227 depth: u8,
228 length: u8,
229 sha256: &[u8],
230) -> Result<Option<Bytes>> {
231 let sha256_hex = sha256_hex(sha256);
232 match source {
233 | MediaSource::Filesystem(media_dir) => {
234 let path = conduit_media_path(media_dir, depth, length, &sha256_hex);
235 match tokio::fs::read(&path).await {
236 | Ok(file) => Ok(Some(file.into())),
237 | Err(e) => {
238 debug_warn!(?path, error = %e, "skipping unreadable Conduit media file");
239 Ok(None)
240 },
241 }
242 },
243 | MediaSource::Provider(provider) =>
244 read_provider_original(provider, &conduit_media_key(depth, length, &sha256_hex)).await,
245 }
246}
247
248async fn read_provider_original(provider: &Arc<Provider>, key: &str) -> Result<Option<Bytes>> {
253 let mut attempt = 0_u32;
254 loop {
255 attempt = attempt.saturating_add(1);
256 match provider.get(key).await {
257 | Ok(file) => return Ok(Some(file)),
258 | Err(e) if is_missing_object(&e) => {
259 debug_warn!(%key, error = %e, "skipping missing Conduit media object");
260 return Ok(None);
261 },
262 | Err(e) if attempt >= PROVIDER_READ_ATTEMPTS => {
263 error!(
264 %key,
265 attempts = PROVIDER_READ_ATTEMPTS,
266 error = %e,
267 "Aborting the Conduit media import: source storage provider unreachable. No \
268 media has been imported in a way that needs cleanup; once the provider is \
269 reachable, restart tuwunel to resume the import from the beginning."
270 );
271 return Err(e);
272 },
273 | Err(e) => {
274 warn!(%key, attempt, error = %e, "Reading Conduit media object failed; retrying");
275 sleep(PROVIDER_READ_RETRY_DELAY).await;
276 },
277 }
278 }
279}
280
281fn is_missing_object(error: &Error) -> bool {
286 matches!(error, Error::ObjectStore(ObjectStoreError::NotFound { .. }))
287}
288
289fn parse_conduit_media_value(value: &[u8]) -> Result<(&[u8], Option<&str>, Option<&str>)> {
294 let (sha256, rest) = value
295 .split_at_checked(32)
296 .ok_or_else(|| err!(Database("Conduit media value shorter than a SHA-256 digest")))?;
297
298 let mut parts = rest.split(|&byte| byte == SEP);
300 let filename = parts.next().unwrap_or_default();
301 let Some(content_type) = parts.next() else {
302 return Err!(Database("Conduit media value has no content-type separator"));
303 };
304 let filename = str::from_utf8(filename)?;
305 let content_type = str::from_utf8(content_type)?;
306 let filename = (!filename.is_empty()).then_some(filename);
307 let content_type = (!content_type.is_empty()).then_some(content_type);
308
309 Ok((sha256, filename, content_type))
310}
311
312async fn conduit_media_owner(
316 owners: Option<&Arc<Map>>,
317 key: &[u8],
318 server_name: &ServerName,
319) -> Option<OwnedUserId> {
320 let localpart = owners?.get(key).await.ok()?;
321
322 UserId::parse_with_server_name(str::from_utf8(&localpart).ok()?, server_name).ok()
323}
324
325async fn conduit_media_blocked(blocklist: Option<&Arc<Map>>, key: &[u8]) -> Result<bool> {
333 let Some(blocklist) = blocklist else {
334 return Ok(false);
335 };
336
337 match blocklist.exists(key).await {
338 | Ok(()) => Ok(true),
339 | Err(e) if e.is_not_found() => Ok(false),
340 | Err(e) => Err(e),
341 }
342}
343
344fn conduit_media_path(media_dir: &Path, depth: u8, length: u8, sha256_hex: &str) -> PathBuf {
348 let mut path = media_dir.to_path_buf();
349 path.extend(conduit_shards(depth, length, sha256_hex));
350 path
351}
352
353fn conduit_media_key(depth: u8, length: u8, sha256_hex: &str) -> String {
357 conduit_shards(depth, length, sha256_hex).join("/")
358}
359
360fn conduit_shards(depth: u8, length: u8, sha256_hex: &str) -> impl Iterator<Item = &str> {
365 let mut rest = Some(sha256_hex);
366 let mut remaining = depth;
367 from_fn(move || {
368 let current = rest?;
369 if remaining == 0 {
370 rest = None;
371 return Some(current);
372 }
373
374 remaining = remaining.saturating_sub(1);
375 match current.split_at_checked(length.into()) {
376 | Some((segment, next)) => {
377 rest = Some(next);
378 Some(segment)
379 },
380 | None => {
381 rest = None;
382 Some(current)
383 },
384 }
385 })
386}
387
388fn sha256_hex(digest: &[u8]) -> String {
390 const HEX: &[u8; 16] = b"0123456789abcdef";
391
392 let mut out = String::with_capacity(digest.len().saturating_mul(2));
393 for &byte in digest {
394 out.push(char::from(HEX[usize::from(byte >> 4)]));
395 out.push(char::from(HEX[usize::from(byte & 0x0F)]));
396 }
397
398 out
399}
400
401pub(super) async fn migrate_conduit_pdus(services: &Services) -> Result {
409 let db = &services.db;
410
411 let rooms: BTreeMap<u64, OwnedRoomId> = db["roomid_shortroomid"]
414 .stream()
415 .ignore_err()
416 .map(|(room_id, short): (&RoomId, u64)| (short, room_id.to_owned()))
417 .collect()
418 .await;
419
420 warn!("Ensuring stored PDUs carry their room_id field...");
421 let cork = db.cork_and_sync();
422
423 let pduid_pdu = &db["pduid_pdu"];
424 let timeline = pduid_pdu
425 .raw_stream()
426 .ignore_err()
427 .ready_fold((0_usize, 0_usize), |acc, (key, value)| {
428 tally(acc, inject_room_id(pduid_pdu, key, value, |_| pduid_room(&rooms, key)))
429 })
430 .await;
431
432 let outlier = &db["eventid_outlierpdu"];
433 let outliers = outlier
434 .raw_stream()
435 .ignore_err()
436 .ready_fold((0_usize, 0_usize), |acc, (key, value)| {
437 tally(acc, inject_room_id(outlier, key, value, |pdu| outlier_room(key, pdu)))
438 })
439 .await;
440
441 drop(cork);
442
443 let fixed = timeline.0.saturating_add(outliers.0);
444 let skipped = timeline.1.saturating_add(outliers.1);
445 if skipped > 0 {
446 warn!(%fixed, %skipped, "Injected room_id into stored PDUs; some were skipped");
447 } else {
448 info!(%fixed, "Ensured stored PDUs carry room_id");
449 }
450
451 Ok(())
452}
453
454fn tally((fixed, skipped): (usize, usize), result: Result<bool>) -> (usize, usize) {
455 match result {
456 | Ok(true) => (fixed.saturating_add(1), skipped),
457 | Ok(false) => (fixed, skipped),
458 | Err(e) => {
459 debug_warn!(error = %e, "skipping unreconcilable Conduit PDU");
460 (fixed, skipped.saturating_add(1))
461 },
462 }
463}
464
465fn inject_room_id(
470 map: &Arc<Map>,
471 key: &[u8],
472 value: &[u8],
473 resolve: impl FnOnce(&CanonicalJsonObject) -> Result<OwnedRoomId>,
474) -> Result<bool> {
475 let probe: HasRoomId = serde_json::from_slice(value)
476 .map_err(|e| err!(Database("Conduit PDU is not canonical JSON: {e}")))?;
477
478 if probe.room_id.is_some() {
479 return Ok(false);
480 }
481
482 let mut pdu: CanonicalJsonObject = serde_json::from_slice(value)
483 .map_err(|e| err!(Database("Conduit PDU is not canonical JSON: {e}")))?;
484
485 let room_id = resolve(&pdu)?;
486 pdu.insert("room_id".into(), CanonicalJsonValue::String(room_id.as_str().into()));
487
488 let bytes = serde_json::to_vec(&pdu)
489 .map_err(|e| err!(Database("re-serializing reconciled Conduit PDU: {e}")))?;
490
491 map.insert(key, bytes);
492
493 Ok(true)
494}
495
496fn pduid_room(rooms: &BTreeMap<u64, OwnedRoomId>, key: &[u8]) -> Result<OwnedRoomId> {
498 let short = key
499 .get(..8)
500 .ok_or_else(|| err!(Database("Conduit pduid is shorter than a short room id")))?;
501
502 rooms
503 .get(&utils::u64_from_u8(short))
504 .cloned()
505 .ok_or_else(|| err!(Database("Conduit pduid short room id maps to no room")))
506}
507
508fn outlier_room(key: &[u8], pdu: &CanonicalJsonObject) -> Result<OwnedRoomId> {
512 let is_create = matches!(
513 pdu.get("type"),
514 Some(CanonicalJsonValue::String(kind)) if kind == "m.room.create"
515 );
516
517 if !is_create {
518 return Err!(Database("Conduit outlier lacks room_id and is not a create event"));
519 }
520
521 let event_id = <&EventId>::try_from(str::from_utf8(key)?)
522 .map_err(|_| err!(Database("Conduit outlier key is not a valid event id")))?;
523
524 RoomId::new_v2(event_id.localpart())
525 .map_err(|e| err!(Database("deriving room id from create event id: {e}")))
526}
527
528pub(super) async fn migrate_conduit_knocks(services: &Services) -> Result {
535 let knocks = copy_cf(services, "roomuserid_knockcount", "roomuserid_knockedcount").await?;
536 copy_cf(services, "userroomid_knockstate", "userroomid_knockedstate").await?;
537
538 if knocks > 0 {
539 warn!(%knocks, "Imported Conduit knocks");
540 }
541
542 Ok(())
543}
544
545pub(super) async fn migrate_conduit_highlight_split(services: &Services) -> Result {
556 let db = &services.db;
557 let highlight = db["userroomid_highlightcount"].clone();
558
559 if pin!(highlight.raw_keys_prefix(b"!"))
562 .next()
563 .await
564 .is_none()
565 {
566 return Ok(());
567 }
568
569 let lastread = db["roomuserid_lastnotificationread"].clone();
570 let cork = db.cork_and_sync();
571 let moved = highlight
572 .raw_stream()
573 .ignore_err()
574 .ready_fold(0_usize, |moved, (key, value)| {
575 if key.first() == Some(&b'!') {
576 lastread.insert(key, value);
577 highlight.remove(key);
578 moved.saturating_add(1)
579 } else {
580 moved
581 }
582 })
583 .await;
584
585 drop(cork);
586
587 if moved > 0 {
588 warn!(%moved, "Split Conduit last-notification-read rows out of the highlight-count column");
589 }
590
591 Ok(())
592}
593
594async fn copy_cf(
597 services: &Services,
598 source_name: &'static str,
599 target_name: &'static str,
600) -> Result<usize> {
601 let db = &services.db;
602 let Some(source) = db.open_cf(source_name)? else {
603 return Ok(0);
604 };
605
606 let target = &db[target_name];
607 let cork = db.cork_and_sync();
608 let copied = source
609 .raw_stream()
610 .ignore_err()
611 .ready_fold(0_usize, |copied, (key, value)| {
612 target.insert(key, value);
613 copied.saturating_add(1)
614 })
615 .await;
616
617 drop(cork);
618
619 Ok(copied)
620}
621
622#[cfg(test)]
623mod tests {
624 use std::path::Path;
625
626 use super::{
627 HasRoomId, conduit_media_key, conduit_media_path, parse_conduit_media_value, sha256_hex,
628 };
629
630 #[test]
631 fn conduit_media_path_deep_matches_conduit_default() {
632 let hex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
635 let path = conduit_media_path(Path::new("/db/media"), 2, 2, hex);
636
637 assert_eq!(
638 path,
639 Path::new(
640 "/db/media/01/23/456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
641 )
642 );
643 }
644
645 #[test]
646 fn conduit_media_path_flat_is_unsharded() {
647 let path = conduit_media_path(Path::new("/db/media"), 0, 2, "abcdef");
648
649 assert_eq!(path, Path::new("/db/media/abcdef"));
650 }
651
652 #[test]
653 fn conduit_media_key_deep_joins_shards_with_slash() {
654 let hex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
657 let key = conduit_media_key(2, 2, hex);
658
659 assert_eq!(key, "01/23/456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef");
660 }
661
662 #[test]
663 fn conduit_media_key_flat_is_bare_digest() {
664 assert_eq!(conduit_media_key(0, 2, "abcdef"), "abcdef");
665 }
666
667 #[test]
668 fn sha256_hex_encodes_lowercase_padded() {
669 assert_eq!(sha256_hex(&[0x00, 0x0F, 0xFF, 0xA5]), "000fffa5");
670 }
671
672 #[test]
673 fn conduit_media_value_ignores_unauthenticated_flag() {
674 let mut value = vec![7_u8; 32];
676 value.extend_from_slice(b"pic.png");
677 value.push(0xFF);
678 value.extend_from_slice(b"image/png");
679 value.push(0xFF);
680
681 let (sha256, filename, content_type) = parse_conduit_media_value(&value).unwrap();
682
683 assert_eq!(sha256, [7_u8; 32].as_slice());
684 assert_eq!(filename, Some("pic.png"));
685 assert_eq!(content_type, Some("image/png"));
686 }
687
688 #[test]
689 fn conduit_media_value_empty_filename_is_none() {
690 let mut value = vec![0_u8; 32];
691 value.push(0xFF);
692 value.extend_from_slice(b"image/png");
693
694 let (_, filename, content_type) = parse_conduit_media_value(&value).unwrap();
695
696 assert_eq!(filename, None);
697 assert_eq!(content_type, Some("image/png"));
698 }
699
700 #[test]
701 fn has_room_id_probe_detects_presence() {
702 let with_room_id = br#"{"room_id":"!r:server","type":"m.room.message"}"#;
703 let without_room_id = br#"{"type":"m.room.create","sender":"@u:server"}"#;
704
705 let present: HasRoomId = serde_json::from_slice(with_room_id).unwrap();
706 let absent: HasRoomId = serde_json::from_slice(without_room_id).unwrap();
707
708 assert!(present.room_id.is_some());
709 assert!(absent.room_id.is_none());
710 }
711}