use std::collections::HashMap;
use thiserror::Error;
use wacore::poll;
use wacore_binary::{Jid, JidExt};
use waproto::whatsapp as wa;
use crate::client::Client;
use crate::send::{SendError, SendResult};
pub use wacore::poll::PollVoteCiphertext;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum PollError {
#[error("{0}")]
Send(#[from] SendError),
#[error("invalid poll: {0}")]
InvalidPoll(String),
#[error("client is not logged in")]
NotLoggedIn,
#[error("poll vote crypto failed: {0}")]
Crypto(#[source] anyhow::Error),
}
#[derive(Debug, Clone)]
pub struct PollOptionResult {
pub name: String,
pub voters: Vec<String>,
}
pub struct Polls<'a> {
client: &'a Client,
}
impl<'a> Polls<'a> {
pub(crate) fn new(client: &'a Client) -> Self {
Self { client }
}
pub async fn create(
&self,
to: impl Into<Jid>,
name: &str,
options: &[String],
selectable_count: u32,
) -> Result<(SendResult, Vec<u8>), PollError> {
let to = &to.into();
self.create_inner(to, name, options, selectable_count, None)
.await
}
pub async fn create_quiz(
&self,
to: impl Into<Jid>,
name: &str,
options: &[String],
correct_index: usize,
) -> Result<(SendResult, Vec<u8>), PollError> {
let to = &to.into();
self.create_inner(to, name, options, 1, Some(correct_index))
.await
}
async fn create_inner(
&self,
to: &Jid,
name: &str,
options: &[String],
selectable_count: u32,
correct_index: Option<usize>,
) -> Result<(SendResult, Vec<u8>), PollError> {
let poll_msg = build_poll_creation_message(name, options, selectable_count, correct_index)?;
let mut message = if selectable_count == 1 {
wa::Message {
poll_creation_message_v3: buffa::MessageField::some(poll_msg),
..Default::default()
}
} else {
wa::Message {
poll_creation_message: buffa::MessageField::some(poll_msg),
..Default::default()
}
};
let message_secret: Vec<u8> = {
use rand::Rng;
let mut secret = vec![0u8; 32];
rand::rng().fill_bytes(&mut secret);
secret
};
message.message_context_info = buffa::MessageField::some(wa::MessageContextInfo {
message_secret: Some(message_secret.clone()),
..Default::default()
});
let result = self.client.send_message(to, message).await?;
Ok((result, message_secret))
}
pub async fn vote(
&self,
chat_jid: impl Into<Jid>,
poll_msg_id: &str,
poll_creator_jid: &Jid,
message_secret: &[u8],
option_names: &[String],
) -> Result<SendResult, PollError> {
let chat_jid = &chat_jid.into();
let my_jid = self.client.pn().ok_or(PollError::NotLoggedIn)?;
let my_base = my_jid.to_non_ad();
let voter_jid = self
.resolve_voter_jid(poll_creator_jid, &my_base, poll_msg_id)
.await;
let voter_jid_str = voter_jid.to_string();
let creator_jid_str = poll_creator_jid.to_non_ad_string();
let selected_hashes: Vec<Vec<u8>> = option_names
.iter()
.map(|name| poll::compute_option_hash(name).to_vec())
.collect();
let (enc_payload, iv) = poll::encrypt_poll_vote_with_secret(
&selected_hashes,
message_secret,
poll_msg_id,
&creator_jid_str,
&voter_jid_str,
)
.map_err(PollError::Crypto)?;
let from_me = my_base.is_same_user_as(poll_creator_jid);
let poll_update = wa::message::PollUpdateMessage {
poll_creation_message_key: buffa::MessageField::some(wa::MessageKey {
remote_jid: Some(chat_jid.to_string()),
from_me: Some(from_me),
id: Some(poll_msg_id.to_string()),
participant: if chat_jid.is_group() {
Some(poll_creator_jid.to_string())
} else {
None
},
}),
vote: buffa::MessageField::some(wa::message::PollEncValue {
enc_payload: Some(enc_payload),
enc_iv: Some(iv.to_vec()),
}),
metadata: buffa::MessageField::none(),
sender_timestamp_ms: Some(wacore::time::now_millis()),
};
let message = wa::Message {
poll_update_message: buffa::MessageField::some(poll_update),
..Default::default()
};
Ok(self.client.send_message(chat_jid, message).await?)
}
async fn resolve_voter_jid(
&self,
poll_creator_jid: &Jid,
own_pn: &Jid,
poll_msg_id: &str,
) -> Jid {
if !poll_creator_jid.is_lid() {
return own_pn.clone();
}
match self.client.lid() {
Some(lid) => lid.to_non_ad(),
None => {
log::warn!(
"Poll {poll_msg_id} is LID-addressed but own LID is unknown; \
falling back to PN voter (host may fail to decrypt)"
);
own_pn.clone()
}
}
}
pub async fn decrypt_vote(
&self,
ciphertext: PollVoteCiphertext<'_>,
message_secret: &[u8],
poll_msg_id: &str,
poll_creator_jid: &Jid,
voter_jid: &Jid,
) -> Result<Vec<Vec<u8>>, PollError> {
let creator = poll_creator_jid.to_non_ad();
let voter = voter_jid.to_non_ad();
let creator_str = creator.to_string();
let voter_str = voter.to_string();
let creator_alt = self.swapped_user(&creator).await;
let voter_alt = self.swapped_user(&voter).await;
let fallback = Self::build_fallback(&creator_alt, &voter_alt);
poll::decrypt_poll_vote_with_fallback(
ciphertext,
message_secret,
poll_msg_id,
poll::PollVoteAddressing {
poll_creator_jid: &creator_str,
voter_jid: &voter_str,
},
fallback,
)
.map_err(PollError::Crypto)
}
async fn swapped_user(&self, jid: &Jid) -> Option<String> {
self.client
.swap_pn_lid_namespace(jid)
.await
.map(|j| j.to_non_ad_string())
}
fn build_fallback<'b>(
creator_alt: &'b Option<String>,
voter_alt: &'b Option<String>,
) -> Option<poll::PollVoteAddressing<'b>> {
match (creator_alt, voter_alt) {
(Some(c), Some(v)) => Some(poll::PollVoteAddressing {
poll_creator_jid: c,
voter_jid: v,
}),
_ => None,
}
}
pub async fn aggregate_votes(
&self,
poll_options: &[String],
votes: &[(&Jid, PollVoteCiphertext<'_>)],
message_secret: &[u8],
poll_msg_id: &str,
poll_creator_jid: &Jid,
) -> Result<Vec<PollOptionResult>, PollError> {
let option_hashes: Vec<([u8; 32], &str)> = poll_options
.iter()
.map(|name| (poll::compute_option_hash(name), name.as_str()))
.collect();
let creator = poll_creator_jid.to_non_ad();
let creator_str = creator.to_string();
let creator_alt = self.swapped_user(&creator).await;
let mut latest_votes: HashMap<String, (Option<String>, Vec<usize>)> =
HashMap::with_capacity(votes.len());
for (voter_jid, ciphertext) in votes {
let voter = voter_jid.to_non_ad();
let voter_str = voter.to_string();
let voter_alt = self.swapped_user(&voter).await;
let fallback = Self::build_fallback(&creator_alt, &voter_alt);
let canonical_voter = if voter.is_lid() {
voter_str.clone()
} else {
voter_alt.clone().unwrap_or_else(|| voter_str.clone())
};
match poll::decrypt_poll_vote_with_fallback(
*ciphertext,
message_secret,
poll_msg_id,
poll::PollVoteAddressing {
poll_creator_jid: &creator_str,
voter_jid: &voter_str,
},
fallback,
) {
Ok(hashes) => {
let display_jid = if voter.is_lid() {
None
} else if voter_alt.is_some() {
Some(voter_str)
} else {
None
};
if hashes.is_empty() {
latest_votes.remove(canonical_voter.as_str());
} else {
let selected_indices: Vec<usize> = hashes
.iter()
.filter_map(|h| {
<[u8; 32]>::try_from(h.as_slice()).ok().and_then(|arr| {
option_hashes.iter().position(|(oh, _)| *oh == arr)
})
})
.collect();
latest_votes.insert(canonical_voter, (display_jid, selected_indices));
}
}
Err(e) => {
log::warn!("Failed to decrypt vote from {voter_jid}: {e}");
}
}
}
let mut results: Vec<PollOptionResult> = poll_options
.iter()
.map(|name| PollOptionResult {
name: name.clone(),
voters: Vec::new(),
})
.collect();
for (canonical_jid, (display_jid, selected_indices)) in latest_votes {
let display_jid = display_jid.unwrap_or(canonical_jid);
if let Some((last_idx, prefix_indices)) = selected_indices.split_last() {
for idx in prefix_indices {
results[*idx].voters.push(display_jid.clone());
}
results[*last_idx].voters.push(display_jid);
}
}
Ok(results)
}
}
impl Client {
pub fn polls(&self) -> Polls<'_> {
Polls::new(self)
}
}
fn build_poll_creation_message(
name: &str,
options: &[String],
selectable_count: u32,
correct_index: Option<usize>,
) -> Result<wa::message::PollCreationMessage, PollError> {
if options.len() < 2 {
return Err(PollError::InvalidPoll(
"poll must have at least 2 options".into(),
));
}
if options.len() > 12 {
return Err(PollError::InvalidPoll(
"polls can have a maximum of 12 options".into(),
));
}
if selectable_count < 1 || selectable_count > options.len() as u32 {
return Err(PollError::InvalidPoll(format!(
"selectable_count must be between 1 and {} (got {selectable_count})",
options.len()
)));
}
let mut seen = std::collections::HashSet::new();
for opt in options {
if !seen.insert(opt) {
return Err(PollError::InvalidPoll(format!(
"duplicate option name: {opt}"
)));
}
}
let (poll_type, correct_answer) = match correct_index {
Some(idx) => {
let correct = options.get(idx).ok_or_else(|| {
PollError::InvalidPoll(format!(
"correct_index {idx} out of range (poll has {} options)",
options.len()
))
})?;
let answer = wa::message::poll_creation_message::Option {
option_name: Some(correct.clone()),
option_hash: Some(hex::encode(poll::compute_option_hash(correct))),
};
(
Some(wa::message::PollType::QUIZ),
buffa::MessageField::some(answer),
)
}
None => (None, buffa::MessageField::none()),
};
let poll_options: Vec<wa::message::poll_creation_message::Option> = options
.iter()
.map(|name| wa::message::poll_creation_message::Option {
option_name: Some(name.clone()),
option_hash: None,
})
.collect();
Ok(wa::message::PollCreationMessage {
enc_key: None,
name: Some(name.to_string()),
options: poll_options,
selectable_options_count: Some(selectable_count),
context_info: buffa::MessageField::none(),
poll_content_type: Some(wa::message::PollContentType::TEXT),
poll_type,
correct_answer,
..Default::default()
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::lid_pn_cache::LearningSource;
use crate::store::commands::DeviceCommand;
use crate::test_utils::create_test_client;
use std::sync::Arc;
#[test]
fn regular_poll_has_no_quiz_fields() {
let options = vec!["A".to_string(), "B".to_string(), "C".to_string()];
let msg = build_poll_creation_message("Q?", &options, 2, None).unwrap();
assert_eq!(msg.poll_type, None);
assert!(msg.correct_answer.is_unset());
assert_eq!(msg.selectable_options_count, Some(2));
assert_eq!(
msg.poll_content_type,
Some(wa::message::PollContentType::TEXT)
);
assert_eq!(msg.options.len(), 3);
assert!(msg.options.iter().all(|o| o.option_hash.is_none()));
}
#[test]
fn quiz_sets_poll_type_and_correct_answer() {
let options = vec!["A".to_string(), "B".to_string(), "C".to_string()];
let msg = build_poll_creation_message("Q?", &options, 1, Some(1)).unwrap();
assert_eq!(msg.poll_type, Some(wa::message::PollType::QUIZ));
let answer = msg
.correct_answer
.as_option()
.expect("quiz must carry a correct answer");
assert_eq!(answer.option_name.as_deref(), Some("B"));
let expected_hash = hex::encode(poll::compute_option_hash("B"));
assert_eq!(answer.option_hash.as_deref(), Some(expected_hash.as_str()));
}
#[test]
fn quiz_rejects_out_of_range_correct_index() {
let options = vec!["A".to_string(), "B".to_string()];
assert!(build_poll_creation_message("Q?", &options, 1, Some(5)).is_err());
}
#[tokio::test]
async fn voter_is_pn_when_poll_creator_is_pn() {
let client: Arc<Client> = create_test_client().await;
let own_pn = Jid::pn("5511999999999");
let creator = Jid::pn("5511777777777");
let voter = client
.polls()
.resolve_voter_jid(&creator, &own_pn, "POLLID")
.await;
assert_eq!(voter, own_pn);
}
#[tokio::test]
async fn voter_is_own_lid_when_poll_creator_is_lid() {
let client: Arc<Client> = create_test_client().await;
let own_lid: Jid = "888000888000888:3@lid".parse().unwrap();
client
.persistence_manager
.process_command(DeviceCommand::SetLid(Some(own_lid.clone())))
.await;
let own_pn = Jid::pn("5511999999999");
let creator = Jid::lid("111000111000111");
let voter = client
.polls()
.resolve_voter_jid(&creator, &own_pn, "POLLID")
.await;
assert!(voter.is_lid(), "voter must be LID-addressed in a LID poll");
assert_eq!(voter.user, own_lid.user);
assert_eq!(voter, own_lid.to_non_ad());
}
#[tokio::test]
async fn voter_falls_back_to_pn_when_own_lid_unknown() {
let client: Arc<Client> = create_test_client().await;
let own_pn = Jid::pn("5511999999999");
let creator = Jid::lid("111000111000111");
let voter = client
.polls()
.resolve_voter_jid(&creator, &own_pn, "POLLID")
.await;
assert_eq!(voter, own_pn);
}
#[tokio::test]
async fn decrypt_vote_recovers_when_fed_lid_but_encrypted_under_pn() {
let client: Arc<Client> = create_test_client().await;
let secret = [0x21u8; 32];
let stanza_id = "3EB0POLLVOTE";
let creator_pn = "5511777777777";
let creator_lid = "111000111000111";
let voter_pn = "5511888888888";
let voter_lid = "222000222000222";
client
.add_lid_pn_mapping(creator_lid, creator_pn, LearningSource::Usync)
.await
.unwrap();
client
.add_lid_pn_mapping(voter_lid, voter_pn, LearningSource::Usync)
.await
.unwrap();
let hashes = vec![poll::compute_option_hash("Yes").to_vec()];
let (enc, iv) = poll::encrypt_poll_vote_with_secret(
&hashes,
&secret,
stanza_id,
&Jid::pn(creator_pn).to_string(),
&Jid::pn(voter_pn).to_string(),
)
.unwrap();
let out = client
.polls()
.decrypt_vote(
PollVoteCiphertext {
enc_payload: &enc,
enc_iv: &iv,
},
&secret,
stanza_id,
&Jid::lid(creator_lid),
&Jid::lid(voter_lid),
)
.await
.expect("fallback should rescue the PN-encrypted vote");
assert_eq!(out, hashes);
}
#[tokio::test]
async fn decrypt_vote_fails_without_mapping() {
let client: Arc<Client> = create_test_client().await;
let secret = [0x21u8; 32];
let stanza_id = "3EB0POLLVOTE";
let (enc, iv) = poll::encrypt_poll_vote_with_secret(
&[poll::compute_option_hash("Yes").to_vec()],
&secret,
stanza_id,
&Jid::pn("5511777777777").to_string(),
&Jid::pn("5511888888888").to_string(),
)
.unwrap();
let res = client
.polls()
.decrypt_vote(
PollVoteCiphertext {
enc_payload: &enc,
enc_iv: &iv,
},
&secret,
stanza_id,
&Jid::lid("111000111000111"),
&Jid::lid("222000222000222"),
)
.await;
assert!(res.is_err(), "no mapping → no fallback → must not decrypt");
}
#[tokio::test]
async fn aggregate_votes_recovers_across_addressing() {
let client: Arc<Client> = create_test_client().await;
let secret = [0x31u8; 32];
let stanza_id = "3EB0AGG";
let options = vec!["Yes".to_string(), "No".to_string()];
let creator_pn = "5511777777777";
let creator_lid = "111000111000111";
let voter_pn = "5511888888888";
let voter_lid = "222000222000222";
client
.add_lid_pn_mapping(creator_lid, creator_pn, LearningSource::Usync)
.await
.unwrap();
client
.add_lid_pn_mapping(voter_lid, voter_pn, LearningSource::Usync)
.await
.unwrap();
let (enc, iv) = poll::encrypt_poll_vote_with_secret(
&[poll::compute_option_hash("Yes").to_vec()],
&secret,
stanza_id,
&Jid::pn(creator_pn).to_string(),
&Jid::pn(voter_pn).to_string(),
)
.unwrap();
let voter_lid_jid = Jid::lid(voter_lid);
let votes: Vec<(&Jid, PollVoteCiphertext)> = vec![(
&voter_lid_jid,
PollVoteCiphertext {
enc_payload: &enc,
enc_iv: &iv,
},
)];
let results = client
.polls()
.aggregate_votes(&options, &votes, &secret, stanza_id, &Jid::lid(creator_lid))
.await
.unwrap();
let yes = results.iter().find(|r| r.name == "Yes").unwrap();
assert_eq!(yes.voters.len(), 1, "the LID voter's 'Yes' must be tallied");
let no = results.iter().find(|r| r.name == "No").unwrap();
assert!(no.voters.is_empty());
}
#[tokio::test]
async fn aggregate_dedupes_revote_across_namespace() {
let client: Arc<Client> = create_test_client().await;
let secret = [0x41u8; 32];
let stanza_id = "3EB0REVOTE";
let options = vec!["Yes".to_string(), "No".to_string()];
let creator_pn = "5511777777777";
let creator_lid = "111000111000111";
let voter_pn = "5511888888888";
let voter_lid = "222000222000222";
client
.add_lid_pn_mapping(creator_lid, creator_pn, LearningSource::Usync)
.await
.unwrap();
client
.add_lid_pn_mapping(voter_lid, voter_pn, LearningSource::Usync)
.await
.unwrap();
let (enc_pn, iv_pn) = poll::encrypt_poll_vote_with_secret(
&[poll::compute_option_hash("Yes").to_vec()],
&secret,
stanza_id,
&Jid::pn(creator_pn).to_string(),
&Jid::pn(voter_pn).to_string(),
)
.unwrap();
let (enc_lid, iv_lid) = poll::encrypt_poll_vote_with_secret(
&[poll::compute_option_hash("No").to_vec()],
&secret,
stanza_id,
&Jid::lid(creator_lid).to_string(),
&Jid::lid(voter_lid).to_string(),
)
.unwrap();
let voter_pn_jid = Jid::pn(voter_pn);
let voter_lid_jid = Jid::lid(voter_lid);
let votes: Vec<(&Jid, PollVoteCiphertext)> = vec![
(
&voter_pn_jid,
PollVoteCiphertext {
enc_payload: &enc_pn,
enc_iv: &iv_pn,
},
),
(
&voter_lid_jid,
PollVoteCiphertext {
enc_payload: &enc_lid,
enc_iv: &iv_lid,
},
),
];
let results = client
.polls()
.aggregate_votes(&options, &votes, &secret, stanza_id, &Jid::lid(creator_lid))
.await
.unwrap();
let yes = results.iter().find(|r| r.name == "Yes").unwrap();
let no = results.iter().find(|r| r.name == "No").unwrap();
assert!(yes.voters.is_empty(), "the PN 'Yes' must be replaced");
assert_eq!(no.voters.len(), 1, "only the re-vote should count, once");
}
#[tokio::test]
async fn aggregate_clears_vote_across_namespace() {
let client: Arc<Client> = create_test_client().await;
let secret = [0x51u8; 32];
let stanza_id = "3EB0CLEAR";
let options = vec!["Yes".to_string(), "No".to_string()];
let creator_pn = "5511777777777";
let creator_lid = "111000111000111";
let voter_pn = "5511888888888";
let voter_lid = "222000222000222";
client
.add_lid_pn_mapping(creator_lid, creator_pn, LearningSource::Usync)
.await
.unwrap();
client
.add_lid_pn_mapping(voter_lid, voter_pn, LearningSource::Usync)
.await
.unwrap();
let (enc_pn, iv_pn) = poll::encrypt_poll_vote_with_secret(
&[poll::compute_option_hash("Yes").to_vec()],
&secret,
stanza_id,
&Jid::pn(creator_pn).to_string(),
&Jid::pn(voter_pn).to_string(),
)
.unwrap();
let (enc_clear, iv_clear) = poll::encrypt_poll_vote_with_secret(
&[],
&secret,
stanza_id,
&Jid::lid(creator_lid).to_string(),
&Jid::lid(voter_lid).to_string(),
)
.unwrap();
let voter_pn_jid = Jid::pn(voter_pn);
let voter_lid_jid = Jid::lid(voter_lid);
let votes: Vec<(&Jid, PollVoteCiphertext)> = vec![
(
&voter_pn_jid,
PollVoteCiphertext {
enc_payload: &enc_pn,
enc_iv: &iv_pn,
},
),
(
&voter_lid_jid,
PollVoteCiphertext {
enc_payload: &enc_clear,
enc_iv: &iv_clear,
},
),
];
let results = client
.polls()
.aggregate_votes(&options, &votes, &secret, stanza_id, &Jid::lid(creator_lid))
.await
.unwrap();
assert!(
results.iter().all(|r| r.voters.is_empty()),
"the LID clear-vote must remove the earlier PN 'Yes'"
);
}
}