use std::time::Duration;
use async_trait::async_trait;
use crate::proto_helpers::MessageExt;
use waproto::whatsapp as wa;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum MsgSecretPolicy {
#[default]
Managed,
BotOnly,
Full,
Disabled,
}
impl MsgSecretPolicy {
pub fn persists(self) -> bool {
!matches!(self, MsgSecretPolicy::Disabled)
}
pub fn bot_only(self) -> bool {
matches!(self, MsgSecretPolicy::BotOnly)
}
pub fn prunes(self) -> bool {
!matches!(self, MsgSecretPolicy::Full)
}
pub fn bounds_retention(self) -> bool {
matches!(self, MsgSecretPolicy::Managed | MsgSecretPolicy::BotOnly)
}
}
#[derive(Debug, Clone, Copy)]
pub struct MsgSecretRetention {
pub text: Duration,
pub poll_event: Duration,
pub bot: Duration,
}
impl Default for MsgSecretRetention {
fn default() -> Self {
Self {
text: Duration::from_secs(30 * 86_400),
poll_event: Duration::from_secs(90 * 86_400),
bot: Duration::from_secs(30 * 86_400),
}
}
}
impl MsgSecretRetention {
fn horizon_secs(&self, class: RetentionClass) -> u64 {
match class {
RetentionClass::Text => self.text.as_secs(),
RetentionClass::PollEvent => self.poll_event.as_secs(),
RetentionClass::Bot => self.bot.as_secs(),
}
}
}
pub const EDIT_PROCESSING_WINDOW_SECS: i64 = 1200;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetentionClass {
Text,
PollEvent,
Bot,
}
pub fn is_bot_context(chat_is_bot: bool, msg: &wa::Message) -> bool {
chat_is_bot || invokes_bot(msg)
}
fn invokes_bot(msg: &wa::Message) -> bool {
let has_bot_metadata = |m: &wa::Message| {
m.message_context_info
.as_option()
.is_some_and(|c| c.bot_metadata.as_option().is_some())
};
has_bot_metadata(msg) || has_bot_metadata(msg.get_base_message()) || msg.mentions_any_bot()
}
fn message_is_poll_or_event(msg: &wa::Message) -> bool {
let base = msg.get_base_message();
base.poll_creation_message.as_option().is_some()
|| base.poll_creation_message_v2.as_option().is_some()
|| base.poll_creation_message_v3.as_option().is_some()
|| base.event_message.as_option().is_some()
}
pub fn classify(msg: &wa::Message, chat_is_bot: bool) -> RetentionClass {
classify_from_flags(
is_bot_context(chat_is_bot, msg),
message_is_poll_or_event(msg),
)
}
pub fn classify_from_flags(bot_context: bool, poll_or_event: bool) -> RetentionClass {
if bot_context {
RetentionClass::Bot
} else if poll_or_event {
RetentionClass::PollEvent
} else {
RetentionClass::Text
}
}
pub fn expires_at(
policy: MsgSecretPolicy,
retention: &MsgSecretRetention,
class: RetentionClass,
message_ts: Option<u64>,
now: i64,
) -> i64 {
if !policy.bounds_retention() {
return 0;
}
let base = message_ts
.and_then(|t| i64::try_from(t).ok())
.unwrap_or(now);
let horizon = i64::try_from(retention.horizon_secs(class)).unwrap_or(i64::MAX);
base.saturating_add(horizon)
}
pub fn within_seed_horizon(
retention: &MsgSecretRetention,
class: RetentionClass,
message_ts: Option<u64>,
now: i64,
) -> bool {
let Some(ts) = message_ts.and_then(|t| i64::try_from(t).ok()) else {
return true;
};
let horizon = i64::try_from(retention.horizon_secs(class)).unwrap_or(i64::MAX);
ts.saturating_add(horizon) > now
}
#[cfg(not(target_arch = "wasm32"))]
#[async_trait]
pub trait OriginalMessageResolver: Send + Sync {
async fn resolve_msg_secret(&self, chat: &str, sender: &str, msg_id: &str) -> Option<[u8; 32]>;
}
#[cfg(target_arch = "wasm32")]
#[async_trait(?Send)]
pub trait OriginalMessageResolver {
async fn resolve_msg_secret(&self, chat: &str, sender: &str, msg_id: &str) -> Option<[u8; 32]>;
}
#[cfg(test)]
mod tests {
use super::*;
const DAY: i64 = 86_400;
#[test]
fn full_and_disabled_never_expire() {
let r = MsgSecretRetention::default();
for policy in [MsgSecretPolicy::Full, MsgSecretPolicy::Disabled] {
assert_eq!(
expires_at(policy, &r, RetentionClass::Text, Some(1_000), 2_000),
0,
"{policy:?} must not set a deadline"
);
}
}
#[test]
fn managed_text_expires_30d_after_message_time() {
let r = MsgSecretRetention::default();
let msg_ts = 1_000_000u64;
let got = expires_at(
MsgSecretPolicy::Managed,
&r,
RetentionClass::Text,
Some(msg_ts),
5_000_000,
);
assert_eq!(got, msg_ts as i64 + 30 * DAY);
}
#[test]
fn poll_event_horizon_is_longer_than_text() {
let r = MsgSecretRetention::default();
let now = 10_000_000i64;
let text = expires_at(
MsgSecretPolicy::Managed,
&r,
RetentionClass::Text,
Some(1_000),
now,
);
let poll = expires_at(
MsgSecretPolicy::Managed,
&r,
RetentionClass::PollEvent,
Some(1_000),
now,
);
assert!(poll > text, "poll/event must outlive text secrets");
assert_eq!(poll - text, (90 - 30) * DAY);
}
#[test]
fn unknown_timestamp_expires_a_horizon_from_now_not_forever() {
let r = MsgSecretRetention::default();
let now = 5_000_000i64;
let got = expires_at(
MsgSecretPolicy::Managed,
&r,
RetentionClass::Text,
None,
now,
);
assert_eq!(got, now + 30 * DAY, "unknown age is bounded, never 0");
}
#[test]
fn seed_horizon_drops_old_text_keeps_recent_and_unknown() {
let r = MsgSecretRetention::default();
let now = 100 * DAY;
assert!(!within_seed_horizon(
&r,
RetentionClass::Text,
Some((60 * DAY) as u64),
now
));
assert!(within_seed_horizon(
&r,
RetentionClass::Text,
Some((90 * DAY) as u64),
now
));
assert!(within_seed_horizon(
&r,
RetentionClass::PollEvent,
Some((60 * DAY) as u64),
now
));
assert!(within_seed_horizon(&r, RetentionClass::Text, None, now));
}
#[test]
fn classify_from_flags_precedence() {
assert_eq!(classify_from_flags(true, true), RetentionClass::Bot);
assert_eq!(classify_from_flags(true, false), RetentionClass::Bot);
assert_eq!(classify_from_flags(false, true), RetentionClass::PollEvent);
assert_eq!(classify_from_flags(false, false), RetentionClass::Text);
}
#[test]
fn is_bot_context_detects_chat_and_invocation() {
assert!(is_bot_context(true, &wa::Message::default()));
let prompt = wa::Message {
message_context_info: buffa::MessageField::some(wa::MessageContextInfo {
bot_metadata: buffa::MessageField::some(wa::BotMetadata::default()),
..Default::default()
}),
..Default::default()
};
assert!(is_bot_context(false, &prompt));
let plain = wa::Message {
conversation: Some("hi".into()),
..Default::default()
};
assert!(!is_bot_context(false, &plain));
}
#[test]
fn policy_predicates() {
assert!(MsgSecretPolicy::Managed.persists());
assert!(MsgSecretPolicy::Managed.prunes());
assert!(MsgSecretPolicy::BotOnly.prunes());
assert!(!MsgSecretPolicy::Full.prunes());
assert!(MsgSecretPolicy::Disabled.prunes());
assert!(!MsgSecretPolicy::Disabled.persists());
assert!(MsgSecretPolicy::BotOnly.bot_only());
}
}