// SPDX-License-Identifier: AGPL-3.0-or-later use crate::mention_extractor::{ MessageMentions, extend_mentions_from_markdown, extract_mentions_from_markdown, }; use crate::types::{ ApiChannelMentionResponse, ApiEmbedAuthorResponse, ApiEmbedFieldResponse, ApiEmbedFooterResponse, ApiEmbedMediaResponse, ApiEmbedProviderResponse, ApiMessageAttachmentResponse, ApiMessageCallResponse, ApiMessageEmbedChildResponse, ApiMessageEmbedResponse, ApiMessageReactionResponse, ApiMessageReferenceResponse, ApiMessageResponse, ApiMessageSnapshotResponse, ApiMessageStickerResponse, ApiReactionEmojiResponse, ApiUserPartial, Message, MessageAttachment, MessageCall, MessageEmbed, MessageEmbedAuthor, MessageEmbedChild, MessageEmbedField, MessageEmbedFooter, MessageEmbedMedia, MessageEmbedProvider, MessageReference, MessageRequest, MessageResponse, MessageSnapshot, MessageStickerItem, }; use crate::udt; use base64::prelude::{BASE64_STANDARD, Engine}; use chrono::{DateTime, Utc}; use fluxer_common::time::{unix_millis, unix_seconds}; use fluxer_common::user_partial::{FLUXER_SYSTEM_USER_ID, UserPartialFields, fluxer_system_user}; use fluxer_svc::shard::ShardService; use fluxer_svc::transport::Transport; use fluxer_svc::{postgres, postgres::BigIntBound, postgres::KeyPart}; use futures::stream::{self, StreamExt}; #[cfg(feature = "scylla")] use scylla::DeserializeRow; #[cfg(feature = "scylla")] use scylla::client::session::Session; #[cfg(feature = "scylla")] use scylla::response::query_result::QueryRowsResult; #[cfg(feature = "scylla")] use scylla::statement::prepared::PreparedStatement; use serde::Deserialize; use std::collections::{HashMap, HashSet}; #[cfg(feature = "scylla")] use std::sync::Arc; use url::Url; const BUCKET_DURATION_MS: i64 = 864_000_000; const FLUXER_EPOCH_MS: i64 = 1_420_070_400_000; const SERVICE_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); const MESSAGE_REFERENCE_TYPE_DEFAULT: i32 = 0; const MESSAGE_REFERENCE_TYPE_FORWARD: i32 = 1; const HIDDEN_REFILL_ROUNDS: usize = 5; const HIDDEN_REFILL_MAX_PAGE: u32 = 800; const MESSAGE_FLAG_IS_CROSSPOST: i64 = 1 << 1; fn effective_reference_type(reference: &MessageReference) -> i32 { reference .reference_type .unwrap_or(MESSAGE_REFERENCE_TYPE_DEFAULT) } fn is_crosspost_copy(message: &Message) -> bool { (message.flags.unwrap_or_default() & MESSAGE_FLAG_IS_CROSSPOST) != 0 } fn attachment_storage_channel_id(message: &Message) -> i64 { if !is_crosspost_copy(message) { return message.channel_id; } message .message_reference .as_ref() .and_then(|reference| reference.channel_id) .unwrap_or(message.channel_id) } fn reply_target(message: &Message) -> Option<(i64, i64)> { if is_crosspost_copy(message) { return None; } let reference = message.message_reference.as_ref()?; if effective_reference_type(reference) != MESSAGE_REFERENCE_TYPE_DEFAULT { return None; } Some((reference.channel_id?, reference.message_id?)) } fn copy_source(message: &Message) -> Option<(i64, i64)> { let reference = message.message_reference.as_ref()?; if !is_crosspost_copy(message) && effective_reference_type(reference) != MESSAGE_REFERENCE_TYPE_FORWARD { return None; } Some((reference.channel_id?, reference.message_id?)) } const MESSAGE_FLAG_SUPPRESS_EMBEDS: i64 = 1 << 2; const MESSAGE_TYPE_THREAD_CREATED: i32 = 18; const THREAD_MESSAGE_FLAG_MASK: i64 = 0x120; const THREAD_FEATURE_CHANNEL_TYPES: [i32; 5] = [10, 11, 12, 15, 16]; const EMBED_MEDIA_OWNED_ATTACHMENT_FLAG: i32 = 1 << 30; const DELETED_USER_USERNAME: &str = "DeletedUser"; const DELETED_USER_GLOBAL_NAME: &str = "Deleted User"; const BUCKET_SCAN_CONCURRENCY: usize = 16; const BUCKET_SCAN_WAVE: usize = 4; const ENRICHMENT_QUERY_CONCURRENCY: usize = 16; const REACTION_MESSAGE_BATCH_SIZE: usize = 64; const ATTACHMENT_DECAY_BATCH_SIZE: usize = 128; const BUCKET_INDEX_PAGE_SIZE: u32 = 200; const EXCLUDED_TYPE_REFILL_PAGE_SIZE: u32 = 100; const EXCLUDED_TYPE_SCAN_ROW_CAP: usize = 5_000; const JS_MAX_SAFE_INTEGER: i64 = 9_007_199_254_740_991; fn assert_safe_byte_size(value: i64) -> i64 { assert!( (0..=JS_MAX_SAFE_INTEGER).contains(&value), "attachment size must fit a non-negative JavaScript safe integer" ); value } #[cfg(feature = "scylla")] const MESSAGE_COLUMNS: &str = "\ message_id, channel_id, bucket, author_id, type, \ webhook_id, webhook_name, webhook_avatar_hash, \ content, edited_timestamp, pinned_timestamp, flags, mention_everyone, \ mention_users, mention_roles, mention_channels, \ has_reaction, version, \ attachments, embeds, sticker_items, message_reference, call, message_snapshots"; pub struct MessagesShard { storage: MessagesStorage, transport: T, } #[cfg(test)] type DeletedMessageKeys = std::sync::Arc>>; #[derive(Clone)] enum MessagesStorage { Postgres(PostgresMessagesStorage), #[cfg(feature = "scylla")] Scylla(Box), #[cfg(test)] Deletions(DeletedMessageKeys), #[cfg(test)] Memory(std::sync::Arc>), #[cfg(test)] Expiries(std::sync::Arc>>), } #[cfg(test)] fn memory_buckets( messages: &[Message], channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, descending: bool, ) -> Vec { let mut buckets = messages .iter() .filter(|message| message.channel_id == channel_id) .map(|message| snowflake_to_bucket(message.message_id)) .filter(|bucket| (min_bucket..=max_bucket).contains(bucket)) .collect::>(); buckets.sort_unstable(); buckets.dedup(); if descending { buckets.reverse(); } buckets.truncate(limit as usize); buckets } #[cfg(test)] fn memory_bucket( messages: &[Message], channel_id: i64, bucket: i32, bound: Option, limit: i32, ) -> Vec { let mut rows = messages .iter() .filter(|message| { message.channel_id == channel_id && snowflake_to_bucket(message.message_id) == bucket && match bound { Some(BucketBound::Before(id)) => message.message_id < id, Some(BucketBound::After(id)) => message.message_id > id, None => true, } }) .cloned() .collect::>(); if matches!(bound, Some(BucketBound::After(_))) { rows.sort_unstable_by_key(|message| message.message_id); } else { rows.sort_unstable_by_key(|message| std::cmp::Reverse(message.message_id)); } rows.truncate(limit.max(0) as usize); rows } #[derive(Clone)] struct PostgresMessagesStorage { kv: postgres::KvClient, } #[cfg(feature = "scylla")] #[derive(Clone)] struct ScyllaMessagesStorage { db: Arc, stmt_get_by_id: PreparedStatement, stmt_get_latest: PreparedStatement, stmt_get_before: PreparedStatement, stmt_get_after: PreparedStatement, stmt_list_buckets_desc: PreparedStatement, stmt_list_buckets_asc: PreparedStatement, stmt_get_reactions: PreparedStatement, stmt_get_reactions_for_messages: PreparedStatement, stmt_get_attachment_decay: PreparedStatement, stmt_get_attachment_decay_many: PreparedStatement, stmt_delete_message: PreparedStatement, } #[cfg_attr(feature = "scylla", derive(DeserializeRow))] #[derive(Debug)] struct MessageDbRow { message_id: i64, channel_id: i64, bucket: i32, author_id: Option, r#type: Option, webhook_id: Option, webhook_name: Option, webhook_avatar_hash: Option, content: Option, edited_timestamp: Option>, pinned_timestamp: Option>, flags: Option, mention_everyone: Option, mention_users: Option>, mention_roles: Option>, mention_channels: Option>, has_reaction: Option, version: Option, attachments: Option>, embeds: Option>, sticker_items: Option>, message_reference: Option, call: Option, message_snapshots: Option>, } #[cfg_attr(feature = "scylla", derive(DeserializeRow))] #[derive(Debug)] struct MessageReactionDbRow { user_id: i64, emoji_id: i64, emoji_name: String, emoji_animated: Option, created_at: Option>, } #[cfg(feature = "scylla")] #[derive(Debug, DeserializeRow)] struct ChannelMessageBucketDbRow { bucket: i32, } #[cfg(feature = "scylla")] #[derive(Debug, DeserializeRow)] struct BatchedMessageReactionDbRow { message_id: i64, user_id: i64, emoji_id: i64, emoji_name: String, emoji_animated: Option, created_at: Option>, } #[cfg(feature = "scylla")] #[derive(Debug, DeserializeRow)] struct AttachmentDecayDbRow { attachment_id: i64, expires_at: DateTime, } #[derive(Debug, Deserialize)] struct ChannelMessageBucketKvRow { bucket: i32, } #[derive(Debug, Deserialize)] struct MessageReactionKvRow { message_id: i64, user_id: i64, emoji_id: i64, emoji_name: String, emoji_animated: Option, created_at: Option, } #[derive(Debug, Deserialize)] struct AttachmentDecayKvRow { attachment_id: i64, expires_at: i64, } #[derive(Debug, Clone, Deserialize)] struct UserPartialServiceResponse { user_id: i64, username: String, discriminator: i32, global_name: Option, avatar_hash: Option, bot: Option, system: Option, flags: Option, avatar_color: Option, mention_flags: Option, #[serde(default)] content_hidden_since: Option, } #[derive(Debug, Default)] struct UserLookup { partials: HashMap, requested: HashSet, } impl UserLookup { fn hides(&self, author_id: Option, message_id: i64) -> bool { author_id .and_then(|author_id| self.partials.get(&author_id)) .and_then(|partial| partial.content_hidden_since) .is_some_and(|since| snowflake_to_epoch_millis(message_id) >= since) } fn merge(&mut self, other: UserLookup) { self.requested.extend(other.requested); self.partials.extend(other.partials); } } #[derive(Clone, Copy)] enum PageCursor { Latest, Before(i64), After(i64), } impl PageCursor { fn next(self, page: &[Message]) -> Option { let ids = page.iter().map(|message| message.message_id); match self { PageCursor::Latest | PageCursor::Before(_) => ids.min().map(PageCursor::Before), PageCursor::After(_) => ids.max().map(PageCursor::After), } } } #[derive(Debug, Deserialize)] enum UserServiceResponse { FoundPartials(Vec), } #[derive(Debug, Deserialize)] struct GatewayRpcResponse { ok: bool, result: Option, } #[derive(Debug, Deserialize)] struct GatewayChannelMentionsResult { channels: Option>, } #[derive(Debug, Deserialize)] struct GatewayChannelMention { id: String, name: String, #[serde(rename = "type")] channel_type: i32, } #[derive(Debug, Clone)] struct ResponseBuildOptions { viewer_user_id: i64, source_guild_id: Option, message_history_cutoff_ms: Option, can_read_message_history: bool, media_endpoint: String, media_proxy_secret_key: String, attachment_url_secret: Vec, include_reactions: bool, nonce: Option, tts: bool, include_hidden: bool, threads_mask: bool, excluded_types: Vec, attachment_decay: bool, } impl ResponseBuildOptions { fn expiring_attachment_ids(&self, messages: &[&Message]) -> HashSet { if self.attachment_decay { collect_attachment_ids(messages) } else { HashSet::new() } } fn hides(&self, message: &Message) -> bool { self.threads_mask && message.message_type == MESSAGE_TYPE_THREAD_CREATED } fn mask_flags(&self, flags: i64) -> i64 { if self.threads_mask { flags & !THREAD_MESSAGE_FLAG_MASK } else { flags } } } #[derive(Debug, Default)] struct ResponseContext { users: HashMap, reactions: HashMap>, referenced_messages: HashMap<(i64, i64), Message>, attachment_decay: HashMap>, channel_mentions: HashMap, mention_context: HashMap, } #[derive(Debug, Default)] struct MessageMentionContext { content: MessageMentions, snapshots: Vec, embed_users: HashSet, } impl MessagesShard { pub fn new_postgres(kv: postgres::KvClient, transport: T) -> anyhow::Result { Ok(Self { storage: MessagesStorage::Postgres(PostgresMessagesStorage { kv }), transport, }) } #[cfg(feature = "scylla")] pub async fn new_scylla(db: Arc, transport: T) -> anyhow::Result { let stmt_get_by_id = db .prepare(format!( "SELECT {MESSAGE_COLUMNS} FROM messages WHERE channel_id = ? AND bucket = ? AND message_id = ? LIMIT 1" )) .await?; let stmt_get_latest = db .prepare(format!( "SELECT {MESSAGE_COLUMNS} FROM messages WHERE channel_id = ? AND bucket = ? ORDER BY message_id DESC LIMIT ?" )) .await?; let stmt_get_before = db .prepare(format!( "SELECT {MESSAGE_COLUMNS} FROM messages WHERE channel_id = ? AND bucket = ? AND message_id < ? ORDER BY message_id DESC LIMIT ?" )) .await?; let stmt_get_after = db .prepare(format!( "SELECT {MESSAGE_COLUMNS} FROM messages WHERE channel_id = ? AND bucket = ? AND message_id > ? ORDER BY message_id ASC LIMIT ?" )) .await?; let stmt_list_buckets_desc = db .prepare( "SELECT bucket FROM channel_message_buckets WHERE channel_id = ? AND bucket >= ? AND bucket <= ? ORDER BY bucket DESC LIMIT ?" ) .await?; let stmt_list_buckets_asc = db .prepare( "SELECT bucket FROM channel_message_buckets WHERE channel_id = ? AND bucket >= ? AND bucket <= ? ORDER BY bucket ASC LIMIT ?" ) .await?; let stmt_get_reactions = db .prepare( "SELECT user_id, emoji_id, emoji_name, emoji_animated, created_at FROM message_reactions WHERE channel_id = ? AND bucket = ? AND message_id = ?" ) .await?; let stmt_get_reactions_for_messages = db .prepare( "SELECT message_id, user_id, emoji_id, emoji_name, emoji_animated, created_at FROM message_reactions WHERE channel_id = ? AND bucket = ? AND message_id IN ?" ) .await?; let stmt_get_attachment_decay = db .prepare( "SELECT attachment_id, expires_at FROM attachment_decay_by_id WHERE attachment_id = ? LIMIT 1" ) .await?; let stmt_get_attachment_decay_many = db .prepare( "SELECT attachment_id, expires_at FROM attachment_decay_by_id WHERE attachment_id IN ?" ) .await?; let stmt_delete_message = db .prepare("DELETE FROM messages WHERE channel_id = ? AND bucket = ? AND message_id = ?") .await?; Ok(Self { storage: MessagesStorage::Scylla(Box::new(ScyllaMessagesStorage { db, stmt_get_by_id, stmt_get_latest, stmt_get_before, stmt_get_after, stmt_list_buckets_desc, stmt_list_buckets_asc, stmt_get_reactions, stmt_get_reactions_for_messages, stmt_get_attachment_decay, stmt_get_attachment_decay_many, stmt_delete_message, })), transport, }) } async fn get_by_id(&self, channel_id: i64, message_id: i64) -> anyhow::Result> { let bucket = snowflake_to_bucket(message_id); self.storage.get_by_id(channel_id, bucket, message_id).await } async fn get_latest(&self, channel_id: i64, limit: u32) -> anyhow::Result> { if limit == 0 { return Ok(Vec::new()); } self.scan_indexed_buckets_desc(channel_id, 0, current_bucket(), limit, None) .await } async fn get_before( &self, channel_id: i64, before_id: i64, limit: u32, ) -> anyhow::Result> { if limit == 0 { return Ok(Vec::new()); } self.scan_indexed_buckets_desc( channel_id, 0, snowflake_to_bucket(before_id), limit, Some(before_id), ) .await } async fn get_after( &self, channel_id: i64, after_id: i64, limit: u32, ) -> anyhow::Result> { if limit == 0 { return Ok(Vec::new()); } let min_bucket = snowflake_to_bucket(after_id); let max_bucket = current_bucket().max(min_bucket); self.scan_indexed_buckets_asc(channel_id, min_bucket, max_bucket, limit, after_id) .await } async fn scan_indexed_buckets_desc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, before_id: Option, ) -> anyhow::Result> { let mut messages = Vec::new(); let mut cursor_max = max_bucket; let limit_i32 = limit.min(i32::MAX as u32) as i32; while messages.len() < limit as usize && cursor_max >= min_bucket { let buckets = self .list_buckets_desc(channel_id, min_bucket, cursor_max, bucket_page_limit(limit)) .await?; let Some(last_bucket) = buckets.last().copied() else { break; }; collect_bucket_waves( &buckets, limit as usize, &mut messages, |bucket| async move { if let Some(before_id) = before_id { self.fetch_before_bucket(channel_id, bucket, before_id, limit_i32) .await } else { self.fetch_latest_bucket(channel_id, bucket, limit_i32) .await } }, ) .await?; if messages.len() >= limit as usize || last_bucket <= min_bucket { break; } cursor_max = last_bucket.saturating_sub(1); } messages.sort_unstable_by_key(|message| std::cmp::Reverse(message.message_id)); messages.truncate(limit as usize); Ok(messages) } async fn scan_indexed_buckets_asc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, after_id: i64, ) -> anyhow::Result> { let mut messages = Vec::new(); let mut cursor_min = min_bucket; let limit_i32 = limit.min(i32::MAX as u32) as i32; while messages.len() < limit as usize && cursor_min <= max_bucket { let buckets = self .list_buckets_asc(channel_id, cursor_min, max_bucket, bucket_page_limit(limit)) .await?; let Some(last_bucket) = buckets.last().copied() else { break; }; collect_bucket_waves( &buckets, limit as usize, &mut messages, |bucket| async move { self.fetch_after_bucket(channel_id, bucket, after_id, limit_i32) .await }, ) .await?; if messages.len() >= limit as usize || last_bucket >= max_bucket { break; } cursor_min = last_bucket.saturating_add(1); } messages.sort_unstable_by_key(|left| left.message_id); messages.truncate(limit as usize); Ok(messages) } async fn list_buckets_desc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { if min_bucket > max_bucket || limit == 0 { return Ok(Vec::new()); } self.storage .list_buckets_desc(channel_id, min_bucket, max_bucket, limit) .await } async fn list_buckets_asc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { if min_bucket > max_bucket || limit == 0 { return Ok(Vec::new()); } self.storage .list_buckets_asc(channel_id, min_bucket, max_bucket, limit) .await } async fn fetch_latest_bucket( &self, channel_id: i64, bucket: i32, limit: i32, ) -> anyhow::Result> { self.storage .fetch_latest_bucket(channel_id, bucket, limit) .await } async fn fetch_before_bucket( &self, channel_id: i64, bucket: i32, before_id: i64, limit: i32, ) -> anyhow::Result> { self.storage .fetch_before_bucket(channel_id, bucket, before_id, limit) .await } async fn fetch_after_bucket( &self, channel_id: i64, bucket: i32, after_id: i64, limit: i32, ) -> anyhow::Result> { self.storage .fetch_after_bucket(channel_id, bucket, after_id, limit) .await } async fn fetch_page( &self, channel_id: i64, cursor: PageCursor, limit: u32, floor: Option, excluded_types: &[i32], ) -> anyhow::Result> { if excluded_types.is_empty() { return match cursor { PageCursor::Latest => self.get_latest(channel_id, limit).await, PageCursor::Before(before_id) => { self.get_before(channel_id, before_id, limit).await } PageCursor::After(after_id) => self.get_after(channel_id, after_id, limit).await, }; } match cursor { PageCursor::Latest => { collect_excluding(limit, excluded_types, None, |cursor, size| async move { match cursor { Some(cursor) => self.get_before(channel_id, cursor, size).await, None => self.get_latest(channel_id, size).await, } }) .await } PageCursor::Before(before_id) => { collect_excluding( limit, excluded_types, Some(before_id), |cursor, size| async move { let mut page = self .get_before(channel_id, cursor.unwrap_or(before_id), size) .await?; if let Some(floor) = floor { page.retain(|message| message.message_id > floor); } Ok(page) }, ) .await } PageCursor::After(after_id) => { collect_excluding(limit, excluded_types, Some(after_id), |cursor, size| { self.get_after(channel_id, cursor.unwrap_or(after_id), size) }) .await } } } async fn visible_page( &self, channel_id: i64, start: PageCursor, limit: u32, floor: Option, options: &ResponseBuildOptions, users: &mut UserLookup, ) -> anyhow::Result> { let mut out = Vec::new(); let mut cursor = start; let mut page_size = limit; for _ in 0..HIDDEN_REFILL_ROUNDS { if limit == 0 { break; } let page = self .fetch_page( channel_id, cursor, page_size, floor, &options.excluded_types, ) .await?; let exhausted = page.len() < page_size as usize; let reached_floor = floor.is_some_and(|floor| page.iter().any(|message| message.message_id <= floor)); let next = cursor.next(&page); let hidden = self.hidden_message_ids(&page, options, users).await; let refill = !hidden.is_empty(); out.extend( page.into_iter() .filter(|message| !hidden.contains(&message.message_id)), ); let Some(next) = next else { break; }; if !refill || exhausted || reached_floor || out.len() >= limit as usize { break; } cursor = next; page_size = page_size .saturating_mul(2) .min(HIDDEN_REFILL_MAX_PAGE.max(limit)); } if matches!(start, PageCursor::After(_)) { out.sort_unstable_by_key(|message| message.message_id); } else { out.sort_unstable_by_key(|message| std::cmp::Reverse(message.message_id)); } out.truncate(limit as usize); Ok(out) } async fn visible_around( &self, channel_id: i64, around_id: i64, limit: u32, options: &ResponseBuildOptions, users: &mut UserLookup, ) -> anyhow::Result> { if limit == 0 { return Ok(Vec::new()); } let (newer_limit, older_limit) = around_window_limits(limit); let mut target_users = UserLookup::default(); let mut newer_users = UserLookup::default(); let mut older_users = UserLookup::default(); let (target, newer, older) = tokio::try_join!( async { let target = self.get_by_id(channel_id, around_id).await?; let Some(target) = target else { return anyhow::Ok(None); }; let hidden = self .hidden_message_ids(std::slice::from_ref(&target), options, &mut target_users) .await; Ok((!hidden.contains(&target.message_id) && !options.excluded_types.contains(&target.message_type)) .then_some(target)) }, self.visible_page( channel_id, PageCursor::After(around_id), newer_limit, None, options, &mut newer_users, ), self.visible_page( channel_id, PageCursor::Before(around_id), older_limit, None, options, &mut older_users, ) )?; users.merge(target_users); users.merge(newer_users); users.merge(older_users); let mut out = Vec::new(); let mut seen = HashSet::new(); for message in newer.into_iter().rev() { if seen.insert(message.message_id) { out.push(message); } } if let Some(message) = target && seen.insert(message.message_id) { out.push(message); } for message in older { if seen.insert(message.message_id) { out.push(message); } } out.truncate(limit as usize); Ok(out) } async fn list_api_responses( &self, channel_id: i64, limit: u32, before_id: Option, after_id: Option, around_id: Option, options: ResponseBuildOptions, ) -> anyhow::Result> { if !options.can_read_message_history && options.message_history_cutoff_ms.is_none() { return Ok(Vec::new()); } let mut users = UserLookup::default(); let mut messages = if let Some(around_id) = around_id { self.visible_around(channel_id, around_id, limit, &options, &mut users) .await? } else if let (Some(before_id), Some(after_id)) = (before_id, after_id) { let mut before = self .visible_page( channel_id, PageCursor::Before(before_id), limit, Some(after_id), &options, &mut users, ) .await?; before.retain(|message| message.message_id > after_id); before } else { let cursor = match (before_id, after_id) { (Some(before_id), _) => PageCursor::Before(before_id), (None, Some(after_id)) => PageCursor::After(after_id), (None, None) => PageCursor::Latest, }; self.visible_page(channel_id, cursor, limit, None, &options, &mut users) .await? }; messages .retain(|message| self.is_message_visible_to_requester(message.message_id, &options)); let (mut messages, orphaned_messages): (Vec<_>, Vec<_>) = messages .into_iter() .partition(|message| message.author_id.is_some() || message.webhook_id.is_some()); self.cleanup_orphaned_messages(orphaned_messages).await; messages.sort_unstable_by_key(|message| std::cmp::Reverse(message.message_id)); let context = self .build_response_context(&messages, &options, true, users) .await?; Ok(messages .iter() .map(|message| self.map_message_response(message, &options, &context, true)) .collect()) } async fn get_api_response( &self, channel_id: i64, message_id: i64, options: ResponseBuildOptions, ) -> anyhow::Result> { if !self.is_message_visible_to_requester(message_id, &options) { return Ok(None); } let Some(message) = self.get_by_id(channel_id, message_id).await? else { return Ok(None); }; if options.hides(&message) { return Ok(None); } if message.author_id.is_none() && message.webhook_id.is_none() { self.cleanup_orphaned_messages(vec![message]).await; return Ok(None); } let mut users = UserLookup::default(); let hidden = self .hidden_message_ids(std::slice::from_ref(&message), &options, &mut users) .await; if !hidden.is_empty() { return Ok(None); } let context = self .build_response_context(std::slice::from_ref(&message), &options, true, users) .await?; Ok(Some( self.map_message_response(&message, &options, &context, true), )) } async fn build_api_response_from_message( &self, message: Message, options: ResponseBuildOptions, ) -> anyhow::Result> { if !self.is_message_visible_to_requester(message.message_id, &options) || options.hides(&message) { return Ok(None); } if message.author_id.is_none() && message.webhook_id.is_none() { self.cleanup_orphaned_messages(vec![message]).await; return Ok(None); } let context = self .build_response_context( std::slice::from_ref(&message), &options, true, UserLookup::default(), ) .await?; Ok(Some( self.map_message_response(&message, &options, &context, true), )) } async fn build_api_responses_from_messages( &self, messages: Vec, options: ResponseBuildOptions, ) -> anyhow::Result> { let (messages, orphaned_messages): (Vec<_>, Vec<_>) = messages .into_iter() .filter(|message| { self.is_message_visible_to_requester(message.message_id, &options) && !options.hides(message) }) .partition(|message| message.author_id.is_some() || message.webhook_id.is_some()); self.cleanup_orphaned_messages(orphaned_messages).await; let mut users = UserLookup::default(); let hidden = self .hidden_message_ids(&messages, &options, &mut users) .await; let messages = messages .into_iter() .filter(|message| !hidden.contains(&message.message_id)) .collect::>(); let context = self .build_response_context(&messages, &options, true, users) .await?; Ok(messages .iter() .map(|message| self.map_message_response(message, &options, &context, true)) .collect()) } async fn hidden_message_ids( &self, messages: &[Message], options: &ResponseBuildOptions, users: &mut UserLookup, ) -> HashSet { if options.include_hidden || messages.is_empty() { return HashSet::new(); } let copies = messages .iter() .filter_map(|message| copy_source(message).map(|source| (message.message_id, source))) .collect::>(); let sources = stream::iter(copies) .map(|(message_id, (channel_id, source_id))| async move { match self.get_by_id(channel_id, source_id).await { Ok(Some(source)) => Some((message_id, (source.author_id, source.message_id))), _ => None, } }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .filter_map(|source| async move { source }) .collect::>() .await; let user_ids = messages .iter() .flat_map(|message| { message .author_id .into_iter() .chain(message.mention_users.iter().copied()) }) .chain(sources.values().filter_map(|(author_id, _)| *author_id)) .collect::>(); self.load_user_partials(user_ids, users).await; messages .iter() .filter(|message| { users.hides(message.author_id, message.message_id) || sources .get(&message.message_id) .is_some_and(|(author_id, source_id)| users.hides(*author_id, *source_id)) }) .map(|message| message.message_id) .collect() } fn is_message_visible_to_requester( &self, message_id: i64, options: &ResponseBuildOptions, ) -> bool { if options.can_read_message_history { return true; } let Some(cutoff) = options.message_history_cutoff_ms else { return false; }; snowflake_to_epoch_millis(message_id) >= cutoff } async fn cleanup_orphaned_messages(&self, messages: Vec) { if messages.is_empty() { return; } let count = messages.len(); let failures = stream::iter(messages) .map(|message| async move { let bucket = snowflake_to_bucket(message.message_id); self.storage .delete_message(message.channel_id, bucket, message.message_id) .await .map_err(|error| (message.channel_id, message.message_id, error)) }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .filter_map(|result| async move { result.err() }) .collect::>() .await; if failures.is_empty() { tracing::info!(count, "cleaned orphaned messages during response build"); } else { tracing::warn!( count, failure_count = failures.len(), first_channel_id = failures[0].0, first_message_id = failures[0].1, error = %failures[0].2, "failed to clean some orphaned messages during response build" ); } } async fn build_response_context( &self, messages: &[Message], options: &ResponseBuildOptions, include_referenced_messages: bool, mut users: UserLookup, ) -> anyhow::Result { let mut referenced_messages = if include_referenced_messages { self.fetch_referenced_messages(messages, options).await } else { HashMap::new() }; let all_messages = messages .iter() .chain(referenced_messages.values()) .collect::>(); let mention_context = build_message_mention_context(&all_messages); let attachment_ids = options.expiring_attachment_ids(&all_messages); let channel_ids = collect_channel_mention_ids(&all_messages, &mention_context); let user_ids = collect_user_ids(&all_messages, &mention_context); let reactions_future = self.fetch_reactions_for_messages(messages, options); let attachment_decay_future = self.fetch_attachment_decay(attachment_ids); let channel_mentions_future = self.resolve_channel_mentions(channel_ids, options); let users_future = self.load_user_partials(user_ids, &mut users); let (reactions, attachment_decay, channel_mentions, ()) = tokio::join!( reactions_future, attachment_decay_future, channel_mentions_future, users_future ); let attachment_decay = attachment_decay?; if !options.include_hidden { referenced_messages .retain(|_, referenced| !users.hides(referenced.author_id, referenced.message_id)); } let users = users .partials .into_values() .map(|partial| (partial.user_id, map_user_partial(partial))) .collect(); Ok(ResponseContext { users, reactions, referenced_messages, attachment_decay, channel_mentions, mention_context, }) } async fn fetch_referenced_messages( &self, messages: &[Message], options: &ResponseBuildOptions, ) -> HashMap<(i64, i64), Message> { let mut refs = HashSet::new(); for message in messages { let Some((channel_id, message_id)) = reply_target(message) else { continue; }; if !self.is_message_visible_to_requester(message_id, options) { continue; } refs.insert((channel_id, message_id)); } stream::iter(refs) .map(|(channel_id, message_id)| async move { match self.get_by_id(channel_id, message_id).await { Ok(Some(message)) if message.author_id.is_some() || message.webhook_id.is_some() => { Some(((channel_id, message_id), message)) } _ => None, } }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .collect::>() .await .into_iter() .flatten() .collect() } async fn fetch_reactions_for_messages( &self, messages: &[Message], options: &ResponseBuildOptions, ) -> HashMap> { if !options.include_reactions { return HashMap::new(); } let mut groups: HashMap<(i64, i32), Vec> = HashMap::new(); for message in messages { if message.has_reaction == Some(false) { continue; } let bucket = snowflake_to_bucket(message.message_id); groups .entry((message.channel_id, bucket)) .or_default() .push(message.message_id); } let mut batches = Vec::new(); for ((channel_id, bucket), mut message_ids) in groups { message_ids.sort_unstable(); message_ids.dedup(); for chunk in message_ids.chunks(REACTION_MESSAGE_BATCH_SIZE) { batches.push((channel_id, bucket, chunk.to_vec())); } } stream::iter(batches) .map(|(channel_id, bucket, message_ids)| async move { self.fetch_reactions_for_message_batch( channel_id, bucket, message_ids, options.viewer_user_id, ) .await }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .collect::>() .await .into_iter() .flatten() .collect() } async fn fetch_reactions_for_message_batch( &self, channel_id: i64, bucket: i32, message_ids: Vec, viewer_user_id: i64, ) -> HashMap> { if message_ids.is_empty() { return HashMap::new(); } self.storage .fetch_reactions_for_message_batch(channel_id, bucket, message_ids, viewer_user_id) .await } async fn fetch_attachment_decay( &self, attachment_ids: HashSet, ) -> anyhow::Result>> { if attachment_ids.is_empty() { return Ok(HashMap::new()); } let mut attachment_ids = attachment_ids.into_iter().collect::>(); attachment_ids.sort_unstable(); attachment_ids.dedup(); let batches = attachment_ids .chunks(ATTACHMENT_DECAY_BATCH_SIZE) .map(<[i64]>::to_vec) .collect::>(); let maps = stream::iter(batches) .map(|attachment_ids| async move { self.fetch_attachment_decay_batch(attachment_ids).await }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .collect::>() .await .into_iter() .collect::>>()?; Ok(maps.into_iter().flatten().collect()) } async fn fetch_attachment_decay_batch( &self, attachment_ids: Vec, ) -> anyhow::Result>> { self.storage .fetch_attachment_decay_batch(attachment_ids) .await } async fn load_user_partials(&self, user_ids: HashSet, users: &mut UserLookup) { let mut missing = user_ids .into_iter() .filter(|user_id| users.requested.insert(*user_id)) .collect::>(); if missing.is_empty() { return; } missing.sort_unstable(); let payload = serde_json::json!({ "op": "GetPartialsByIds", "user_ids": missing, }); let payload_bytes = serde_json::to_vec(&payload).unwrap_or_default(); let response = self .transport .request( "svc.users", payload_bytes.as_slice(), SERVICE_REQUEST_TIMEOUT, ) .await .ok() .and_then(|bytes| serde_json::from_slice::(&bytes).ok()); let partials = match response { Some(UserServiceResponse::FoundPartials(partials)) => partials, _ => Vec::new(), }; users.partials.extend( partials .into_iter() .map(|partial| (partial.user_id, partial)), ); } async fn resolve_channel_mentions( &self, channel_ids: HashSet, options: &ResponseBuildOptions, ) -> HashMap { let Some(guild_id) = options.source_guild_id else { return HashMap::new(); }; if channel_ids.is_empty() { return HashMap::new(); } let mut channel_ids: Vec = channel_ids.into_iter().map(|id| id.to_string()).collect(); channel_ids.sort(); channel_ids.dedup(); let mut payload = serde_json::json!({ "guild_id": guild_id.to_string(), "channel_ids": channel_ids, }); if options.threads_mask { payload["exclude_thread_channels"] = serde_json::Value::Bool(true); } let payload_bytes = serde_json::to_vec(&payload).unwrap_or_default(); let response = self .transport .request( "rpc.gateway.guild.resolve_channel_mentions", payload_bytes.as_slice(), SERVICE_REQUEST_TIMEOUT, ) .await .ok() .and_then(|bytes| { serde_json::from_slice::>(&bytes) .ok() }); let Some(response) = response else { return HashMap::new(); }; if !response.ok { return HashMap::new(); } response .result .and_then(|result| result.channels) .unwrap_or_default() .into_iter() .filter(|channel| { !(options.threads_mask && THREAD_FEATURE_CHANNEL_TYPES.contains(&channel.channel_type)) }) .map(|channel| { let item = ApiChannelMentionResponse { id: channel.id, name: channel.name, channel_type: channel.channel_type, }; (item.id.clone(), item) }) .collect() } fn map_message_response( &self, message: &Message, options: &ResponseBuildOptions, context: &ResponseContext, include_referenced_message: bool, ) -> ApiMessageResponse { let author = self.resolve_author(message, context); let storage_channel_id = attachment_storage_channel_id(message); let attachments = message .attachments .as_deref() .unwrap_or_default() .iter() .filter_map(|attachment| { self.map_attachment(storage_channel_id, attachment, options, context) }) .collect(); let embeds = if (message.flags.unwrap_or_default() & MESSAGE_FLAG_SUPPRESS_EMBEDS) == 0 { message .embeds .as_deref() .unwrap_or_default() .iter() .map(|embed| self.map_embed(embed, options)) .collect() } else { Vec::new() }; let stickers = message .sticker_items .as_deref() .unwrap_or_default() .iter() .filter_map(map_sticker) .collect(); let fallback_mentions; let message_mentions = match context.mention_context.get(&message.message_id) { Some(mentions) => mentions, None => { fallback_mentions = build_mention_context_entry(message); &fallback_mentions } }; let content_mentions = &message_mentions.content; let mention_roles = ids_present_in_set(&message.mention_roles, &content_mentions.roles); let mention_channels = ids_present_in_set(&message.mention_channels, &content_mentions.channels) .into_iter() .filter_map(|id| context.channel_mentions.get(&id).cloned()) .collect::>(); let mut referenced_user_ids = content_mentions.users.clone(); referenced_user_ids.extend(message_mentions.embed_users.iter().copied()); for snapshot_mentions in &message_mentions.snapshots { referenced_user_ids.extend(snapshot_mentions.users.iter().copied()); } let mentioned_user_ids = message .mention_users .iter() .copied() .filter(|id| content_mentions.users.contains(id) || !referenced_user_ids.contains(id)) .collect::>(); let mentions = mentioned_user_ids .iter() .filter_map(|id| context.users.get(id).cloned()) .collect::>(); for id in &mentioned_user_ids { referenced_user_ids.remove(id); } let users = referenced_user_ids .into_iter() .filter_map(|id| context.users.get(&id).cloned()) .collect::>(); let referenced_message = if include_referenced_message { reply_target(message).map(|(channel_id, message_id)| { context .referenced_messages .get(&(channel_id, message_id)) .filter(|referenced| !options.hides(referenced)) .map(|referenced| { let mut referenced_options = options.clone(); referenced_options.nonce = None; referenced_options.tts = false; Box::new(self.map_message_response( referenced, &referenced_options, context, false, )) }) }) } else { None }; ApiMessageResponse { id: message.message_id.to_string(), channel_id: message.channel_id.to_string(), author, webhook_id: message.webhook_id.map(|id| id.to_string()), message_type: message.message_type, flags: options.mask_flags(message.flags.unwrap_or_default()), content: message.content.clone().unwrap_or_default(), timestamp: epoch_millis_to_iso(snowflake_to_epoch_millis(message.message_id)), edited_timestamp: message.edited_timestamp.map(epoch_millis_to_iso), pinned: message.pinned.unwrap_or(false), mention_everyone: message.mention_everyone.unwrap_or(false), tts: options.tts, mentions, mention_roles, mention_channels: (!mention_channels.is_empty()).then_some(mention_channels), users: (!users.is_empty()).then_some(users), embeds, attachments, stickers, reactions: context.reactions.get(&message.message_id).cloned(), message_reference: message .message_reference .as_ref() .and_then(map_message_reference), message_snapshots: message.message_snapshots.as_ref().and_then(|snapshots| { let mapped: Vec<_> = snapshots .iter() .enumerate() .map(|(index, snapshot)| { self.map_snapshot( message.message_id, message.channel_id, index, snapshot, options, context, ) }) .collect(); (!mapped.is_empty()).then_some(mapped) }), nonce: options.nonce.clone(), call: message.call.as_ref().map(map_call), referenced_message, } } fn resolve_author(&self, message: &Message, context: &ResponseContext) -> ApiUserPartial { if let Some(author_id) = message.author_id { return context .users .get(&author_id) .cloned() .unwrap_or_else(|| deleted_user(author_id)); } if let (Some(webhook_id), Some(webhook_name)) = (message.webhook_id, message.webhook_name.clone()) { return ApiUserPartial { id: webhook_id.to_string(), username: webhook_name, discriminator: "0000".to_owned(), global_name: None, avatar: message.webhook_avatar_hash.clone(), avatar_color: None, bot: Some(true), system: None, flags: 0, mention_flags: None, }; } deleted_user(0) } fn map_attachment( &self, channel_id: i64, attachment: &MessageAttachment, options: &ResponseBuildOptions, context: &ResponseContext, ) -> Option { let attachment_id = attachment.attachment_id?; let filename = attachment.filename.clone().unwrap_or_default(); let decay = context.attachment_decay.get(&attachment_id); let expired = decay.is_some_and(|expires_at| expires_at.timestamp_millis() <= unix_millis()); let url = (!expired).then(|| { attachment_cdn_url( channel_id, attachment_id, &filename, options, unix_seconds(), ) }); let content_type = attachment.content_type.clone().unwrap_or_else(|| { mime_guess::from_path(&filename) .first_or_octet_stream() .essence_str() .to_owned() }); let is_audio = content_type.to_ascii_lowercase().starts_with("audio/"); Some(ApiMessageAttachmentResponse { id: attachment_id.to_string(), filename, title: attachment.title.clone(), description: attachment.description.clone(), content_type: Some(content_type), content_hash: attachment.content_hash.clone(), size: assert_safe_byte_size(attachment.size.unwrap_or_default()), url: url.clone(), proxy_url: url, width: (!is_audio).then_some(attachment.width).flatten(), height: (!is_audio).then_some(attachment.height).flatten(), placeholder: attachment.placeholder.clone(), flags: attachment.flags.unwrap_or_default(), nsfw: attachment.nsfw, duration: attachment.duration_secs.or(attachment.duration), waveform: attachment.waveform.clone(), expires_at: decay.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)), expired: expired.then_some(true), }) } fn map_embed( &self, embed: &MessageEmbed, options: &ResponseBuildOptions, ) -> ApiMessageEmbedResponse { let base = self.map_embed_child_like( embed.embed_type.clone(), embed.title.clone(), embed.description.clone(), embed.url.clone(), embed.timestamp, embed.color, embed.author.clone(), embed.provider.clone(), embed.thumbnail.clone(), embed.image.clone(), embed.video.clone(), embed.audio.clone(), embed.footer.clone(), embed.fields.clone(), embed.nsfw, embed.html.clone(), embed.html_width, embed.html_height, options, ); let children = embed.children.as_ref().map(|children| { children .iter() .map(|child| { self.map_embed_child_like( child.embed_type.clone(), child.title.clone(), child.description.clone(), child.url.clone(), child.timestamp, child.color, child.author.clone(), child.provider.clone(), child.thumbnail.clone(), child.image.clone(), child.video.clone(), child.audio.clone(), child.footer.clone(), child.fields.clone(), child.nsfw, child.html.clone(), child.html_width, child.html_height, options, ) }) .collect::>() }); ApiMessageEmbedResponse { base, children: children.filter(|children| !children.is_empty()), } } #[allow(clippy::too_many_arguments)] fn map_embed_child_like( &self, embed_type: Option, title: Option, description: Option, url: Option, timestamp: Option, color: Option, author: Option, provider: Option, thumbnail: Option, image: Option, video: Option, audio: Option, footer: Option, fields: Option>, nsfw: Option, html: Option, html_width: Option, html_height: Option, options: &ResponseBuildOptions, ) -> ApiMessageEmbedChildResponse { ApiMessageEmbedChildResponse { embed_type: embed_type.unwrap_or_else(|| "rich".to_owned()), title, description, url: url.map(|url| sign_own_url(&url, options)), timestamp: timestamp.map(epoch_millis_to_iso), color, author: author.and_then(|author| { author.name.map(|name| ApiEmbedAuthorResponse { name, url: author.url.map(|url| sign_own_url(&url, options)), proxy_icon_url: author .icon_url .as_ref() .map(|url| external_media_proxy_url(url, options)), icon_url: author.icon_url.map(|url| sign_own_url(&url, options)), }) }), provider: provider.and_then(|provider| { provider.name.map(|name| ApiEmbedProviderResponse { name, url: provider.url.map(|url| sign_own_url(&url, options)), icon_url: None, proxy_icon_url: None, }) }), thumbnail: thumbnail.and_then(|media| self.map_embed_media(media, options)), image: image.and_then(|media| self.map_embed_media(media, options)), video: video.and_then(|media| self.map_embed_media(media, options)), audio: audio.and_then(|media| self.map_embed_media(media, options)), footer: footer.and_then(|footer| { footer.text.map(|text| ApiEmbedFooterResponse { text, proxy_icon_url: footer .icon_url .as_ref() .map(|url| external_media_proxy_url(url, options)), icon_url: footer.icon_url.map(|url| sign_own_url(&url, options)), }) }), fields: fields.map(|fields| fields.into_iter().map(map_embed_field_response).collect()), nsfw, html, html_width, html_height, } } fn map_embed_media( &self, media: MessageEmbedMedia, options: &ResponseBuildOptions, ) -> Option { let url = media.url?; Some(ApiEmbedMediaResponse { proxy_url: external_media_proxy_url(&url, options), url: sign_own_url(&url, options), width: media.width, height: media.height, duration: media.duration, description: media.description, content_type: media.content_type, content_hash: media.content_hash, placeholder: media.placeholder, flags: media .flags .map(|flags| flags & !EMBED_MEDIA_OWNED_ATTACHMENT_FLAG), }) } fn map_snapshot( &self, message_id: i64, channel_id: i64, snapshot_index: usize, snapshot: &MessageSnapshot, options: &ResponseBuildOptions, context: &ResponseContext, ) -> ApiMessageSnapshotResponse { let snapshot_mentions = context .mention_context .get(&message_id) .and_then(|mentions| mentions.snapshots.get(snapshot_index)) .cloned() .unwrap_or_else(|| extract_mentions_from_markdown(snapshot.content.as_deref())); let mention_channels = ids_present_in_set(&snapshot.mention_channels, &snapshot_mentions.channels) .into_iter() .filter_map(|id| context.channel_mentions.get(&id).cloned()) .collect::>(); let embeds = if (snapshot.flags.unwrap_or_default() & MESSAGE_FLAG_SUPPRESS_EMBEDS as i32) == 0 { snapshot.embeds.as_deref().map(|embeds| { embeds .iter() .map(|embed| self.map_embed(embed, options)) .collect() }) } else { None }; let attachments = snapshot.attachments.as_ref().map(|attachments| { attachments .iter() .filter_map(|attachment| { self.map_attachment(channel_id, attachment, options, context) }) .collect() }); ApiMessageSnapshotResponse { content: snapshot.content.clone(), timestamp: snapshot .timestamp .map(epoch_millis_to_iso) .unwrap_or_else(|| epoch_millis_to_iso(0)), edited_timestamp: snapshot.edited_timestamp.map(epoch_millis_to_iso), mentions: (!snapshot.mention_users.is_empty()).then(|| { snapshot .mention_users .iter() .map(ToString::to_string) .collect() }), mention_roles: (!snapshot.mention_roles.is_empty()).then(|| { snapshot .mention_roles .iter() .map(ToString::to_string) .collect() }), mention_channels: (!mention_channels.is_empty()).then_some(mention_channels), embeds, attachments, stickers: snapshot .sticker_items .as_ref() .map(|stickers| stickers.iter().filter_map(map_sticker).collect::>()), snapshot_type: snapshot.snapshot_type.unwrap_or_default(), flags: options.mask_flags(i64::from(snapshot.flags.unwrap_or_default())) as i32, } } } impl MessagesStorage { async fn get_by_id( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage.get_by_id(channel_id, bucket, message_id).await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage.get_by_id(channel_id, bucket, message_id).await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(None), #[cfg(test)] MessagesStorage::Memory(messages) => Ok(messages .iter() .find(|message| { message.channel_id == channel_id && message.message_id == message_id }) .cloned()), } } async fn list_buckets_desc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage .list_buckets(channel_id, min_bucket, max_bucket, limit, true) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage .list_buckets_desc(channel_id, min_bucket, max_bucket, limit) .await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(Vec::new()), #[cfg(test)] MessagesStorage::Memory(messages) => Ok(memory_buckets( messages, channel_id, min_bucket, max_bucket, limit, true, )), } } async fn list_buckets_asc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage .list_buckets(channel_id, min_bucket, max_bucket, limit, false) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage .list_buckets_asc(channel_id, min_bucket, max_bucket, limit) .await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(Vec::new()), #[cfg(test)] MessagesStorage::Memory(messages) => Ok(memory_buckets( messages, channel_id, min_bucket, max_bucket, limit, false, )), } } async fn fetch_latest_bucket( &self, channel_id: i64, bucket: i32, limit: i32, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage .fetch_bucket(channel_id, bucket, None, true, limit) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage.fetch_latest_bucket(channel_id, bucket, limit).await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(Vec::new()), #[cfg(test)] MessagesStorage::Memory(messages) => { Ok(memory_bucket(messages, channel_id, bucket, None, limit)) } } } async fn fetch_before_bucket( &self, channel_id: i64, bucket: i32, before_id: i64, limit: i32, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage .fetch_bucket( channel_id, bucket, Some(BucketBound::Before(before_id)), true, limit, ) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage .fetch_before_bucket(channel_id, bucket, before_id, limit) .await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(Vec::new()), #[cfg(test)] MessagesStorage::Memory(messages) => Ok(memory_bucket( messages, channel_id, bucket, Some(BucketBound::Before(before_id)), limit, )), } } async fn fetch_after_bucket( &self, channel_id: i64, bucket: i32, after_id: i64, limit: i32, ) -> anyhow::Result> { match self { MessagesStorage::Postgres(storage) => { storage .fetch_bucket( channel_id, bucket, Some(BucketBound::After(after_id)), false, limit, ) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage .fetch_after_bucket(channel_id, bucket, after_id, limit) .await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Expiries(_) => Ok(Vec::new()), #[cfg(test)] MessagesStorage::Memory(messages) => Ok(memory_bucket( messages, channel_id, bucket, Some(BucketBound::After(after_id)), limit, )), } } async fn delete_message( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result<()> { match self { MessagesStorage::Postgres(storage) => { storage.delete_message(channel_id, bucket, message_id).await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage.delete_message(channel_id, bucket, message_id).await } #[cfg(test)] MessagesStorage::Deletions(deleted) => { deleted .lock() .expect("deleted message keys mutex poisoned") .push((channel_id, bucket, message_id)); Ok(()) } #[cfg(test)] MessagesStorage::Memory(_) | MessagesStorage::Expiries(_) => Ok(()), } } async fn fetch_reactions_for_message_batch( &self, channel_id: i64, bucket: i32, message_ids: Vec, viewer_user_id: i64, ) -> HashMap> { match self { MessagesStorage::Postgres(storage) => { storage .fetch_reactions_for_message_batch( channel_id, bucket, message_ids, viewer_user_id, ) .await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage .fetch_reactions_for_message_batch( channel_id, bucket, message_ids, viewer_user_id, ) .await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Memory(_) | MessagesStorage::Expiries(_) => HashMap::new(), } } async fn fetch_attachment_decay_batch( &self, attachment_ids: Vec, ) -> anyhow::Result>> { match self { MessagesStorage::Postgres(storage) => { storage.fetch_attachment_decay_batch(attachment_ids).await } #[cfg(feature = "scylla")] MessagesStorage::Scylla(storage) => { storage.fetch_attachment_decay_batch(attachment_ids).await } #[cfg(test)] MessagesStorage::Deletions(_) | MessagesStorage::Memory(_) => Ok(HashMap::new()), #[cfg(test)] MessagesStorage::Expiries(expiries) => Ok(attachment_ids .into_iter() .filter_map(|id| expiries.get(&id).map(|expires_at| (id, *expires_at))) .collect()), } } } #[derive(Clone, Copy)] enum BucketBound { Before(i64), After(i64), } impl PostgresMessagesStorage { async fn get_by_id( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result> { let key = message_row_key(channel_id, bucket, message_id)?; let Some(row) = self.kv.get_row("messages", &key).await? else { return Ok(None); }; decode_postgres_message(row).map(Some) } async fn list_buckets( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, desc: bool, ) -> anyhow::Result> { if min_bucket > max_bucket || limit == 0 { return Ok(Vec::new()); } let partition_key = postgres::kv_key(&[KeyPart::BigInt(channel_id)])?; let mut buckets = self .kv .get_partition_rows("channel_message_buckets", &partition_key) .await? .into_iter() .map(|(_, row)| decode_postgres_bucket(row)) .collect::>>()?; buckets.retain(|bucket| *bucket >= min_bucket && *bucket <= max_bucket); if desc { buckets.sort_unstable_by(|left, right| right.cmp(left)); } else { buckets.sort_unstable(); } buckets.truncate(limit.min(i32::MAX as u32) as usize); Ok(buckets) } async fn fetch_bucket( &self, channel_id: i64, bucket: i32, bound: Option, desc: bool, limit: i32, ) -> anyhow::Result> { if limit <= 0 { return Ok(Vec::new()); } let partition_key = postgres::kv_key(&[KeyPart::BigInt(channel_id), KeyPart::Number(bucket as i64)])?; let bound = match bound { Some(BucketBound::Before(message_id)) => Some(BigIntBound::LessThan(message_id)), Some(BucketBound::After(message_id)) => Some(BigIntBound::GreaterThan(message_id)), None => None, }; self.kv .get_partition_rows_by_bigint_field( "messages", &partition_key, "message_id", bound, desc, limit as i64, ) .await? .into_iter() .map(|(_, row)| decode_postgres_message(row)) .collect::>>() } async fn delete_message( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result<()> { let key = message_row_key(channel_id, bucket, message_id)?; self.kv.delete_row("messages", &key).await } async fn fetch_reactions_for_message_batch( &self, channel_id: i64, bucket: i32, message_ids: Vec, viewer_user_id: i64, ) -> HashMap> { let wanted = message_ids.into_iter().collect::>(); if wanted.is_empty() { return HashMap::new(); } let partition_key = match postgres::kv_key(&[ KeyPart::BigInt(channel_id), KeyPart::Number(bucket as i64), ]) { Ok(key) => key, Err(_) => return HashMap::new(), }; let wanted_values = wanted.iter().copied().collect::>(); let rows = match self .kv .get_partition_rows_by_bigint_field_values( "message_reactions", &partition_key, "message_id", &wanted_values, ) .await { Ok(rows) => rows, Err(_) => return HashMap::new(), }; let mut by_message: HashMap> = HashMap::new(); for (_, row) in rows { let Ok((message_id, reaction)) = decode_postgres_reaction(row) else { continue; }; if wanted.contains(&message_id) { by_message.entry(message_id).or_default().push(reaction); } } by_message .into_iter() .filter_map(|(message_id, reactions)| { let reactions = map_reactions(reactions, viewer_user_id); (!reactions.is_empty()).then_some((message_id, reactions)) }) .collect() } async fn fetch_attachment_decay_batch( &self, attachment_ids: Vec, ) -> anyhow::Result>> { if attachment_ids.is_empty() { return Ok(HashMap::new()); } let keys = attachment_ids .iter() .map(|attachment_id| postgres::kv_key(&[KeyPart::BigInt(*attachment_id)])) .collect::>>()?; let rows = self.kv.get_rows("attachment_decay_by_id", &keys).await?; rows.into_iter() .map(|(_, row)| decode_postgres_attachment_decay(row)) .collect::>>() } } #[cfg(feature = "scylla")] impl ScyllaMessagesStorage { async fn get_by_id( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result> { let result = self .db .execute_unpaged(&self.stmt_get_by_id, (channel_id, bucket, message_id)) .await?; let rows = result.into_rows_result()?; Ok(rows.maybe_first_row::()?.map(Into::into)) } async fn list_buckets_desc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { let result = self .db .execute_unpaged( &self.stmt_list_buckets_desc, ( channel_id, min_bucket, max_bucket, limit.min(i32::MAX as u32) as i32, ), ) .await?; bucket_rows_to_vec(result.into_rows_result()?) } async fn list_buckets_asc( &self, channel_id: i64, min_bucket: i32, max_bucket: i32, limit: u32, ) -> anyhow::Result> { let result = self .db .execute_unpaged( &self.stmt_list_buckets_asc, ( channel_id, min_bucket, max_bucket, limit.min(i32::MAX as u32) as i32, ), ) .await?; bucket_rows_to_vec(result.into_rows_result()?) } async fn fetch_latest_bucket( &self, channel_id: i64, bucket: i32, limit: i32, ) -> anyhow::Result> { let result = self .db .execute_unpaged(&self.stmt_get_latest, (channel_id, bucket, limit)) .await?; rows_to_messages(result.into_rows_result()?) } async fn fetch_before_bucket( &self, channel_id: i64, bucket: i32, before_id: i64, limit: i32, ) -> anyhow::Result> { let result = self .db .execute_unpaged( &self.stmt_get_before, (channel_id, bucket, before_id, limit), ) .await?; rows_to_messages(result.into_rows_result()?) } async fn fetch_after_bucket( &self, channel_id: i64, bucket: i32, after_id: i64, limit: i32, ) -> anyhow::Result> { let result = self .db .execute_unpaged(&self.stmt_get_after, (channel_id, bucket, after_id, limit)) .await?; rows_to_messages(result.into_rows_result()?) } async fn delete_message( &self, channel_id: i64, bucket: i32, message_id: i64, ) -> anyhow::Result<()> { self.db .execute_unpaged(&self.stmt_delete_message, (channel_id, bucket, message_id)) .await?; Ok(()) } async fn fetch_reactions_for_message_batch( &self, channel_id: i64, bucket: i32, message_ids: Vec, viewer_user_id: i64, ) -> HashMap> { let result = self .db .execute_unpaged( &self.stmt_get_reactions_for_messages, (channel_id, bucket, message_ids.clone()), ) .await; let Ok(result) = result else { return self .fetch_reactions_for_messages_individually( channel_id, bucket, message_ids, viewer_user_id, ) .await; }; let Ok(rows) = result.into_rows_result() else { return self .fetch_reactions_for_messages_individually( channel_id, bucket, message_ids, viewer_user_id, ) .await; }; let Ok(rows) = rows.rows::() else { return HashMap::new(); }; let mut by_message: HashMap> = HashMap::new(); for row in rows.filter_map(Result::ok) { by_message .entry(row.message_id) .or_default() .push(MessageReactionDbRow { user_id: row.user_id, emoji_id: row.emoji_id, emoji_name: row.emoji_name, emoji_animated: row.emoji_animated, created_at: row.created_at, }); } by_message .into_iter() .filter_map(|(message_id, reactions)| { let reactions = map_reactions(reactions, viewer_user_id); (!reactions.is_empty()).then_some((message_id, reactions)) }) .collect() } async fn fetch_reactions_for_messages_individually( &self, channel_id: i64, bucket: i32, message_ids: Vec, viewer_user_id: i64, ) -> HashMap> { stream::iter(message_ids) .map(|message_id| async move { let result = self .db .execute_unpaged(&self.stmt_get_reactions, (channel_id, bucket, message_id)) .await .ok()?; let rows = result.into_rows_result().ok()?; let reactions: Vec = rows .rows::() .ok()? .filter_map(Result::ok) .collect(); let reactions = map_reactions(reactions, viewer_user_id); (!reactions.is_empty()).then_some((message_id, reactions)) }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .collect::>() .await .into_iter() .flatten() .collect() } async fn fetch_attachment_decay_batch( &self, attachment_ids: Vec, ) -> anyhow::Result>> { let result = self .db .execute_unpaged( &self.stmt_get_attachment_decay_many, (attachment_ids.clone(),), ) .await; let Ok(result) = result else { return self .fetch_attachment_decay_individually(attachment_ids) .await; }; let Ok(rows) = result.into_rows_result() else { return self .fetch_attachment_decay_individually(attachment_ids) .await; }; let rows = rows.rows::()?; rows.map(|row| { row.map(|row| (row.attachment_id, row.expires_at)) .map_err(Into::into) }) .collect::>>() } async fn fetch_attachment_decay_individually( &self, attachment_ids: Vec, ) -> anyhow::Result>> { let rows = stream::iter(attachment_ids) .map(|attachment_id| async move { let result = self .db .execute_unpaged(&self.stmt_get_attachment_decay, (attachment_id,)) .await?; let rows = result.into_rows_result()?; let row = rows.maybe_first_row::()?; Ok::<_, anyhow::Error>(row.map(|row| (row.attachment_id, row.expires_at))) }) .buffer_unordered(ENRICHMENT_QUERY_CONCURRENCY) .collect::>() .await .into_iter() .collect::>>()?; Ok(rows.into_iter().flatten().collect()) } } impl ShardService for MessagesShard { type Request = MessageRequest; type Response = MessageResponse; fn service_name(&self) -> &str { "messages" } async fn handle(&self, request: MessageRequest) -> anyhow::Result { match request { MessageRequest::GetResponseById { channel_id, message_id, viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret_base64, include_reactions, attachment_decay, nonce, tts, include_hidden, threads_mask, } => { let channel_id = parse_i64(&channel_id, "channel_id")?; let message_id = parse_i64(&message_id, "message_id")?; let viewer_user_id = parse_i64(&viewer_user_id, "viewer_user_id")?; let source_guild_id = source_guild_id .as_deref() .map(|id| parse_i64(id, "source_guild_id")) .transpose()?; let response = self .get_api_response( channel_id, message_id, ResponseBuildOptions { viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret: decode_attachment_url_secret( attachment_url_secret_base64.as_deref(), ), include_reactions: include_reactions.unwrap_or(true), attachment_decay: attachment_decay.unwrap_or(true), nonce, tts: tts.unwrap_or(false), include_hidden, threads_mask, excluded_types: Vec::new(), }, ) .await?; match response { Some(response) => Ok(MessageResponse::FoundApi(response)), None => Ok(MessageResponse::NotFound), } } MessageRequest::BuildResponse { message, viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret_base64, include_reactions, attachment_decay, nonce, tts, include_hidden, threads_mask, } => { let viewer_user_id = parse_i64(&viewer_user_id, "viewer_user_id")?; let source_guild_id = source_guild_id .as_deref() .map(|id| parse_i64(id, "source_guild_id")) .transpose()?; let response = self .build_api_response_from_message( message, ResponseBuildOptions { viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret: decode_attachment_url_secret( attachment_url_secret_base64.as_deref(), ), include_reactions: include_reactions.unwrap_or(true), attachment_decay: attachment_decay.unwrap_or(true), nonce, tts: tts.unwrap_or(false), include_hidden, threads_mask, excluded_types: Vec::new(), }, ) .await?; match response { Some(response) => Ok(MessageResponse::FoundApi(response)), None => Ok(MessageResponse::NotFound), } } MessageRequest::BuildResponses { messages, viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret_base64, include_reactions, attachment_decay, include_hidden, threads_mask, } => { let viewer_user_id = parse_i64(&viewer_user_id, "viewer_user_id")?; let source_guild_id = source_guild_id .as_deref() .map(|id| parse_i64(id, "source_guild_id")) .transpose()?; let responses = self .build_api_responses_from_messages( messages, ResponseBuildOptions { viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret: decode_attachment_url_secret( attachment_url_secret_base64.as_deref(), ), include_reactions: include_reactions.unwrap_or(true), attachment_decay: attachment_decay.unwrap_or(true), nonce: None, tts: false, include_hidden, threads_mask, excluded_types: Vec::new(), }, ) .await?; Ok(MessageResponse::FoundApiMany(responses)) } MessageRequest::ListResponses { channel_id, viewer_user_id, limit, before_id, after_id, around_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret_base64, include_reactions, attachment_decay, include_hidden, threads_mask, mut exclude_types, } => { let channel_id = parse_i64(&channel_id, "channel_id")?; if threads_mask && !exclude_types.contains(&MESSAGE_TYPE_THREAD_CREATED) { exclude_types.push(MESSAGE_TYPE_THREAD_CREATED); } let viewer_user_id = parse_i64(&viewer_user_id, "viewer_user_id")?; let before_id = before_id .as_deref() .map(|id| parse_i64(id, "before_id")) .transpose()?; let after_id = after_id .as_deref() .map(|id| parse_i64(id, "after_id")) .transpose()?; let around_id = around_id .as_deref() .map(|id| parse_i64(id, "around_id")) .transpose()?; let source_guild_id = source_guild_id .as_deref() .map(|id| parse_i64(id, "source_guild_id")) .transpose()?; let responses = self .list_api_responses( channel_id, limit, before_id, after_id, around_id, ResponseBuildOptions { viewer_user_id, source_guild_id, message_history_cutoff_ms, can_read_message_history, media_endpoint, media_proxy_secret_key, attachment_url_secret: decode_attachment_url_secret( attachment_url_secret_base64.as_deref(), ), include_reactions: include_reactions.unwrap_or(true), attachment_decay: attachment_decay.unwrap_or(true), nonce: None, tts: false, include_hidden, threads_mask, excluded_types: exclude_types, }, ) .await?; Ok(MessageResponse::FoundApiMany(responses)) } MessageRequest::ExtractMentions { contents } => { let mentions = contents .iter() .map(|content| { crate::types::ExtractedMentionsResponse::from( extract_mentions_from_markdown(Some(content)), ) }) .collect(); Ok(MessageResponse::FoundMentions(mentions)) } } } } fn message_row_key(channel_id: i64, bucket: i32, message_id: i64) -> anyhow::Result { postgres::kv_key(&[ KeyPart::BigInt(channel_id), KeyPart::Number(bucket as i64), KeyPart::BigInt(message_id), ]) } fn decode_postgres_bucket(row: serde_json::Value) -> anyhow::Result { let row = postgres::decode_row_dates_as_millis(row)?; let row: ChannelMessageBucketKvRow = serde_json::from_value(row)?; Ok(row.bucket) } fn decode_postgres_message(row: serde_json::Value) -> anyhow::Result { let row = postgres::decode_row_dates_as_millis(row)?; let serde_json::Value::Object(mut row) = row else { anyhow::bail!("Postgres message row payload is not an object"); }; let pinned = row .get("pinned_timestamp") .is_some_and(|value| !value.is_null()); row.insert("pinned".to_owned(), serde_json::Value::Bool(pinned)); Ok(serde_json::from_value(serde_json::Value::Object(row))?) } fn decode_postgres_reaction(row: serde_json::Value) -> anyhow::Result<(i64, MessageReactionDbRow)> { let row = postgres::decode_row_dates_as_millis(row)?; let row: MessageReactionKvRow = serde_json::from_value(row)?; let created_at = row .created_at .and_then(DateTime::::from_timestamp_millis); Ok(( row.message_id, MessageReactionDbRow { user_id: row.user_id, emoji_id: row.emoji_id, emoji_name: row.emoji_name, emoji_animated: row.emoji_animated, created_at, }, )) } fn decode_postgres_attachment_decay( row: serde_json::Value, ) -> anyhow::Result<(i64, DateTime)> { let row = postgres::decode_row_dates_as_millis(row)?; let row: AttachmentDecayKvRow = serde_json::from_value(row)?; let expires_at = DateTime::::from_timestamp_millis(row.expires_at) .ok_or_else(|| anyhow::anyhow!("invalid attachment decay timestamp"))?; Ok((row.attachment_id, expires_at)) } #[cfg(feature = "scylla")] fn rows_to_messages(rows: QueryRowsResult) -> anyhow::Result> { let rows: Vec = rows.rows::()?.collect::>()?; Ok(rows.into_iter().map(Message::from).collect()) } #[cfg(feature = "scylla")] fn bucket_rows_to_vec(rows: QueryRowsResult) -> anyhow::Result> { let rows: Vec = rows .rows::()? .collect::>()?; Ok(rows.into_iter().map(|row| row.bucket).collect()) } fn bucket_page_limit(message_limit: u32) -> u32 { message_limit.clamp(32, BUCKET_INDEX_PAGE_SIZE) } fn next_bucket_wave(wave: usize) -> usize { wave.saturating_mul(2).min(BUCKET_SCAN_CONCURRENCY) } async fn collect_bucket_waves( buckets: &[i32], limit: usize, collected: &mut Vec, fetch: Fetch, ) -> anyhow::Result<()> where Fetch: Fn(i32) -> Fut, Fut: std::future::Future>>, { let mut offset = 0; let mut wave = BUCKET_SCAN_WAVE; while offset < buckets.len() && collected.len() < limit { let end = buckets.len().min(offset + wave); let results = stream::iter(buckets[offset..end].iter().copied()) .map(&fetch) .buffer_unordered(wave) .collect::>() .await; for result in results { collected.extend(result?); } offset = end; wave = next_bucket_wave(wave); } Ok(()) } async fn collect_excluding( limit: u32, excluded_types: &[i32], mut cursor: Option, fetch: Fetch, ) -> anyhow::Result> where Fetch: Fn(Option, u32) -> Fut, Fut: std::future::Future>>, { let mut collected = Vec::new(); let mut scanned = 0_usize; let mut page_size = limit; while collected.len() < limit as usize && scanned < EXCLUDED_TYPE_SCAN_ROW_CAP { let page = fetch(cursor, page_size).await?; let exhausted = page.len() < page_size as usize; let Some(last) = page.last() else { break; }; cursor = Some(last.message_id); scanned += page.len(); collected.extend( page.into_iter() .filter(|message| !excluded_types.contains(&message.message_type)), ); if exhausted { break; } page_size = limit.max(EXCLUDED_TYPE_REFILL_PAGE_SIZE); } collected.truncate(limit as usize); Ok(collected) } fn around_window_limits(limit: u32) -> (u32, u32) { let newer_limit = limit / 2; let older_limit = limit.saturating_sub(1).saturating_sub(newer_limit); (newer_limit, older_limit) } fn parse_i64(value: &str, field_name: &str) -> anyhow::Result { value .parse::() .map_err(|error| anyhow::anyhow!("invalid {field_name}: {error}")) } fn snowflake_to_epoch_millis(snowflake: i64) -> i64 { (snowflake >> 22) + FLUXER_EPOCH_MS } fn epoch_millis_to_bucket(epoch_millis: i64) -> i32 { ((epoch_millis - FLUXER_EPOCH_MS) / BUCKET_DURATION_MS) as i32 } fn snowflake_to_bucket(snowflake: i64) -> i32 { ((snowflake >> 22) / BUCKET_DURATION_MS) as i32 } fn current_bucket() -> i32 { epoch_millis_to_bucket(unix_millis()) } fn epoch_millis_to_iso(epoch_millis: i64) -> String { DateTime::::from_timestamp_millis(epoch_millis) .unwrap_or(DateTime::::UNIX_EPOCH) .to_rfc3339_opts(chrono::SecondsFormat::Millis, true) } fn map_embed_field_response(field: MessageEmbedField) -> ApiEmbedFieldResponse { ApiEmbedFieldResponse { name: field.name.unwrap_or_default(), value: field.value.unwrap_or_default(), is_inline: field.is_inline.unwrap_or(false), } } fn build_message_mention_context(messages: &[&Message]) -> HashMap { messages .iter() .map(|message| (message.message_id, build_mention_context_entry(message))) .collect() } fn build_mention_context_entry(message: &Message) -> MessageMentionContext { let message_snapshots = message.message_snapshots.as_deref().unwrap_or_default(); let snapshots = message_snapshots .iter() .map(|snapshot| extract_mentions_from_markdown(snapshot.content.as_deref())) .collect(); let mut embed_users = HashSet::new(); for embed in message.embeds.as_deref().unwrap_or_default() { collect_user_ids_from_embed(embed, &mut embed_users); } for snapshot in message_snapshots { for embed in snapshot.embeds.as_deref().unwrap_or_default() { collect_user_ids_from_embed(embed, &mut embed_users); } } MessageMentionContext { content: extract_mentions_from_markdown(message.content.as_deref()), snapshots, embed_users, } } fn ids_present_in_set(ids: &[i64], present: &HashSet) -> Vec { ids.iter() .filter(|id| present.contains(id)) .map(ToString::to_string) .collect() } fn collect_attachment_ids(messages: &[&Message]) -> HashSet { let mut ids = HashSet::new(); for message in messages { if let Some(attachments) = &message.attachments { for attachment in attachments { if let Some(id) = attachment.attachment_id { ids.insert(id); } } } if let Some(snapshots) = &message.message_snapshots { for snapshot in snapshots { if let Some(attachments) = &snapshot.attachments { for attachment in attachments { if let Some(id) = attachment.attachment_id { ids.insert(id); } } } } } } ids } fn collect_channel_mention_ids( messages: &[&Message], mention_context: &HashMap, ) -> HashSet { let mut ids = HashSet::new(); for message in messages { let content_mentions = mention_context .get(&message.message_id) .map(|mentions| &mentions.content); for id in &message.mention_channels { if content_mentions.is_some_and(|mentions| mentions.channels.contains(id)) { ids.insert(*id); } } if let Some(snapshots) = &message.message_snapshots { for (index, snapshot) in snapshots.iter().enumerate() { let snapshot_mentions = mention_context .get(&message.message_id) .and_then(|mentions| mentions.snapshots.get(index)); for id in &snapshot.mention_channels { if snapshot_mentions.is_some_and(|mentions| mentions.channels.contains(id)) { ids.insert(*id); } } } } } ids } fn collect_user_ids( messages: &[&Message], mention_context: &HashMap, ) -> HashSet { let mut ids = HashSet::new(); for message in messages { if let Some(author_id) = message.author_id { ids.insert(author_id); } for id in &message.mention_users { ids.insert(*id); } if let Some(mentions) = mention_context.get(&message.message_id) { ids.extend(mentions.content.users.iter().copied()); ids.extend(mentions.embed_users.iter().copied()); } if let Some(snapshots) = &message.message_snapshots { for (index, snapshot) in snapshots.iter().enumerate() { for id in &snapshot.mention_users { ids.insert(*id); } if let Some(mentions) = mention_context .get(&message.message_id) .and_then(|mentions| mentions.snapshots.get(index)) { ids.extend(mentions.users.iter().copied()); } } } } ids } fn collect_user_ids_from_embed(embed: &MessageEmbed, target: &mut HashSet) { let mut mentions = MessageMentions::default(); extend_mentions_from_markdown(embed.title.as_deref(), &mut mentions); extend_mentions_from_markdown(embed.description.as_deref(), &mut mentions); if let Some(author) = &embed.author { extend_mentions_from_markdown(author.name.as_deref(), &mut mentions); } if let Some(footer) = &embed.footer { extend_mentions_from_markdown(footer.text.as_deref(), &mut mentions); } for media in [&embed.image, &embed.thumbnail, &embed.video, &embed.audio] .into_iter() .flatten() { extend_mentions_from_markdown(media.description.as_deref(), &mut mentions); } if let Some(fields) = &embed.fields { for field in fields { extend_mentions_from_markdown(field.name.as_deref(), &mut mentions); extend_mentions_from_markdown(field.value.as_deref(), &mut mentions); } } target.extend(mentions.users); if let Some(children) = &embed.children { for child in children { collect_user_ids_from_embed_child(child, target); } } } fn collect_user_ids_from_embed_child(embed: &MessageEmbedChild, target: &mut HashSet) { let mut mentions = MessageMentions::default(); extend_mentions_from_markdown(embed.title.as_deref(), &mut mentions); extend_mentions_from_markdown(embed.description.as_deref(), &mut mentions); if let Some(author) = &embed.author { extend_mentions_from_markdown(author.name.as_deref(), &mut mentions); } if let Some(footer) = &embed.footer { extend_mentions_from_markdown(footer.text.as_deref(), &mut mentions); } for media in [&embed.image, &embed.thumbnail, &embed.video, &embed.audio] .into_iter() .flatten() { extend_mentions_from_markdown(media.description.as_deref(), &mut mentions); } if let Some(fields) = &embed.fields { for field in fields { extend_mentions_from_markdown(field.name.as_deref(), &mut mentions); extend_mentions_from_markdown(field.value.as_deref(), &mut mentions); } } target.extend(mentions.users); } fn map_user_partial(partial: UserPartialServiceResponse) -> ApiUserPartial { ApiUserPartial::from(UserPartialFields { user_id: partial.user_id, username: partial.username, discriminator: partial.discriminator, global_name: partial.global_name, avatar_hash: partial.avatar_hash, avatar_color: partial.avatar_color, bot: partial.bot, system: partial.system, flags: partial.flags, mention_flags: partial.mention_flags, }) } fn deleted_user(user_id: i64) -> ApiUserPartial { if user_id == FLUXER_SYSTEM_USER_ID { return fluxer_system_user(); } ApiUserPartial { id: user_id.to_string(), username: DELETED_USER_USERNAME.to_owned(), discriminator: "0000".to_owned(), global_name: Some(DELETED_USER_GLOBAL_NAME.to_owned()), avatar: None, avatar_color: None, bot: None, system: None, flags: 0, mention_flags: None, } } fn map_reactions( reactions: Vec, viewer_user_id: i64, ) -> Vec { #[derive(Clone)] struct Group { emoji_id: Option, emoji_name: String, animated: bool, count: i32, me: bool, min_created_at: i64, } let mut groups: HashMap = HashMap::new(); for reaction in reactions { let is_custom = reaction.emoji_id != 0; let key = if is_custom { format!("custom_{}", reaction.emoji_id) } else { format!("unicode_{}", reaction.emoji_name) }; let created_at = reaction .created_at .map(|dt| dt.timestamp_millis()) .unwrap_or_default(); groups .entry(key) .and_modify(|group| { group.count += 1; group.me = group.me || reaction.user_id == viewer_user_id; if created_at < group.min_created_at { group.min_created_at = created_at; } }) .or_insert_with(|| Group { emoji_id: is_custom.then_some(reaction.emoji_id), emoji_name: reaction.emoji_name, animated: reaction.emoji_animated.unwrap_or(false), count: 1, me: reaction.user_id == viewer_user_id, min_created_at: created_at, }); } let mut groups: Vec = groups.into_values().collect(); groups.sort_by(|left, right| { left.min_created_at .cmp(&right.min_created_at) .then_with(|| left.emoji_name.cmp(&right.emoji_name)) .then_with(|| left.emoji_id.cmp(&right.emoji_id)) }); groups .into_iter() .map(|group| ApiMessageReactionResponse { emoji: ApiReactionEmojiResponse { id: group.emoji_id.map(|id| id.to_string()), name: group.emoji_name, animated: group.animated.then_some(true), }, count: group.count, me: group.me.then_some(true), }) .collect() } fn map_message_reference(reference: &MessageReference) -> Option { Some(ApiMessageReferenceResponse { channel_id: reference.channel_id?.to_string(), message_id: reference.message_id.map(|id| id.to_string()), guild_id: reference.guild_id.map(|id| id.to_string()), reference_type: effective_reference_type(reference), }) } fn map_sticker(sticker: &MessageStickerItem) -> Option { Some(ApiMessageStickerResponse { id: sticker.sticker_id?.to_string(), name: sticker.name.clone().unwrap_or_default(), animated: sticker.animated.unwrap_or(false), }) } fn map_call(call: &MessageCall) -> ApiMessageCallResponse { ApiMessageCallResponse { participants: call .participant_ids .iter() .map(ToString::to_string) .collect(), ended_timestamp: call.ended_timestamp.map(epoch_millis_to_iso), } } fn make_attachment_cdn_key(channel_id: i64, attachment_id: i64, filename: &str) -> String { format!("attachments/{channel_id}/{attachment_id}/{filename}") } fn make_attachment_cdn_url( media_endpoint: &str, channel_id: i64, attachment_id: i64, filename: &str, ) -> String { format!( "{}/{}", media_endpoint.trim_end_matches('/'), make_attachment_cdn_key(channel_id, attachment_id, filename) ) } fn decode_attachment_url_secret(encoded: Option<&str>) -> Vec { let Some(encoded) = encoded.filter(|value| !value.is_empty()) else { return Vec::new(); }; match BASE64_STANDARD.decode(encoded) { Ok(secret) => secret, Err(error) => { tracing::warn!(%error, "attachment url secret is not standard base64"); Vec::new() } } } fn attachment_anchor_secs(attachment_id: i64) -> u64 { snowflake_to_epoch_millis(attachment_id).max(0) as u64 / 1_000 } fn sign_attachment_url( url: &str, storage_key: &str, attachment_id: i64, now_secs: u64, secret: &[u8], ) -> String { if secret.is_empty() { return url.to_owned(); } fluxer_common::attachment_url_signature::with_signature( url, storage_key, attachment_anchor_secs(attachment_id), now_secs, secret, ) } fn attachment_cdn_url( channel_id: i64, attachment_id: i64, filename: &str, options: &ResponseBuildOptions, now_secs: u64, ) -> String { let storage_key = make_attachment_cdn_key(channel_id, attachment_id, filename); let url = make_attachment_cdn_url(&options.media_endpoint, channel_id, attachment_id, filename); sign_attachment_url( &url, &storage_key, attachment_id, now_secs, &options.attachment_url_secret, ) } fn raw_url_path(url: &str) -> &str { let Some((_, after_scheme)) = url.split_once("://") else { return ""; }; let Some(start) = after_scheme.find(['/', '?', '#']) else { return ""; }; let path = &after_scheme[start..]; if !path.starts_with('/') { return ""; } path.find(['?', '#']).map_or(path, |end| &path[..end]) } fn own_attachment_key(input_url: &str, media_endpoint: &str) -> Option<(String, i64)> { let endpoint_path = raw_url_path(media_endpoint).trim_end_matches('/'); let path = raw_url_path(input_url).strip_prefix(endpoint_path)?; if !path.starts_with("/attachments/") { return None; } let storage_key = fluxer_common::attachment_url_signature::decode_key(path)?; let attachment_id = attachment_key_id(&storage_key)?; Some((storage_key, attachment_id)) } fn attachment_key_id(storage_key: &str) -> Option { let segments: Vec<&str> = storage_key.split('/').collect(); let ["attachments", channel_id, attachment_id, filename @ ..] = segments.as_slice() else { return None; }; if !is_snowflake_segment(channel_id) || !is_snowflake_segment(attachment_id) || filename.is_empty() || filename .iter() .any(|part| matches!(*part, "" | "." | "..") || part.contains('\0')) { return None; } attachment_id.parse().ok() } fn is_snowflake_segment(segment: &str) -> bool { (1..=20).contains(&segment.len()) && segment.bytes().all(|byte| byte.is_ascii_digit()) } fn sign_own_media_url( input_url: &str, media_endpoint: &str, now_secs: u64, secret: &[u8], ) -> String { let Some((storage_key, attachment_id)) = own_attachment_key(input_url, media_endpoint) else { return input_url.to_owned(); }; sign_attachment_url(input_url, &storage_key, attachment_id, now_secs, secret) } fn sign_own_url(input_url: &str, options: &ResponseBuildOptions) -> String { sign_own_media_url( input_url, options.media_endpoint.trim_end_matches('/'), unix_seconds(), &options.attachment_url_secret, ) } fn external_media_proxy_url(input_url: &str, options: &ResponseBuildOptions) -> String { media_proxy_url_at(input_url, options, unix_seconds()) } fn media_proxy_url_at(input_url: &str, options: &ResponseBuildOptions, now_secs: u64) -> String { let media_endpoint = options.media_endpoint.trim_end_matches('/'); let parsed_url = match Url::parse(input_url) { Ok(url) => url, Err(_) => return input_url.to_owned(), }; if fluxer_common::external_media_path::is_own_media_url(media_endpoint, &parsed_url) { return sign_own_media_url( input_url, media_endpoint, now_secs, &options.attachment_url_secret, ); } if options.media_proxy_secret_key.is_empty() { return input_url.to_owned(); } fluxer_common::external_media_path::build_external_media_proxy_url( media_endpoint, parsed_url.as_str(), options.media_proxy_secret_key.as_bytes(), ) .unwrap_or_else(|| input_url.to_owned()) } fn dt_to_epoch_millis(dt: &DateTime) -> i64 { dt.timestamp_millis() } fn convert_attachment(a: udt::AttachmentUdt) -> MessageAttachment { MessageAttachment { attachment_id: a.attachment_id, filename: a.filename, size: a.size, title: a.title, description: a.description, width: a.width, height: a.height, duration: a.duration, content_type: a.content_type, content_hash: a.content_hash, placeholder: a.placeholder, flags: a.flags, nsfw: a.nsfw, duration_secs: a.duration_secs, waveform: a.waveform, } } fn convert_embed_author(a: udt::EmbedAuthorUdt) -> MessageEmbedAuthor { MessageEmbedAuthor { name: a.name, url: a.url, icon_url: a.icon_url, } } fn convert_embed_field(f: udt::EmbedFieldUdt) -> MessageEmbedField { MessageEmbedField { name: f.name, value: f.value, is_inline: f.is_inline, } } fn convert_embed_footer(f: udt::EmbedFooterUdt) -> MessageEmbedFooter { MessageEmbedFooter { text: f.text, icon_url: f.icon_url, } } fn convert_embed_media(m: udt::EmbedMediaUdt) -> MessageEmbedMedia { MessageEmbedMedia { url: m.url, width: m.width, height: m.height, duration: m.duration, description: m.description, content_type: m.content_type, content_hash: m.content_hash, placeholder: m.placeholder, flags: m.flags, } } fn convert_embed_provider(p: udt::EmbedProviderUdt) -> MessageEmbedProvider { MessageEmbedProvider { name: p.name, url: p.url, } } fn convert_embed_child(c: udt::EmbedChildUdt) -> MessageEmbedChild { MessageEmbedChild { embed_type: c.embed_type, title: c.title, description: c.description, url: c.url, timestamp: c.timestamp.as_ref().map(dt_to_epoch_millis), color: c.color, author: c.author.map(convert_embed_author), provider: c.provider.map(convert_embed_provider), thumbnail: c.thumbnail.map(convert_embed_media), image: c.image.map(convert_embed_media), video: c.video.map(convert_embed_media), footer: c.footer.map(convert_embed_footer), fields: c .fields .map(|v| v.into_iter().map(convert_embed_field).collect()), nsfw: c.nsfw, audio: c.audio.map(convert_embed_media), html: c.html, html_width: c.html_width, html_height: c.html_height, } } fn convert_embed(e: udt::EmbedUdt) -> MessageEmbed { MessageEmbed { embed_type: e.embed_type, title: e.title, description: e.description, url: e.url, timestamp: e.timestamp.as_ref().map(dt_to_epoch_millis), color: e.color, author: e.author.map(convert_embed_author), provider: e.provider.map(convert_embed_provider), thumbnail: e.thumbnail.map(convert_embed_media), image: e.image.map(convert_embed_media), video: e.video.map(convert_embed_media), footer: e.footer.map(convert_embed_footer), fields: e .fields .map(|v| v.into_iter().map(convert_embed_field).collect()), nsfw: e.nsfw, children: e .children .map(|v| v.into_iter().map(convert_embed_child).collect()), audio: e.audio.map(convert_embed_media), html: e.html, html_width: e.html_width, html_height: e.html_height, } } fn convert_sticker_item(s: udt::StickerItemUdt) -> MessageStickerItem { MessageStickerItem { sticker_id: s.sticker_id, name: s.name, format_type: s.format_type, animated: s.animated, } } fn convert_message_reference(r: udt::MessageReferenceUdt) -> MessageReference { MessageReference { channel_id: r.channel_id, message_id: r.message_id, guild_id: r.guild_id, reference_type: r.reference_type, } } fn convert_message_call(c: udt::MessageCallUdt) -> MessageCall { MessageCall { participant_ids: c .participant_ids .map(|s| s.into_iter().collect()) .unwrap_or_default(), ended_timestamp: c.ended_timestamp.as_ref().map(dt_to_epoch_millis), } } fn convert_message_snapshot(s: udt::MessageSnapshotUdt) -> MessageSnapshot { let edited_ts = s .edited_timestamp .as_ref() .or(s.edited_timestmap.as_ref()) .map(dt_to_epoch_millis); MessageSnapshot { content: s.content, timestamp: s.timestamp.as_ref().map(dt_to_epoch_millis), edited_timestamp: edited_ts, mention_users: s .mention_users .map(|s| s.into_iter().collect()) .unwrap_or_default(), mention_roles: s .mention_roles .map(|s| s.into_iter().collect()) .unwrap_or_default(), mention_channels: s .mention_channels .map(|s| s.into_iter().collect()) .unwrap_or_default(), attachments: s .attachments .map(|v| v.into_iter().map(convert_attachment).collect()), embeds: s.embeds.map(|v| v.into_iter().map(convert_embed).collect()), sticker_items: s .sticker_items .map(|v| v.into_iter().map(convert_sticker_item).collect()), snapshot_type: s.snapshot_type, flags: s.flags, } } impl From for Message { fn from(row: MessageDbRow) -> Self { Self { message_id: row.message_id, channel_id: row.channel_id, bucket: row.bucket, author_id: row.author_id, message_type: row.r#type.unwrap_or_default(), webhook_id: row.webhook_id, webhook_name: row.webhook_name, webhook_avatar_hash: row.webhook_avatar_hash, content: row.content, edited_timestamp: row.edited_timestamp.as_ref().map(dt_to_epoch_millis), pinned: Some(row.pinned_timestamp.is_some()), flags: row.flags.map(i64::from), mention_everyone: row.mention_everyone, mention_users: row .mention_users .map(|s| s.into_iter().collect()) .unwrap_or_default(), mention_roles: row .mention_roles .map(|s| s.into_iter().collect()) .unwrap_or_default(), mention_channels: row .mention_channels .map(|s| s.into_iter().collect()) .unwrap_or_default(), has_reaction: row.has_reaction, version: row.version.unwrap_or_default(), attachments: row .attachments .map(|v| v.into_iter().map(convert_attachment).collect()), embeds: row .embeds .map(|v| v.into_iter().map(convert_embed).collect()), sticker_items: row .sticker_items .map(|v| v.into_iter().map(convert_sticker_item).collect()), message_reference: row.message_reference.map(convert_message_reference), call: row.call.map(convert_message_call), message_snapshots: row .message_snapshots .map(|v| v.into_iter().map(convert_message_snapshot).collect()), } } } #[cfg(test)] mod tests { use super::*; use fluxer_svc::transport::{ InMemoryTransport, TransportMessage, TransportSubscriber, reply_message, }; use serde_json::json; #[test] fn missing_id_zero_fallback_is_fluxer_system_user() { let mapped = deleted_user(0); assert_eq!(mapped.id, "0"); assert_eq!(mapped.username, "Fluxer"); assert_eq!(mapped.global_name, None); assert_eq!(mapped.bot, Some(true)); assert_eq!(mapped.system, Some(true)); } #[test] fn missing_regular_user_fallback_is_not_bot_or_system() { let mapped = deleted_user(42); assert_eq!(mapped.id, "42"); assert_eq!(mapped.username, DELETED_USER_USERNAME); assert_eq!( mapped.global_name, Some(DELETED_USER_GLOBAL_NAME.to_owned()) ); assert_eq!(mapped.bot, None); assert_eq!(mapped.system, None); } fn reaction( user_id: i64, emoji_id: i64, emoji_name: &str, created_at_ms: i64, ) -> MessageReactionDbRow { MessageReactionDbRow { user_id, emoji_id, emoji_name: emoji_name.to_owned(), emoji_animated: Some(false), created_at: DateTime::::from_timestamp_millis(created_at_ms), } } #[test] fn postgres_message_decoder_maps_tagged_kv_payload() { let message = decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "1497639278555484216"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": "1509197195776110592"}, "author_id": {"__fluxer_type": "bigint", "value": "1472426752046002208"}, "type": null, "content": "hello", "edited_timestamp": {"__fluxer_type": "date", "value": "2026-06-15T12:34:56.789Z"}, "pinned_timestamp": {"__fluxer_type": "date", "value": "2026-06-15T12:35:00.000Z"}, "mention_users": {"__fluxer_type": "set", "value": [ {"__fluxer_type": "bigint", "value": "1472426752046002208"} ]}, "mention_roles": {"__fluxer_type": "set", "value": []}, "mention_channels": {"__fluxer_type": "set", "value": []}, "attachments": [{ "attachment_id": {"__fluxer_type": "bigint", "value": "1509197195776110593"}, "filename": "a.png", "size": {"__fluxer_type": "bigint", "value": "12345"}, "content_type": "image/png", "flags": 0 }], "call": { "participant_ids": {"__fluxer_type": "set", "value": [ {"__fluxer_type": "bigint", "value": "1472426752046002208"} ]}, "ended_timestamp": {"__fluxer_type": "date", "value": "2026-06-15T12:36:00.000Z"} }, "version": null })) .unwrap(); assert_eq!(message.message_id, 1_509_197_195_776_110_592); assert_eq!(message.channel_id, 1_497_639_278_555_484_216); assert_eq!(message.author_id, Some(1_472_426_752_046_002_208)); assert_eq!(message.message_type, 0); assert_eq!(message.edited_timestamp, Some(1_781_526_896_789)); assert_eq!(message.pinned, Some(true)); assert_eq!(message.version, 0); assert_eq!(message.mention_users, vec![1_472_426_752_046_002_208]); assert_eq!( message.attachments.unwrap()[0].attachment_id, Some(1_509_197_195_776_110_593) ); assert_eq!( message.call.unwrap().participant_ids, vec![1_472_426_752_046_002_208] ); } #[test] fn build_responses_request_accepts_legacy_null_version_rows() { let request: MessageRequest = serde_json::from_value(json!({ "op": "BuildResponses", "messages": [{ "message_id": "1449544529132171273", "channel_id": "1431572375251247158", "bucket": 399, "author_id": "1130650140672000000", "type": null, "version": null, "content": "" }], "viewer_user_id": "1130650140672000000", "source_guild_id": null, "message_history_cutoff_ms": null, "can_read_message_history": true, "media_endpoint": "https://media.example", "media_proxy_secret_key": "secret", "include_reactions": true })) .unwrap(); let MessageRequest::BuildResponses { messages, .. } = request else { panic!("expected BuildResponses"); }; assert_eq!(messages[0].message_type, 0); assert_eq!(messages[0].version, 0); } #[test] fn mention_context_includes_embed_user_ids_for_message_and_snapshots() { let message: Message = serde_json::from_value(json!({ "message_id": "10", "channel_id": "20", "bucket": 1, "author_id": "30", "type": 0, "version": 0, "content": "hello <@40>", "mention_users": ["40"], "embeds": [{ "title": "title <@50>", "description": "description <@60>", "footer": {"text": "footer <@70>"}, "fields": [{"name": "field <@80>", "value": "value <@90>"}] }], "message_snapshots": [{ "content": "snapshot <@100>", "embeds": [{"description": "snapshot embed <@110>"}] }] })) .unwrap(); let message_ref = &message; let messages = std::slice::from_ref(&message_ref); let mention_context = build_message_mention_context(messages); let entry = mention_context.get(&10).unwrap(); assert_eq!(entry.content.users, HashSet::from([40])); assert_eq!(entry.embed_users, HashSet::from([50, 60, 70, 80, 90, 110])); assert_eq!(entry.snapshots[0].users, HashSet::from([100])); assert_eq!( collect_user_ids(messages, &mention_context), HashSet::from([30, 40, 50, 60, 70, 80, 90, 100, 110]) ); } #[test] fn postgres_reaction_decoder_maps_created_at() { let (message_id, reaction) = decode_postgres_reaction(json!({ "message_id": {"__fluxer_type": "bigint", "value": "1509197195776110592"}, "user_id": {"__fluxer_type": "bigint", "value": "1472426752046002208"}, "emoji_id": {"__fluxer_type": "bigint", "value": "0"}, "emoji_name": "wave", "emoji_animated": false, "created_at": {"__fluxer_type": "date", "value": "2026-06-15T12:34:56.789Z"} })) .unwrap(); assert_eq!(message_id, 1_509_197_195_776_110_592); assert_eq!(reaction.user_id, 1_472_426_752_046_002_208); assert_eq!(reaction.emoji_name, "wave"); assert_eq!( reaction.created_at.map(|dt| dt.timestamp_millis()), Some(1_781_526_896_789) ); } #[tokio::test] async fn bucket_scan_stops_once_the_newest_buckets_fill_the_page() { let buckets = (0..50).map(|index| 500 - index).collect::>(); let scanned = std::sync::Mutex::new(Vec::new()); let mut collected: Vec = Vec::new(); collect_bucket_waves(&buckets, 50, &mut collected, |bucket| { let scanned = &scanned; async move { scanned.lock().expect("scan log").push(bucket); let rows = (0..200) .map(|row| i64::from(bucket) * 1_000 + row) .collect::>(); Ok(rows) } }) .await .expect("bucket scan succeeds"); let mut scanned = scanned.into_inner().expect("scan log"); scanned.sort_unstable_by(|left, right| right.cmp(left)); assert_eq!(scanned, buckets[..BUCKET_SCAN_WAVE]); let newest_skipped_id = i64::from(buckets[BUCKET_SCAN_WAVE]) * 1_000 + 199; assert!(collected.iter().all(|id| *id > newest_skipped_id)); } #[tokio::test] async fn bucket_scan_walks_the_whole_page_when_buckets_are_sparse() { let buckets = (0..50).map(|index| 500 - index).collect::>(); let scanned = std::sync::Mutex::new(0_usize); let mut collected: Vec = Vec::new(); collect_bucket_waves(&buckets, 50, &mut collected, |bucket| { let scanned = &scanned; async move { *scanned.lock().expect("scan count") += 1; Ok(vec![i64::from(bucket)]) } }) .await .expect("bucket scan succeeds"); assert_eq!(scanned.into_inner().expect("scan count"), buckets.len()); assert_eq!(collected.len(), buckets.len()); } #[test] fn reference_with_no_stored_type_is_treated_as_a_reply_everywhere() { let legacy = MessageReference { channel_id: Some(1), message_id: Some(2), guild_id: None, reference_type: None, }; let explicit = MessageReference { reference_type: Some(MESSAGE_REFERENCE_TYPE_DEFAULT), ..legacy.clone() }; let forward = MessageReference { reference_type: Some(MESSAGE_REFERENCE_TYPE_DEFAULT + 1), ..legacy.clone() }; assert_eq!( effective_reference_type(&legacy), MESSAGE_REFERENCE_TYPE_DEFAULT ); assert_eq!( effective_reference_type(&legacy), effective_reference_type(&explicit) ); assert_ne!( effective_reference_type(&forward), MESSAGE_REFERENCE_TYPE_DEFAULT ); let mapped = map_message_reference(&legacy).expect("legacy reference maps"); assert_eq!(mapped.reference_type, MESSAGE_REFERENCE_TYPE_DEFAULT); } #[test] fn around_window_limits_match_reference_api() { assert_eq!(around_window_limits(0), (0, 0)); assert_eq!(around_window_limits(1), (0, 0)); assert_eq!(around_window_limits(2), (1, 0)); assert_eq!(around_window_limits(3), (1, 1)); assert_eq!(around_window_limits(50), (25, 24)); assert_eq!(around_window_limits(51), (25, 25)); } #[test] fn epoch_bucket_matches_snowflake_bucket() { let message_id = 1_509_197_195_776_110_592; assert_eq!( epoch_millis_to_bucket(snowflake_to_epoch_millis(message_id)), snowflake_to_bucket(message_id) ); } #[test] fn snowflake_bucket_matches_existing_channel_rows() { assert_eq!(snowflake_to_bucket(1_474_193_838_282_432_581), 406); assert_eq!(snowflake_to_bucket(1_488_946_116_942_273_651), 410); assert_eq!(snowflake_to_bucket(1_509_256_674_043_502_592), 416); } #[test] fn embed_field_response_defaults_null_name_and_value_to_empty_strings() { let response = map_embed_field_response(MessageEmbedField { name: None, value: None, is_inline: None, }); assert_eq!(response.name, ""); assert_eq!(response.value, ""); assert!(!response.is_inline); } #[test] fn embed_media_response_hides_the_owned_attachment_flag() { let shard = recording_shard(&DeletedMessageKeys::default()); let response = shard .map_embed_media( MessageEmbedMedia { url: Some("https://media.example.com/attachments/1/2/a.png".to_owned()), width: None, height: None, duration: None, description: None, content_type: None, content_hash: None, placeholder: None, flags: Some(EMBED_MEDIA_OWNED_ATTACHMENT_FLAG | (1 << 3)), }, &build_options(), ) .unwrap(); assert_eq!(response.flags, Some(1 << 3)); } #[test] fn reactions_are_grouped_sorted_and_viewer_aware() { let mapped = map_reactions( vec![ reaction(10, 0, "z", 300), reaction(11, 0, "a", 100), reaction(12, 0, "a", 200), reaction(11, 99, "party", 50), ], 12, ); assert_eq!(mapped.len(), 3); assert_eq!(mapped[0].emoji.id.as_deref(), Some("99")); assert_eq!(mapped[0].emoji.name, "party"); assert_eq!(mapped[0].count, 1); assert_eq!(mapped[1].emoji.name, "a"); assert_eq!(mapped[1].count, 2); assert_eq!(mapped[1].me, Some(true)); assert_eq!(mapped[2].emoji.name, "z"); } #[test] fn response_context_id_collection_spans_referenced_messages() { let message = decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": "100"}, "author_id": {"__fluxer_type": "bigint", "value": "1"}, "content": "hey <@2> over in <#20>", "mention_users": {"__fluxer_type": "set", "value": [ {"__fluxer_type": "bigint", "value": "2"} ]}, "mention_channels": {"__fluxer_type": "set", "value": [ {"__fluxer_type": "bigint", "value": "20"} ]}, "attachments": [{ "attachment_id": {"__fluxer_type": "bigint", "value": "1000"}, "filename": "a.png" }] })) .unwrap(); let referenced = decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": "99"}, "author_id": {"__fluxer_type": "bigint", "value": "3"}, "content": "look at <#21>", "mention_channels": {"__fluxer_type": "set", "value": [ {"__fluxer_type": "bigint", "value": "21"} ]}, "attachments": [{ "attachment_id": {"__fluxer_type": "bigint", "value": "1001"}, "filename": "b.png" }] })) .unwrap(); let messages = [message]; let referenced_messages: HashMap<(i64, i64), Message> = [((10, 99), referenced)].into_iter().collect(); let all_messages = messages .iter() .chain(referenced_messages.values()) .collect::>(); let mention_context = build_message_mention_context(&all_messages); assert_eq!( collect_attachment_ids(&all_messages), HashSet::from([1000, 1001]) ); assert_eq!( collect_channel_mention_ids(&all_messages, &mention_context), HashSet::from([20, 21]) ); assert_eq!( collect_user_ids(&all_messages, &mention_context), HashSet::from([1, 2, 3]) ); } fn recording_shard(deleted: &DeletedMessageKeys) -> MessagesShard { MessagesShard { storage: MessagesStorage::Deletions(deleted.clone()), transport: InMemoryTransport::new(), } } fn build_options() -> ResponseBuildOptions { ResponseBuildOptions { viewer_user_id: 1, source_guild_id: None, message_history_cutoff_ms: None, can_read_message_history: true, media_endpoint: "https://media.example.com".to_owned(), media_proxy_secret_key: "secret".to_owned(), attachment_url_secret: Vec::new(), include_reactions: false, nonce: None, tts: false, include_hidden: false, threads_mask: false, excluded_types: Vec::new(), attachment_decay: true, } } fn orphaned_message(message_id: i64) -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": message_id.to_string()}, "content": "orphan" })) .unwrap() } fn webhook_message(message_id: i64) -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": message_id.to_string()}, "webhook_id": {"__fluxer_type": "bigint", "value": "77"}, "webhook_name": "hook", "content": "kept" })) .unwrap() } fn authored_message(message_id: i64) -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": message_id.to_string()}, "author_id": {"__fluxer_type": "bigint", "value": "1472426752046002208"}, "content": "kept" })) .unwrap() } fn legacy_string_author_message(message_id: i64) -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": message_id.to_string()}, "author_id": "1472426752046002208", "content": "kept" })) .unwrap() } async fn stub_user_service(transport: &InMemoryTransport) -> tokio::task::JoinHandle<()> { let mut subscriber = transport.subscribe("svc.users").await.unwrap(); let transport = transport.clone(); tokio::spawn(async move { while let Some(message) = subscriber.next().await { let _ = reply_message(&message, &transport, b"\"NotFound\"").await; } }) } fn recorded_deletions(deleted: &DeletedMessageKeys) -> Vec<(i64, i32, i64)> { deleted.lock().unwrap().clone() } #[tokio::test] async fn build_path_keeps_authored_rows_from_older_releases() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let users = stub_user_service(&shard.transport).await; let wrapped_id = 1_509_197_195_776_110_592; let legacy_id = 1_509_197_195_776_110_593; let responses = shard .build_api_responses_from_messages( vec![ authored_message(wrapped_id), legacy_string_author_message(legacy_id), ], build_options(), ) .await .unwrap(); users.abort(); assert_eq!(responses.len(), 2); assert_eq!(responses[0].id, wrapped_id.to_string()); assert_eq!(responses[1].id, legacy_id.to_string()); assert_eq!(responses[0].author.id, "1472426752046002208"); assert_eq!(responses[1].author.id, "1472426752046002208"); assert!(recorded_deletions(&deleted).is_empty()); } #[tokio::test] async fn build_path_reaps_orphaned_messages() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let orphan_id = 1_509_197_195_776_110_592; let kept_id = 1_509_197_195_776_110_593; let responses = shard .build_api_responses_from_messages( vec![orphaned_message(orphan_id), webhook_message(kept_id)], build_options(), ) .await .unwrap(); assert_eq!(responses.len(), 1); assert_eq!(responses[0].id, kept_id.to_string()); assert_eq!( recorded_deletions(&deleted), vec![(10, snowflake_to_bucket(orphan_id), orphan_id)] ); } #[tokio::test] async fn build_path_reaps_orphans_of_each_batch_separately() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let first_orphan_id = 1_509_197_195_776_110_592; let second_orphan_id = 1_509_197_195_776_110_594; shard .build_api_responses_from_messages( vec![ orphaned_message(first_orphan_id), webhook_message(1_509_197_195_776_110_593), ], build_options(), ) .await .unwrap(); shard .build_api_responses_from_messages( vec![orphaned_message(second_orphan_id)], build_options(), ) .await .unwrap(); assert_eq!( recorded_deletions(&deleted), vec![ (10, snowflake_to_bucket(first_orphan_id), first_orphan_id), (10, snowflake_to_bucket(second_orphan_id), second_orphan_id), ] ); } #[tokio::test] async fn build_path_reaps_a_repeated_orphan_without_failing() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let orphan_id = 1_509_197_195_776_110_592; for _ in 0..2 { let responses = shard .build_api_responses_from_messages( vec![orphaned_message(orphan_id)], build_options(), ) .await .unwrap(); assert!(responses.is_empty()); } assert_eq!( recorded_deletions(&deleted), vec![ (10, snowflake_to_bucket(orphan_id), orphan_id), (10, snowflake_to_bucket(orphan_id), orphan_id), ] ); } #[tokio::test] async fn build_path_leaves_orphans_the_requester_cannot_see() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let orphan_id = 1_509_197_195_776_110_592; let options = ResponseBuildOptions { can_read_message_history: false, message_history_cutoff_ms: Some(snowflake_to_epoch_millis(orphan_id) + 1), ..build_options() }; let responses = shard .build_api_responses_from_messages(vec![orphaned_message(orphan_id)], options) .await .unwrap(); assert!(responses.is_empty()); assert!(recorded_deletions(&deleted).is_empty()); } #[tokio::test] async fn single_build_path_reaps_an_orphaned_message() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let orphan_id = 1_509_197_195_776_110_592; let kept_id = 1_509_197_195_776_110_593; let orphan_response = shard .build_api_response_from_message(orphaned_message(orphan_id), build_options()) .await .unwrap(); let kept_response = shard .build_api_response_from_message(webhook_message(kept_id), build_options()) .await .unwrap(); assert!(orphan_response.is_none()); assert_eq!(kept_response.unwrap().id, kept_id.to_string()); assert_eq!( recorded_deletions(&deleted), vec![(10, snowflake_to_bucket(orphan_id), orphan_id)] ); } #[tokio::test] async fn single_build_path_leaves_an_orphan_the_requester_cannot_see() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let orphan_id = 1_509_197_195_776_110_592; let options = ResponseBuildOptions { can_read_message_history: false, message_history_cutoff_ms: Some(snowflake_to_epoch_millis(orphan_id) + 1), ..build_options() }; let response = shard .build_api_response_from_message(orphaned_message(orphan_id), options) .await .unwrap(); assert!(response.is_none()); assert!(recorded_deletions(&deleted).is_empty()); } const SIGNATURE_VECTORS: &str = include_str!("../../fluxer_common/src/testdata/attachment_url_signature_vectors.json"); const SIGNED_CHANNEL_ID: i64 = 1_544_725_486_800_732_163; const SIGNED_ATTACHMENT_ID: i64 = 1_544_971_349_200_470_016; const SIGNED_FILENAME: &str = "cat.gif"; const SIGNED_ANCHOR_SECS: u64 = 1_788_420_273; fn signing_options() -> ResponseBuildOptions { ResponseBuildOptions { media_endpoint: "https://media.test".to_owned(), attachment_url_secret: (0u8..32).collect(), ..build_options() } } fn signed_attachment() -> MessageAttachment { serde_json::from_value(json!({ "attachment_id": SIGNED_ATTACHMENT_ID.to_string(), "filename": SIGNED_FILENAME, "size": 1024, })) .unwrap() } fn signed_storage_key() -> String { make_attachment_cdn_key(SIGNED_CHANNEL_ID, SIGNED_ATTACHMENT_ID, SIGNED_FILENAME) } fn unsigned_attachment_url(options: &ResponseBuildOptions) -> String { make_attachment_cdn_url( &options.media_endpoint, SIGNED_CHANNEL_ID, SIGNED_ATTACHMENT_ID, SIGNED_FILENAME, ) } fn assert_signs(url: &str, unsigned_prefix: &str, options: &ResponseBuildOptions, now: u64) { assert_eq!( Some(format!("{unsigned_prefix}?")), url.split_once("ex=").map(|(head, _)| head.to_owned()), "{url}" ); assert!(!url.contains("/external/"), "{url}"); assert!(!url.ends_with('&'), "{url}"); assert!(!url.contains("&&"), "{url}"); let query = url.split_once('?').expect("a signed url has a query").1; assert_eq!( fluxer_common::attachment_url_signature::Verdict::Valid, fluxer_common::attachment_url_signature::verify( &signed_storage_key(), Some(query), &[&options.attachment_url_secret], now, ) .verdict, "{url}" ); } #[test] fn the_signing_anchor_is_the_attachment_snowflake_second() { assert_eq!( SIGNED_ANCHOR_SECS, attachment_anchor_secs(SIGNED_ATTACHMENT_ID) ); assert_eq!(0, attachment_anchor_secs(i64::MIN)); } #[test] fn attachment_urls_match_the_shared_signature_fixture() { let fixture: serde_json::Value = serde_json::from_str(SIGNATURE_VECTORS).unwrap(); let secret = BASE64_STANDARD .decode(fixture["secrets_base64"][0].as_str().unwrap()) .unwrap(); let options = ResponseBuildOptions { attachment_url_secret: secret, ..signing_options() }; let parsed = fixture["sign"].as_array().unwrap(); let mut run = 0; let mut data_packages = 0; for case in parsed { let name = case["name"].as_str().unwrap(); if case["uc"].as_str().unwrap() == "dp" { data_packages += 1; continue; } assert_eq!("", case["uc"].as_str().unwrap(), "{name} uc"); let channel_id: i64 = case["channel_id"].as_str().unwrap().parse().unwrap(); let attachment_id: i64 = case["attachment_id"].as_str().unwrap().parse().unwrap(); let filename = case["filename"].as_str().unwrap(); let now = case["now"].as_u64().unwrap(); let url = case["url"].as_str().unwrap(); let signed = case["signed"].as_str().unwrap(); assert_eq!( case["anchor"].as_u64().unwrap(), attachment_anchor_secs(attachment_id), "{name} anchor" ); assert_eq!( case["storage_key"].as_str().unwrap(), make_attachment_cdn_key(channel_id, attachment_id, filename), "{name} key" ); let unsigned = make_attachment_cdn_url( &options.media_endpoint, channel_id, attachment_id, filename, ); assert_eq!( Some(unsigned.as_str()), url.split(['?', '#']).next(), "{name} unsigned url" ); if url == unsigned { assert_eq!( signed, attachment_cdn_url(channel_id, attachment_id, filename, &options, now), "{name} attachment url" ); } assert_eq!( signed, media_proxy_url_at(url, &options, now), "{name} embed url" ); run += 1; } assert_eq!(parsed.len(), run + data_packages); assert!(run >= 5); assert!(data_packages >= 1); } #[test] fn an_attachment_signs_url_and_proxy_url_identically() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let options = signing_options(); let mapped = shard .map_attachment( SIGNED_CHANNEL_ID, &signed_attachment(), &options, &ResponseContext::default(), ) .expect("an attachment with an id maps"); let url = mapped.url.expect("a live attachment has a url"); assert_eq!(Some(url.clone()), mapped.proxy_url); assert_signs( &url, &unsigned_attachment_url(&options), &options, unix_seconds(), ); } #[test] fn an_own_url_whose_filename_contains_a_slash_is_signed() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS + 10; let key = "attachments/1544725486800732163/1544971349200470016/a/b.gif"; for url in [ format!("{}/{key}", options.media_endpoint), format!( "{}/attachments/1544725486800732163/1544971349200470016/a%2Fb.gif", options.media_endpoint ), ] { let signed = media_proxy_url_at(&url, &options, now); assert_eq!( Some(format!("{url}?")), signed.split_once("ex=").map(|(head, _)| head.to_owned()), "{url}" ); let query = signed.split_once('?').expect("a signed url has a query").1; assert_eq!( fluxer_common::attachment_url_signature::Verdict::Valid, fluxer_common::attachment_url_signature::verify( key, Some(query), &[&options.attachment_url_secret], now, ) .verdict, "{url}" ); } for refused in [ format!( "{}/attachments/1544725486800732163/1544971349200470016/a//b.gif", options.media_endpoint ), format!( "{}/attachments/1544725486800732163/1544971349200470016/a/../b.gif", options.media_endpoint ), format!( "{}/attachments/1544725486800732163/1544971349200470016/a/", options.media_endpoint ), ] { assert_eq!( refused, media_proxy_url_at(&refused, &options, now), "{refused}" ); } } #[test] fn a_decayed_attachment_is_never_signed() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let options = signing_options(); let context = ResponseContext { attachment_decay: [( SIGNED_ATTACHMENT_ID, DateTime::::from_timestamp_millis(unix_millis() - 1_000).unwrap(), )] .into_iter() .collect(), ..ResponseContext::default() }; let mapped = shard .map_attachment(SIGNED_CHANNEL_ID, &signed_attachment(), &options, &context) .expect("a decayed attachment still maps"); assert_eq!(None, mapped.url); assert_eq!(None, mapped.proxy_url); assert_eq!(Some(true), mapped.expired); } #[test] fn attachment_expiry_is_only_looked_up_while_decay_is_on() { let mut message = orphaned_message(1); message.attachments = Some(vec![signed_attachment()]); let off = ResponseBuildOptions { attachment_decay: false, ..build_options() }; assert_eq!( HashSet::from([SIGNED_ATTACHMENT_ID]), build_options().expiring_attachment_ids(&[&message]) ); assert!(off.expiring_attachment_ids(&[&message]).is_empty()); } #[tokio::test] async fn a_passed_expiry_is_ignored_in_responses_while_decay_is_off() { let shard = MessagesShard { storage: MessagesStorage::Expiries(std::sync::Arc::new(HashMap::from([( SIGNED_ATTACHMENT_ID, DateTime::::from_timestamp_millis(unix_millis() - 1_000).unwrap(), )]))), transport: InMemoryTransport::new(), }; let users = stub_user_service(&shard.transport).await; let mut message = authored_message(1_509_197_195_776_110_592); message.attachments = Some(vec![signed_attachment()]); let off = ResponseBuildOptions { attachment_decay: false, ..signing_options() }; let while_on = shard .build_api_response_from_message(message.clone(), signing_options()) .await .unwrap() .expect("the message is visible"); let while_off = shard .build_api_response_from_message(message.clone(), off.clone()) .await .unwrap() .expect("the message is visible"); let listed_while_off = shard .build_api_responses_from_messages(vec![message], off) .await .unwrap(); users.abort(); let expired = &while_on.attachments[0]; assert_eq!(Some(true), expired.expired); assert_eq!(None, expired.url); assert!(expired.expires_at.is_some()); for response in [&while_off, &listed_while_off[0]] { let kept = &response.attachments[0]; assert_eq!(None, kept.expired); assert_eq!(None, kept.expires_at); assert!(kept.url.is_some()); } } #[test] fn an_attachment_backed_embed_signs_every_own_url_field() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let options = signing_options(); let unsigned = unsigned_attachment_url(&options); let embed: MessageEmbed = serde_json::from_value(json!({ "type": "rich", "url": unsigned, "author": {"name": "author", "url": unsigned, "icon_url": unsigned}, "provider": {"name": "provider", "url": unsigned}, "footer": {"text": "footer", "icon_url": unsigned}, "image": {"url": unsigned}, "thumbnail": {"url": unsigned}, })) .unwrap(); let mapped = shard.map_embed(&embed, &options); let now = unix_seconds(); let base = mapped.base; let author = base.author.expect("the embed has an author"); let provider = base.provider.expect("the embed has a provider"); let footer = base.footer.expect("the embed has a footer"); let image = base.image.expect("the embed has an image"); let thumbnail = base.thumbnail.expect("the embed has a thumbnail"); for signed in [ base.url.expect("the embed has a url"), author.url.expect("the author has a url"), author.icon_url.expect("the author has an icon url"), author .proxy_icon_url .expect("the author has a proxy icon url"), provider.url.expect("the provider has a url"), footer.icon_url.expect("the footer has an icon url"), footer .proxy_icon_url .expect("the footer has a proxy icon url"), image.url.clone(), image.proxy_url.clone(), thumbnail.url, thumbnail.proxy_url, ] { assert_signs(&signed, &unsigned, &options, now); } assert_eq!(image.url, image.proxy_url); } #[test] fn a_foreign_embed_url_is_never_rewritten() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let options = signing_options(); let embed: MessageEmbed = serde_json::from_value(json!({ "type": "link", "url": "https://example.com/article", "author": {"name": "author", "url": "https://example.com/author"}, "provider": {"name": "provider", "url": "https://example.com"}, })) .unwrap(); let mapped = shard.map_embed(&embed, &options); let base = mapped.base; assert_eq!( Some("https://example.com/article".to_owned()), base.url, "a page link never reaches the media proxy" ); assert_eq!( Some("https://example.com/author".to_owned()), base.author.expect("the embed has an author").url ); assert_eq!( Some("https://example.com".to_owned()), base.provider.expect("the embed has a provider").url ); } #[test] fn an_own_attachment_embed_url_is_signed_for_every_origin_spelling() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; for prefix in [ "https://media.test", "https://MEDIA.test", "https://Media.Test", "https://media.test:443", ] { let input = format!("{prefix}/{}", signed_storage_key()); let signed = media_proxy_url_at(&input, &options, now); assert_signs(&signed, &input, &options, now); } } #[test] fn an_own_attachment_embed_url_keeps_its_transform_parameters_and_fragment() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; let input = format!( "https://media.test/{}?width=64&format=webp#frame", signed_storage_key() ); let signed = media_proxy_url_at(&input, &options, now); assert!(signed.ends_with("&width=64&format=webp#frame"), "{signed}"); assert!(!signed.contains("/external/"), "{signed}"); let query = signed .split_once('?') .expect("a signed url has a query") .1 .split_once('#') .expect("the fragment is kept") .0; assert_eq!( fluxer_common::attachment_url_signature::Verdict::Valid, fluxer_common::attachment_url_signature::verify( &signed_storage_key(), Some(query), &[&options.attachment_url_secret], now, ) .verdict ); } #[test] fn an_own_endpoint_url_outside_attachments_passes_through_unsigned() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; for input in [ "https://media.test/emojis/1544725486800732163.webp", "https://media.test/avatars/1/abc.png", "https://media.test/stickers/1/abc.png", "https://media.test/attachments/1544725486800732163", "https://media.test/attachments/1544725486800732163/1544971349200470016", "https://media.test/attachments/1544725486800732163/1544971349200470016/", "https://media.test/attachments/1544725486800732163/not-a-snowflake/cat.gif", ] { assert_eq!(input, media_proxy_url_at(input, &options, now), "{input}"); } } #[test] fn a_foreign_url_is_still_wrapped_into_the_external_path() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; for input in [ "https://example.com/attachments/1544725486800732163/1544971349200470016/cat.gif", "https://media.test.example.com/attachments/1/2/cat.gif", "http://media.test/attachments/1/2/cat.gif", "https://media.test:8443/attachments/1/2/cat.gif", ] { let wrapped = media_proxy_url_at(input, &options, now); assert!( wrapped.starts_with("https://media.test/external/"), "{input}" ); assert!(!wrapped.contains("ex="), "{input}"); } } #[test] fn an_absent_attachment_secret_leaves_every_url_unsigned() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let options = ResponseBuildOptions { attachment_url_secret: Vec::new(), ..signing_options() }; let unsigned = unsigned_attachment_url(&options); let mapped = shard .map_attachment( SIGNED_CHANNEL_ID, &signed_attachment(), &options, &ResponseContext::default(), ) .expect("an attachment with an id maps"); assert_eq!(Some(unsigned.clone()), mapped.url); assert_eq!(Some(unsigned.clone()), mapped.proxy_url); assert_eq!( unsigned, media_proxy_url_at(&unsigned, &options, SIGNED_ANCHOR_SECS) ); assert_eq!( unsigned, attachment_cdn_url( SIGNED_CHANNEL_ID, SIGNED_ATTACHMENT_ID, SIGNED_FILENAME, &options, SIGNED_ANCHOR_SECS ) ); } #[test] fn a_self_hosted_path_prefix_only_matches_whole_segments() { let options = ResponseBuildOptions { media_endpoint: "https://self.test/media".to_owned(), ..signing_options() }; let now = SIGNED_ANCHOR_SECS; let own = format!("https://self.test/media/{}", signed_storage_key()); let signed = media_proxy_url_at(&own, &options, now); assert_signs(&signed, &own, &options, now); let foreign = format!("https://self.test/mediaxyz/{}", signed_storage_key()); assert!( media_proxy_url_at(&foreign, &options, now) .starts_with("https://self.test/media/external/") ); } #[test] fn a_percent_encoded_filename_signs_the_decoded_key() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; let encoded = "https://media.test/attachments/1544725486800732163/1544971349200470016/caf%C3%A9.gif"; let decoded = "https://media.test/attachments/1544725486800732163/1544971349200470016/café.gif"; let key = "attachments/1544725486800732163/1544971349200470016/café.gif"; for input in [encoded, decoded] { let signed = media_proxy_url_at(input, &options, now); let query = signed.split_once('?').expect("a signed url has a query").1; assert_eq!( fluxer_common::attachment_url_signature::Verdict::Valid, fluxer_common::attachment_url_signature::verify( key, Some(query), &[&options.attachment_url_secret], now, ) .verdict, "{input}" ); } } #[test] fn a_percent_encoded_slash_signs_the_key_the_media_proxy_decodes() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; for input in [ "https://media.test/attachments/1544725486800732163%2F1544971349200470016/cat.gif", "https://media.test/attachments/1544725486800732163%2f1544971349200470016%2Fcat.gif", "https://media.test/attachments/%31544725486800732163/1544971349200470016/cat.gif", ] { let raw_path = input.strip_prefix("https://media.test").unwrap(); assert_eq!( Some(signed_storage_key()), fluxer_common::attachment_url_signature::decode_key(raw_path), "{input}" ); let signed = media_proxy_url_at(input, &options, now); assert_signs(&signed, input, &options, now); } for input in [ "https://media.test/attachments/1544725486800732163/1544971349200470016/%2F", "https://media.test/attachments/1544725486800732163/1544971349200470016/cat.gif%00", ] { assert_eq!(input, media_proxy_url_at(input, &options, now), "{input}"); } } #[test] fn a_dot_segment_is_never_signed_under_its_normalised_key() { let options = signing_options(); let now = SIGNED_ANCHOR_SECS; for input in [ "https://media.test/attachments/1544725486800732163/9/../1544971349200470016/cat.gif", "https://media.test/attachments/1544725486800732163/1544971349200470016/./cat.gif", "https://media.test/attachments/1544725486800732163/1544971349200470016/%2E/cat.gif", "https://media.test/attachments/1544725486800732163/9/%2e%2e/1544971349200470016/cat.gif", ] { assert_eq!( format!("/{}", signed_storage_key()), Url::parse(input).unwrap().path(), "{input}" ); assert_eq!(input, media_proxy_url_at(input, &options, now), "{input}"); } for input in [ "https://media.test/attachments/1544725486800732163/1544971349200470016/..", "https://media.test/attachments/+1544725486800732163/1544971349200470016/cat.gif", "https://media.test/attachments/1544725486800732163/-1544971349200470016/cat.gif", "https://media.test/attachments/1544725486800732163/1544971349200470016/caf%C3%28.gif", ] { assert_eq!(input, media_proxy_url_at(input, &options, now), "{input}"); } } #[test] fn a_request_without_the_secret_field_still_decodes() { let request: MessageRequest = serde_json::from_value(json!({ "op": "ListResponses", "channel_id": "10", "viewer_user_id": "1", "limit": 50, "can_read_message_history": true, "media_endpoint": "https://media.test", "media_proxy_secret_key": "secret", })) .expect("an older caller omits the attachment secret"); match request { MessageRequest::ListResponses { attachment_url_secret_base64, .. } => assert_eq!(None, attachment_url_secret_base64), other => panic!("unexpected request {other:?}"), } } #[test] fn a_secret_that_is_not_base64_disables_signing() { assert!(decode_attachment_url_secret(None).is_empty()); assert!(decode_attachment_url_secret(Some("")).is_empty()); assert!(decode_attachment_url_secret(Some("not base64!")).is_empty()); assert_eq!( (0u8..32).collect::>(), decode_attachment_url_secret(Some("AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8=")) ); } fn typed_message(message_id: i64, message_type: i32, flags: i64) -> Message { serde_json::from_value(json!({ "message_id": message_id.to_string(), "channel_id": "10", "bucket": 416, "author_id": "1472426752046002208", "type": message_type, "flags": flags, "version": 1 })) .unwrap() } fn masked_options() -> ResponseBuildOptions { ResponseBuildOptions { threads_mask: true, ..build_options() } } #[tokio::test] async fn excluded_types_refill_the_page_past_hidden_rows() { let channel = [ (100, 18), (99, 0), (98, 18), (97, 18), (96, 0), (95, 0), (94, 0), ] .map(|(id, message_type)| typed_message(id, message_type, 0)); let fetches = std::sync::Mutex::new(Vec::new()); let collected = collect_excluding(3, &[18], None, |cursor, size| { let fetches = &fetches; let channel = &channel; async move { fetches.lock().unwrap().push((cursor, size)); Ok(channel .iter() .filter(|message| cursor.is_none_or(|cursor| message.message_id < cursor)) .take(size as usize) .cloned() .collect()) } }) .await .unwrap(); assert_eq!( collected.iter().map(|m| m.message_id).collect::>(), vec![99, 96, 95] ); assert_eq!( *fetches.lock().unwrap(), vec![(None, 3), (Some(98), EXCLUDED_TYPE_REFILL_PAGE_SIZE)] ); } #[tokio::test] async fn excluded_types_stop_at_the_end_of_the_channel() { let channel = [(100, 18), (99, 0), (98, 18)] .map(|(id, message_type)| typed_message(id, message_type, 0)); let fetches = std::sync::Mutex::new(0_usize); let collected = collect_excluding(2, &[18], None, |cursor, size| { let fetches = &fetches; let channel = &channel; async move { *fetches.lock().unwrap() += 1; Ok(channel .iter() .filter(|message| cursor.is_none_or(|cursor| message.message_id < cursor)) .take(size as usize) .cloned() .collect()) } }) .await .unwrap(); assert_eq!( collected.iter().map(|m| m.message_id).collect::>(), vec![99] ); assert_eq!(*fetches.lock().unwrap(), 2); } #[tokio::test] async fn excluded_types_reach_past_a_long_hidden_run() { let channel = (0..3_000_i64) .map(|offset| typed_message(10_000 - offset, 18, 0)) .chain([typed_message(6_999, 0, 0), typed_message(6_998, 0, 0)]) .collect::>(); let collected = collect_excluding(2, &[18], None, |cursor, size| { let channel = &channel; async move { Ok(channel .iter() .filter(|message| cursor.is_none_or(|cursor| message.message_id < cursor)) .take(size as usize) .cloned() .collect()) } }) .await .unwrap(); assert_eq!( collected.iter().map(|m| m.message_id).collect::>(), vec![6_999, 6_998] ); } #[tokio::test] async fn excluded_types_stop_at_the_scan_row_cap() { let scanned = std::sync::Mutex::new(0_usize); let collected = collect_excluding(2, &[18], None, |cursor, size| { let scanned = &scanned; async move { *scanned.lock().unwrap() += size as usize; let top = cursor.unwrap_or(1_000_000); Ok((1..=i64::from(size)) .map(|offset| typed_message(top - offset, 18, 0)) .collect()) } }) .await .unwrap(); assert!(collected.is_empty()); let scanned = *scanned.lock().unwrap(); assert!(scanned >= EXCLUDED_TYPE_SCAN_ROW_CAP); assert!(scanned < EXCLUDED_TYPE_SCAN_ROW_CAP + EXCLUDED_TYPE_REFILL_PAGE_SIZE as usize); } const LISTED_BASE_ID: i64 = 1_509_197_195_776_110_592; async fn listed_ids( types: &[i32], limit: u32, before: Option, after: Option, around: Option, excluded_types: &[i32], ) -> Vec { let rows = types .iter() .enumerate() .map(|(index, message_type)| { typed_message(LISTED_BASE_ID + index as i64 + 1, *message_type, 0) }) .collect::>(); let shard = MessagesShard { storage: MessagesStorage::Memory(std::sync::Arc::new(rows)), transport: InMemoryTransport::new(), }; let users = stub_user_service(&shard.transport).await; let at = |offset: i64| LISTED_BASE_ID + offset; let responses = shard .list_api_responses( 10, limit, before.map(at), after.map(at), around.map(at), ResponseBuildOptions { excluded_types: excluded_types.to_vec(), ..build_options() }, ) .await .unwrap(); users.abort(); responses .iter() .map(|response| response.id.parse::().unwrap() - LISTED_BASE_ID) .collect() } #[tokio::test] async fn excluded_types_list_the_latest_page_past_hidden_rows() { let types = [0, 0, 18, 0, 18, 18, 0, 18]; assert_eq!( listed_ids(&types, 3, None, None, None, &[18]).await, [7, 4, 2] ); assert_eq!( listed_ids(&types, 3, None, None, None, &[]).await, [8, 7, 6] ); } #[tokio::test] async fn excluded_types_list_after_an_id_in_ascending_pages() { let types = [18, 18, 18, 0, 18, 0, 0, 0]; assert_eq!( listed_ids(&types, 3, None, Some(0), None, &[18]).await, [7, 6, 4] ); assert_eq!( listed_ids(&types, 3, None, Some(0), None, &[]).await, [3, 2, 1] ); } #[tokio::test] async fn excluded_types_keep_a_before_and_after_window_inside_its_bounds() { let types = [0, 0, 0, 18, 0, 18, 18, 0]; assert_eq!( listed_ids(&types, 3, Some(9), Some(1), None, &[18]).await, [8, 5, 3] ); assert_eq!( listed_ids(&types, 3, Some(9), Some(4), None, &[18]).await, [8, 5] ); assert_eq!( listed_ids(&types, 3, Some(9), Some(1), None, &[]).await, [8, 7, 6] ); } #[tokio::test] async fn excluded_types_drop_a_hidden_around_target_and_refill_both_sides() { let types = [0, 0, 18, 0, 18, 0, 18, 0, 0]; assert_eq!( listed_ids(&types, 5, None, None, Some(5), &[18]).await, [8, 6, 4, 2] ); assert_eq!( listed_ids(&types, 5, None, None, Some(5), &[]).await, [7, 6, 5, 4, 3] ); } #[tokio::test] async fn threads_mask_drops_thread_created_rows_and_clears_thread_flags() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let users = stub_user_service(&shard.transport).await; let mut forwarded = typed_message(1_509_197_195_776_110_594, 0, 0x120 | 0x4); forwarded.message_snapshots = Some(vec![ serde_json::from_value(json!({"content": "x", "type": 0, "flags": 0x124})).unwrap(), ]); let messages = vec![ forwarded, typed_message(1_509_197_195_776_110_593, MESSAGE_TYPE_THREAD_CREATED, 0), ]; let masked = shard .build_api_responses_from_messages(messages.clone(), masked_options()) .await .unwrap(); let control = shard .build_api_responses_from_messages(messages, build_options()) .await .unwrap(); users.abort(); assert_eq!(masked.len(), 1); assert_eq!(masked[0].flags, 0x4); assert_eq!(masked[0].message_snapshots.as_ref().unwrap()[0].flags, 0x4); assert_eq!(control.len(), 2); assert_eq!(control[0].flags, 0x124); assert_eq!( control[0].message_snapshots.as_ref().unwrap()[0].flags, 0x124 ); assert_eq!(control[1].message_type, MESSAGE_TYPE_THREAD_CREATED); assert!(recorded_deletions(&deleted).is_empty()); } #[tokio::test] async fn threads_mask_drops_a_single_thread_created_row() { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let users = stub_user_service(&shard.transport).await; let starter = typed_message(1_509_197_195_776_110_593, MESSAGE_TYPE_THREAD_CREATED, 0); let masked = shard .build_api_response_from_message(starter.clone(), masked_options()) .await .unwrap(); let control = shard .build_api_response_from_message(starter, build_options()) .await .unwrap(); users.abort(); assert!(masked.is_none()); assert!(control.is_some()); } #[test] fn threads_mask_clears_nested_referenced_flags_and_nulls_a_thread_created_reference() { let shard = recording_shard(&DeletedMessageKeys::default()); let reply_to = |target: i64| { let mut reply = typed_message(1_509_197_195_776_110_600, 19, 0x20); reply.message_reference = Some(MessageReference { channel_id: Some(10), message_id: Some(target), guild_id: None, reference_type: Some(0), }); reply }; let referenced = typed_message(1_509_197_195_776_110_590, 0, 0x120); let thread_created = typed_message(1_509_197_195_776_110_591, MESSAGE_TYPE_THREAD_CREATED, 0); let context = ResponseContext { referenced_messages: HashMap::from([ ((10, referenced.message_id), referenced.clone()), ((10, thread_created.message_id), thread_created.clone()), ]), ..ResponseContext::default() }; let masked = shard.map_message_response( &reply_to(referenced.message_id), &masked_options(), &context, true, ); assert_eq!(masked.flags, 0); let nested = masked.referenced_message.unwrap().unwrap(); assert_eq!(nested.flags, 0); let control = shard.map_message_response( &reply_to(referenced.message_id), &build_options(), &context, true, ); assert_eq!(control.flags, 0x20); assert_eq!(control.referenced_message.unwrap().unwrap().flags, 0x120); let masked_starter = shard.map_message_response( &reply_to(thread_created.message_id), &masked_options(), &context, true, ); assert!(matches!(masked_starter.referenced_message, Some(None))); let control_starter = shard.map_message_response( &reply_to(thread_created.message_id), &build_options(), &context, true, ); assert!(matches!(control_starter.referenced_message, Some(Some(_)))); } #[test] fn thread_options_default_to_off_for_older_callers() { let list: MessageRequest = serde_json::from_value(json!({ "op": "ListResponses", "channel_id": "10", "viewer_user_id": "1", "limit": 50, "can_read_message_history": true, "media_endpoint": "https://media.test", "media_proxy_secret_key": "secret", })) .unwrap(); match list { MessageRequest::ListResponses { threads_mask, exclude_types, .. } => { assert!(!threads_mask); assert!(exclude_types.is_empty()); } other => panic!("unexpected request {other:?}"), } let get: MessageRequest = serde_json::from_value(json!({ "op": "GetResponseById", "channel_id": "10", "message_id": "11", "viewer_user_id": "1", "can_read_message_history": true, "media_endpoint": "https://media.test", "media_proxy_secret_key": "secret", })) .unwrap(); match get { MessageRequest::GetResponseById { threads_mask, .. } => assert!(!threads_mask), other => panic!("unexpected request {other:?}"), } let options = build_options(); assert_eq!(options.mask_flags(0x1ff), 0x1ff); assert!(!options.hides(&typed_message(1, MESSAGE_TYPE_THREAD_CREATED, 0))); } #[test] fn a_reference_without_a_message_id_round_trips() { let message: Message = serde_json::from_value(json!({ "message_id": "1509197195776110592", "channel_id": "10", "bucket": 416, "author_id": "1472426752046002208", "type": MESSAGE_TYPE_THREAD_CREATED, "version": 1, "message_reference": {"channel_id": "1509197195776110500", "guild_id": "7", "type": 0} })) .unwrap(); let reference = message.message_reference.as_ref().unwrap(); assert_eq!(reference.message_id, None); let mapped = map_message_reference(reference).expect("a thread reference maps"); assert_eq!(mapped.message_id, None); let encoded = serde_json::to_value(&mapped).unwrap(); assert_eq!( encoded, json!({"channel_id": "1509197195776110500", "guild_id": "7", "type": 0}) ); let decoded: ApiMessageReferenceResponse = serde_json::from_value(encoded).unwrap(); assert_eq!(decoded.message_id, None); let shard = recording_shard(&DeletedMessageKeys::default()); let response = shard.map_message_response( &message, &build_options(), &ResponseContext::default(), true, ); assert!(response.referenced_message.is_none()); let packed = rmp_serde::to_vec_named(&response).unwrap(); let unpacked: ApiMessageResponse = rmp_serde::from_slice(&packed).unwrap(); let unpacked_reference = unpacked.message_reference.unwrap(); assert_eq!(unpacked_reference.message_id, None); assert_eq!(unpacked_reference.channel_id, "1509197195776110500"); let reply = MessageReference { channel_id: Some(1), message_id: Some(2), guild_id: None, reference_type: None, }; assert_eq!( serde_json::to_value(map_message_reference(&reply).unwrap()).unwrap(), json!({"channel_id": "1", "message_id": "2", "type": 0}) ); } async fn stub_channel_mentions( transport: &InMemoryTransport, payloads: std::sync::Arc>>, ) -> tokio::task::JoinHandle<()> { let mut subscriber = transport .subscribe("rpc.gateway.guild.resolve_channel_mentions") .await .unwrap(); let transport = transport.clone(); tokio::spawn(async move { while let Some(message) = subscriber.next().await { payloads .lock() .unwrap() .push(serde_json::from_slice(TransportMessage::payload(&message)).unwrap()); let reply = json!({"ok": true, "result": {"channels": [ {"id": "20", "name": "general", "type": 0}, {"id": "21", "name": "thread", "type": 11}, {"id": "22", "name": "forum", "type": 15}, {"id": "23", "name": "news-thread", "type": 10} ]}}); let _ = reply_message(&message, &transport, reply.to_string().as_bytes()).await; } }) } #[tokio::test] async fn threads_mask_excludes_thread_channels_from_mention_resolution() { let shard = recording_shard(&DeletedMessageKeys::default()); let payloads = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); let stub = stub_channel_mentions(&shard.transport, payloads.clone()).await; let ids = HashSet::from([20, 21, 22, 23]); let masked = shard .resolve_channel_mentions( ids.clone(), &ResponseBuildOptions { source_guild_id: Some(7), ..masked_options() }, ) .await; let control = shard .resolve_channel_mentions( ids, &ResponseBuildOptions { source_guild_id: Some(7), ..build_options() }, ) .await; stub.abort(); let mut masked_ids = masked.into_keys().collect::>(); masked_ids.sort(); assert_eq!(masked_ids, vec!["20"]); assert_eq!(control.len(), 4); assert_eq!(control["23"].channel_type, 10); let payloads = payloads.lock().unwrap(); assert_eq!(payloads[0]["exclude_thread_channels"], json!(true)); assert!(payloads[1].get("exclude_thread_channels").is_none()); } const CROSSPOST_SOURCE_CHANNEL_ID: i64 = 500; const CROSSPOST_SOURCE_GUILD_ID: i64 = 600; const CROSSPOST_SOURCE_MESSAGE_ID: i64 = 1_509_197_195_776_110_590; fn referencing_message(message_type: i32, flags: i64, reference: serde_json::Value) -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": "10"}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": "1509197195776110600"}, "webhook_id": {"__fluxer_type": "bigint", "value": "77"}, "webhook_name": "Source #news", "type": message_type, "flags": flags, "content": "published", "message_reference": reference })) .unwrap() } fn source_reference() -> serde_json::Value { json!({ "channel_id": CROSSPOST_SOURCE_CHANNEL_ID.to_string(), "message_id": CROSSPOST_SOURCE_MESSAGE_ID.to_string(), "guild_id": CROSSPOST_SOURCE_GUILD_ID.to_string(), "type": MESSAGE_REFERENCE_TYPE_DEFAULT }) } fn source_message() -> Message { decode_postgres_message(json!({ "channel_id": {"__fluxer_type": "bigint", "value": CROSSPOST_SOURCE_CHANNEL_ID.to_string()}, "bucket": 416, "message_id": {"__fluxer_type": "bigint", "value": CROSSPOST_SOURCE_MESSAGE_ID.to_string()}, "author_id": {"__fluxer_type": "bigint", "value": "1472426752046002208"}, "flags": 1, "content": "published" })) .unwrap() } fn context_with_source() -> ResponseContext { ResponseContext { referenced_messages: [( (CROSSPOST_SOURCE_CHANNEL_ID, CROSSPOST_SOURCE_MESSAGE_ID), source_message(), )] .into_iter() .collect(), ..ResponseContext::default() } } fn serialized_response(message: &Message, context: &ResponseContext) -> serde_json::Value { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); serde_json::to_value(shard.map_message_response(message, &build_options(), context, true)) .expect("response serialises") } fn signed_attachment_response(message: &Message) -> serde_json::Value { let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let response = serde_json::to_value(shard.map_message_response( message, &signing_options(), &ResponseContext::default(), true, )) .expect("response serialises"); response["attachments"][0].clone() } fn assert_attachment_signed_for_channel(attachment: &serde_json::Value, channel_id: i64) { let options = signing_options(); let storage_key = make_attachment_cdn_key(channel_id, SIGNED_ATTACHMENT_ID, SIGNED_FILENAME); let unsigned = make_attachment_cdn_url( &options.media_endpoint, channel_id, SIGNED_ATTACHMENT_ID, SIGNED_FILENAME, ); let url = attachment["url"] .as_str() .expect("a live attachment has a url"); assert_eq!(attachment["proxy_url"], attachment["url"]); assert_eq!(attachment["id"], json!(SIGNED_ATTACHMENT_ID.to_string())); assert_eq!( Some(format!("{unsigned}?")), url.split_once("ex=").map(|(head, _)| head.to_owned()), "{url}" ); let query = url.split_once('?').expect("a signed url has a query").1; assert_eq!( fluxer_common::attachment_url_signature::Verdict::Valid, fluxer_common::attachment_url_signature::verify( &storage_key, Some(query), &[&options.attachment_url_secret], unix_seconds(), ) .verdict, "{url}" ); } #[test] fn crosspost_copy_attachment_urls_point_at_the_source_channel() { let mut copy = referencing_message(0, MESSAGE_FLAG_IS_CROSSPOST, source_reference()); copy.attachments = Some(vec![signed_attachment()]); assert_eq!( attachment_storage_channel_id(©), CROSSPOST_SOURCE_CHANNEL_ID ); let attachment = signed_attachment_response(©); assert_attachment_signed_for_channel(&attachment, CROSSPOST_SOURCE_CHANNEL_ID); assert!( !attachment["url"] .as_str() .unwrap() .contains(&format!("/attachments/{}/", copy.channel_id)), "{attachment}" ); } #[test] fn a_message_without_the_copy_flag_keeps_its_own_channel_for_attachments() { for message in [ referencing_message(19, 0, source_reference()), referencing_message(0, 0, source_reference()), webhook_message(1_509_197_195_776_110_600), ] { let mut message = message; message.attachments = Some(vec![signed_attachment()]); assert_eq!(attachment_storage_channel_id(&message), message.channel_id); let attachment = signed_attachment_response(&message); assert_attachment_signed_for_channel(&attachment, message.channel_id); } } #[test] fn a_copy_without_a_reference_channel_falls_back_to_its_own_channel() { let mut copy = referencing_message(0, MESSAGE_FLAG_IS_CROSSPOST, source_reference()); copy.message_reference = None; copy.attachments = Some(vec![signed_attachment()]); assert_eq!(attachment_storage_channel_id(©), copy.channel_id); let attachment = signed_attachment_response(©); assert_attachment_signed_for_channel(&attachment, copy.channel_id); } #[tokio::test] async fn crosspost_copy_omits_referenced_message_and_requests_no_fetch() { let copy = referencing_message(0, MESSAGE_FLAG_IS_CROSSPOST, source_reference()); assert_eq!(reply_target(©), None); let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let fetched = shard .fetch_referenced_messages(std::slice::from_ref(©), &build_options()) .await; assert!(fetched.is_empty()); let response = serialized_response(©, &context_with_source()); let object = response.as_object().expect("response is an object"); assert!(!object.contains_key("referenced_message")); assert_eq!( response["message_reference"], json!({ "channel_id": CROSSPOST_SOURCE_CHANNEL_ID.to_string(), "message_id": CROSSPOST_SOURCE_MESSAGE_ID.to_string(), "guild_id": CROSSPOST_SOURCE_GUILD_ID.to_string(), "type": MESSAGE_REFERENCE_TYPE_DEFAULT }) ); assert_eq!(response["flags"], json!(MESSAGE_FLAG_IS_CROSSPOST)); } #[test] fn crosspost_copy_with_other_flags_still_omits_referenced_message() { let copy = referencing_message( 0, MESSAGE_FLAG_IS_CROSSPOST | MESSAGE_FLAG_SUPPRESS_EMBEDS | (1 << 3), source_reference(), ); assert_eq!(reply_target(©), None); let response = serialized_response(©, &context_with_source()); assert!(response.get("referenced_message").is_none()); } #[test] fn reply_with_the_same_reference_still_resolves() { let reply = referencing_message(19, 0, source_reference()); assert_eq!( reply_target(&reply), Some((CROSSPOST_SOURCE_CHANNEL_ID, CROSSPOST_SOURCE_MESSAGE_ID)) ); let response = serialized_response(&reply, &context_with_source()); assert_eq!( response["referenced_message"]["id"], json!(CROSSPOST_SOURCE_MESSAGE_ID.to_string()) ); assert_eq!( response["referenced_message"]["content"], json!("published") ); } #[test] fn reply_whose_target_is_missing_serialises_referenced_message_null() { let reply = referencing_message(19, 0, source_reference()); let response = serialized_response(&reply, &ResponseContext::default()); assert_eq!( response.get("referenced_message"), Some(&serde_json::Value::Null) ); } #[tokio::test] async fn follow_add_reference_without_message_id_keeps_channel_and_guild() { let follow_add = referencing_message( 12, 0, json!({ "channel_id": CROSSPOST_SOURCE_CHANNEL_ID.to_string(), "guild_id": CROSSPOST_SOURCE_GUILD_ID.to_string(), "type": MESSAGE_REFERENCE_TYPE_DEFAULT }), ); assert_eq!(reply_target(&follow_add), None); let deleted = DeletedMessageKeys::default(); let shard = recording_shard(&deleted); let fetched = shard .fetch_referenced_messages(std::slice::from_ref(&follow_add), &build_options()) .await; assert!(fetched.is_empty()); let response = serialized_response(&follow_add, &context_with_source()); assert_eq!(response["type"], json!(12)); assert!(response.get("referenced_message").is_none()); assert_eq!( response["message_reference"], json!({ "channel_id": CROSSPOST_SOURCE_CHANNEL_ID.to_string(), "guild_id": CROSSPOST_SOURCE_GUILD_ID.to_string(), "type": MESSAGE_REFERENCE_TYPE_DEFAULT }) ); } #[test] fn map_message_reference_emits_every_field_for_a_copy() { let mapped = map_message_reference(&MessageReference { channel_id: Some(CROSSPOST_SOURCE_CHANNEL_ID), message_id: Some(CROSSPOST_SOURCE_MESSAGE_ID), guild_id: Some(CROSSPOST_SOURCE_GUILD_ID), reference_type: Some(MESSAGE_REFERENCE_TYPE_DEFAULT), }) .expect("copy reference maps"); assert_eq!(mapped.channel_id, CROSSPOST_SOURCE_CHANNEL_ID.to_string()); assert_eq!( mapped.message_id, Some(CROSSPOST_SOURCE_MESSAGE_ID.to_string()) ); assert_eq!(mapped.guild_id, Some(CROSSPOST_SOURCE_GUILD_ID.to_string())); assert_eq!(mapped.reference_type, MESSAGE_REFERENCE_TYPE_DEFAULT); } #[test] fn map_message_reference_still_requires_a_channel_id() { let mapped = map_message_reference(&MessageReference { channel_id: None, message_id: Some(CROSSPOST_SOURCE_MESSAGE_ID), guild_id: Some(CROSSPOST_SOURCE_GUILD_ID), reference_type: None, }); assert!(mapped.is_none()); } const HIDDEN_AUTHOR: i64 = 1_472_426_752_046_002_301; const OTHER_AUTHOR: i64 = 1_472_426_752_046_002_302; const CHANNEL: i64 = 10; const WINDOW_MS: i64 = 1_790_000_000_000; fn snowflake_at(epoch_millis: i64, sequence: i64) -> i64 { ((epoch_millis - FLUXER_EPOCH_MS) << 22) | sequence } fn message_by(author_id: i64, message_id: i64, extra: serde_json::Value) -> Message { let mut row = json!({ "channel_id": {"__fluxer_type": "bigint", "value": CHANNEL.to_string()}, "bucket": snowflake_to_bucket(message_id), "message_id": {"__fluxer_type": "bigint", "value": message_id.to_string()}, "author_id": {"__fluxer_type": "bigint", "value": author_id.to_string()}, "content": format!("message {message_id}"), }); row.as_object_mut() .unwrap() .extend(extra.as_object().unwrap().clone()); decode_postgres_message(row).unwrap() } async fn user_service( transport: &InMemoryTransport, hidden_since: Option, ) -> tokio::task::JoinHandle<()> { let mut subscriber = transport.subscribe("svc.users").await.unwrap(); let transport = transport.clone(); tokio::spawn(async move { while let Some(message) = subscriber.next().await { let request: serde_json::Value = serde_json::from_slice(message.payload()).unwrap(); let partials = request["user_ids"] .as_array() .unwrap() .iter() .map(|id| { let id = id.as_i64().unwrap(); json!({ "user_id": id, "username": format!("user{id}"), "discriminator": 1, "content_hidden_since": (id == HIDDEN_AUTHOR).then_some(hidden_since).flatten(), }) }) .collect::>(); let reply = serde_json::to_vec(&json!({"FoundPartials": partials})).unwrap(); let _ = reply_message(&message, &transport, &reply).await; } }) } fn memory_shard(messages: Vec) -> MessagesShard { MessagesShard { storage: MessagesStorage::Memory(std::sync::Arc::new(messages)), transport: InMemoryTransport::new(), } } fn channel_history() -> (Vec, Vec, Vec) { let mut messages = Vec::new(); let mut hidden = Vec::new(); let mut visible = Vec::new(); for offset in 0..3 { let id = snowflake_at(WINDOW_MS - 60_000 + offset, 0); messages.push(message_by(HIDDEN_AUTHOR, id, json!({}))); visible.push(id); } for offset in 0..2 { let id = snowflake_at(WINDOW_MS - 30_000 + offset, 0); messages.push(message_by(OTHER_AUTHOR, id, json!({}))); visible.push(id); } hidden.push(snowflake_at(WINDOW_MS, 0)); messages.push(message_by(HIDDEN_AUTHOR, hidden[0], json!({}))); for offset in 1..7 { let id = snowflake_at(WINDOW_MS + offset * 1_000, 0); messages.push(message_by(HIDDEN_AUTHOR, id, json!({}))); hidden.push(id); } visible.sort_unstable_by_key(|id| std::cmp::Reverse(*id)); (messages, hidden, visible) } fn ids(responses: &[ApiMessageResponse]) -> Vec { responses .iter() .map(|response| response.id.parse().unwrap()) .collect() } #[tokio::test] async fn hidden_window_drops_messages_from_every_list_shape_and_refills_the_page() { let (messages, hidden, visible) = channel_history(); let shard = memory_shard(messages); let users = user_service(&shard.transport, Some(WINDOW_MS)).await; let latest = shard .list_api_responses(CHANNEL, 3, None, None, None, build_options()) .await .unwrap(); assert_eq!(ids(&latest), visible[..3].to_vec()); let before = shard .list_api_responses(CHANNEL, 50, Some(hidden[6]), None, None, build_options()) .await .unwrap(); assert_eq!(ids(&before), visible); let after = shard .list_api_responses(CHANNEL, 50, None, Some(visible[1]), None, build_options()) .await .unwrap(); assert_eq!(ids(&after), vec![visible[0]]); let around = shard .list_api_responses(CHANNEL, 4, None, None, Some(hidden[0]), build_options()) .await .unwrap(); assert!(ids(&around).iter().all(|id| visible.contains(id))); assert!(!around.is_empty()); let mut staff_view = build_options(); staff_view.include_hidden = true; let all = shard .list_api_responses(CHANNEL, 50, None, None, None, staff_view) .await .unwrap(); assert_eq!(all.len(), visible.len() + hidden.len()); users.abort(); } #[tokio::test] async fn clearing_the_window_restores_every_message() { let (messages, hidden, visible) = channel_history(); let shard = memory_shard(messages); let users = user_service(&shard.transport, None).await; let latest = shard .list_api_responses(CHANNEL, 50, None, None, None, build_options()) .await .unwrap(); assert_eq!(latest.len(), visible.len() + hidden.len()); let single = shard .get_api_response(CHANNEL, hidden[0], build_options()) .await .unwrap(); assert_eq!(single.unwrap().id, hidden[0].to_string()); users.abort(); } #[tokio::test] async fn single_fetch_hides_messages_inside_the_window_only() { let (messages, hidden, visible) = channel_history(); let shard = memory_shard(messages); let users = user_service(&shard.transport, Some(WINDOW_MS)).await; for id in &hidden { let response = shard .get_api_response(CHANNEL, *id, build_options()) .await .unwrap(); assert!(response.is_none(), "{id}"); } let earlier = visible.last().copied().unwrap(); let response = shard .get_api_response(CHANNEL, earlier, build_options()) .await .unwrap(); assert_eq!(response.unwrap().author.id, HIDDEN_AUTHOR.to_string()); let mut staff_view = build_options(); staff_view.include_hidden = true; assert!( shard .get_api_response(CHANNEL, hidden[0], staff_view) .await .unwrap() .is_some() ); users.abort(); } #[tokio::test] async fn search_pins_and_saved_builds_drop_hidden_messages() { let (messages, hidden, visible) = channel_history(); let shard = memory_shard(messages.clone()); let users = user_service(&shard.transport, Some(WINDOW_MS)).await; let built = shard .build_api_responses_from_messages(messages.clone(), build_options()) .await .unwrap(); let mut built_ids = ids(&built); built_ids.sort_unstable_by_key(|id| std::cmp::Reverse(*id)); assert_eq!(built_ids, visible); assert!(built_ids.iter().all(|id| !hidden.contains(id))); users.abort(); let shard = memory_shard(messages.clone()); let users = user_service(&shard.transport, None).await; let restored = shard .build_api_responses_from_messages(messages, build_options()) .await .unwrap(); assert_eq!(restored.len(), visible.len() + hidden.len()); users.abort(); } #[tokio::test] async fn replies_lose_their_preview_and_forwards_vanish_when_the_source_is_hidden() { let (mut messages, hidden, _) = channel_history(); let reply_id = snowflake_at(WINDOW_MS + 20_000, 0); let forward_id = snowflake_at(WINDOW_MS + 21_000, 0); let reference = |kind: i32| { json!({"message_reference": { "channel_id": {"__fluxer_type": "bigint", "value": CHANNEL.to_string()}, "message_id": {"__fluxer_type": "bigint", "value": hidden[1].to_string()}, "type": kind, }}) }; messages.push(message_by(OTHER_AUTHOR, reply_id, reference(0))); let mut forward = reference(1); forward["message_snapshots"] = json!([{"content": "copied", "type": 0}]); forward["content"] = json!(""); messages.push(message_by(OTHER_AUTHOR, forward_id, forward)); let shard = memory_shard(messages); let users = user_service(&shard.transport, Some(WINDOW_MS)).await; let reply = shard .get_api_response(CHANNEL, reply_id, build_options()) .await .unwrap() .unwrap(); let reply = serde_json::to_value(&reply).unwrap(); assert!(reply["referenced_message"].is_null()); assert_eq!( reply["message_reference"]["message_id"], hidden[1].to_string() ); assert!( shard .get_api_response(CHANNEL, forward_id, build_options()) .await .unwrap() .is_none() ); let latest = shard .list_api_responses(CHANNEL, 2, None, None, None, build_options()) .await .unwrap(); assert_eq!(ids(&latest)[0], reply_id); assert!(!ids(&latest).contains(&forward_id)); users.abort(); } }