imessage-database 5.0.1

/*!
 Poll payloads generated by the Messages Polls app.
*/

use std::collections::{HashMap, HashSet};

use base64::{Engine as _, engine::general_purpose};
use jzon::{self, JsonValue};
use plist::Value;

use crate::{
    error::plist::PlistParseError,
    util::plist::{get_string_from_nested_dict, parse_ns_keyed_archiver},
};

/// Poll option identifier.
pub type PollOptionID = String;

/// Poll option and its counted votes.
#[derive(Debug, PartialEq, Eq)]
pub struct PollOption {
    /// Option text.
    pub text: String,
    /// Handle that created the option.
    pub creator: String,
    /// Votes cast for this option.
    pub votes: Vec<PollVote>,
}

impl PollOption {
    fn from_json(json_data: &JsonValue) -> (PollOptionID, Self) {
        let poll_id = json_data["optionIdentifier"]
            .as_str()
            .unwrap_or_default()
            .to_string();
        let creator = json_data["creatorHandle"]
            .as_str()
            .unwrap_or_default()
            .to_string();
        let text = json_data["text"].as_str().unwrap_or_default().to_string();

        (
            poll_id,
            PollOption {
                text,
                creator,
                votes: Vec::new(),
            },
        )
    }
}

/// One vote in a poll update payload.
#[derive(Debug, PartialEq, Eq)]
pub struct PollVote {
    /// Handle that cast the vote.
    pub voter: String,
    /// Option identifier that received the vote.
    pub option_id: PollOptionID,
}

impl PollVote {
    fn from_json(payload: &JsonValue) -> Result<Self, PlistParseError> {
        let voter = payload["participantHandle"]
            .as_str()
            .ok_or(PlistParseError::MissingKey("participantHandle".to_string()))?
            .to_string();

        let option_id = payload["voteOptionIdentifier"]
            .as_str()
            .ok_or(PlistParseError::MissingKey(
                "voteOptionIdentifier".to_string(),
            ))?
            .to_string();

        Ok(PollVote { voter, option_id })
    }
}

/// Poll state reconstructed from poll and vote payloads.
#[derive(Debug, PartialEq, Eq)]
pub struct Poll {
    /// Options keyed by poll option identifier.
    pub options: HashMap<PollOptionID, PollOption>,
    /// Option identifiers in display order.
    pub order: Vec<PollOptionID>,
}

impl Poll {
    /// Parse the initial poll payload.
    pub fn from_payload(payload: &Value) -> Result<Self, PlistParseError> {
        let parsed = parse_ns_keyed_archiver(payload)?;

        let url = get_string_from_nested_dict(&parsed, "URL")
            .ok_or(PlistParseError::MissingKey("URL".to_string()))?;

        // Parse the JSON
        let parsed_json = base64_url_to_json(url)?;

        let mut options = HashMap::new();
        let mut order = Vec::new();
        let ordered_options = parsed_json["item"]["orderedPollOptions"].as_array().ok_or(
            PlistParseError::MissingKey("orderedPollOptions".to_string()),
        )?;

        for option in ordered_options {
            let (id, poll_option) = PollOption::from_json(option);
            order.push(id.clone());
            options.insert(id, poll_option);
        }

        Ok(Poll { options, order })
    }

    /// Apply a vote payload to the poll options.
    pub fn count_votes(&mut self, payload: &Value) -> Result<(), PlistParseError> {
        let parsed = parse_ns_keyed_archiver(payload)?;

        let url = get_string_from_nested_dict(&parsed, "URL")
            .ok_or(PlistParseError::MissingKey("URL".to_string()))?;

        // Parse the JSON
        let parsed_json = base64_url_to_json(url)?;

        let votes = parsed_json["item"]["votes"]
            .as_array()
            .ok_or(PlistParseError::MissingKey("votes".to_string()))?
            .iter()
            .map(PollVote::from_json)
            .collect::<Result<Vec<PollVote>, PlistParseError>>()?;

        self.apply_votes(votes);

        Ok(())
    }

    /// Apply one payload's votes as a snapshot.
    pub fn apply_votes(&mut self, votes: Vec<PollVote>) {
        let voters: HashSet<&str> = votes.iter().map(|vote| vote.voter.as_str()).collect();
        for option in self.options.values_mut() {
            option
                .votes
                .retain(|vote| !voters.contains(vote.voter.as_str()));
        }

        // Payloads are small, so a linear scan is cheaper than hashing owned pairs.
        let mut distinct: Vec<PollVote> = Vec::with_capacity(votes.len());
        for vote in votes {
            if !distinct.contains(&vote) {
                distinct.push(vote);
            }
        }

        for vote in distinct {
            if let Some(option) = self.options.get_mut(&vote.option_id) {
                option.votes.push(vote);
            }
        }
    }
}

fn base64_url_to_json(data: &str) -> Result<JsonValue, PlistParseError> {
    // Strip the fixed prefix
    let after_prefix = data
        .strip_prefix("data:,")
        .ok_or(PlistParseError::WrongMessageType)?;

    // Extract the base64 part before the first "?" (if present)
    let base64_part = after_prefix
        .split_once('?')
        .map_or(after_prefix, |(before, _)| before);

    // Decode the base64 part
    let bytes = String::from_utf8(
        general_purpose::URL_SAFE
            .decode(base64_part)
            .map_err(|_| PlistParseError::WrongMessageType)?,
    )
    .map_err(|_| PlistParseError::WrongMessageType)?;

    // Parse the JSON
    jzon::parse(&bytes).map_err(|_| PlistParseError::WrongMessageType)
}

#[cfg(test)]
mod tests {
    use std::{env::current_dir, fs::File};

    use plist::Value;

    use crate::message_types::polls::{Poll, PollVote};

    /// Option identifiers from `PollCreate.plist`, in display order.
    const OPTION_1: &str = "630CD0F8-9502-40D4-B2E1-C53E02FD435E";
    const OPTION_2: &str = "58B59D74-EF73-4BA7-AA8E-5A30FBAF3F44";
    const OPTION_3: &str = "D493DDC8-6EE2-408F-888C-3986838A1FBB";
    const OPTION_4: &str = "541DCA6D-686A-4FEB-9092-00BA4A45E2A4";

    fn archive(filename: &str) -> Value {
        let path = current_dir()
            .unwrap()
            .join("test_data/app_message")
            .join(filename);
        Value::from_reader(File::open(path).unwrap()).unwrap()
    }

    fn vote(voter: &str, option_id: &str) -> PollVote {
        PollVote {
            voter: voter.to_string(),
            option_id: option_id.to_string(),
        }
    }

    #[test]
    fn test_parse_poll_creation() {
        let poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();

        assert_eq!(poll.order, vec![OPTION_1, OPTION_2, OPTION_3, OPTION_4]);
        let options: Vec<(&str, &str)> = poll
            .order
            .iter()
            .map(|id| {
                (
                    poll.options[id].text.as_str(),
                    poll.options[id].creator.as_str(),
                )
            })
            .collect();
        assert_eq!(
            options,
            vec![
                ("Option 1", "steven@applenet.com"),
                ("Option 2", "steven@applenet.com"),
                ("Option 3", "steven@applenet.com"),
                ("Option 4", "+12223334444"),
            ]
        );
        assert!(poll.options.values().all(|option| option.votes.is_empty()));
    }

    #[test]
    fn test_parse_poll_votes() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.count_votes(&archive("PollVote.plist")).unwrap();
        poll.count_votes(&archive("PollVote2.plist")).unwrap();

        // Votes accumulate across payloads in arrival order.
        assert_eq!(
            poll.options[OPTION_1].votes,
            vec![vote("steven@applenet.com", OPTION_1)]
        );
        assert_eq!(
            poll.options[OPTION_2].votes,
            vec![
                vote("+12223334444", OPTION_2),
                vote("steven@applenet.com", OPTION_2),
            ]
        );
        assert!(poll.options[OPTION_3].votes.is_empty());
        assert!(poll.options[OPTION_4].votes.is_empty());
    }

    #[test]
    fn test_parse_poll_vote_unknown_option_dropped() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.count_votes(&archive("PollRemovedVote.plist")).unwrap();

        // The fixture carries one vote for option
        // 89913F11-D12C-4A06-BB4B-C26C6C4D1A1F, which `PollCreate.plist` does
        // not define: it was captured from a different poll. It proves only
        // that a vote for an unknown option is dropped; a vote-then-retract
        // capture from one poll is still missing (GitHub issue #823).
        assert!(poll.options.values().all(|option| option.votes.is_empty()));
    }

    #[test]
    fn test_parse_poll_votes_reapplied_snapshot_does_not_double_count() {
        let mut once = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        once.count_votes(&archive("PollVote2.plist")).unwrap();

        let mut twice = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        twice.count_votes(&archive("PollVote2.plist")).unwrap();
        twice.count_votes(&archive("PollVote2.plist")).unwrap();

        assert_eq!(twice, once);
        assert_eq!(
            once.options[OPTION_1].votes,
            vec![vote("steven@applenet.com", OPTION_1)]
        );
        assert_eq!(
            once.options[OPTION_2].votes,
            vec![vote("steven@applenet.com", OPTION_2)]
        );
    }

    #[test]
    fn test_apply_votes_snapshot_replaces_across_options() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![vote("a", OPTION_1)]);
        poll.apply_votes(vec![vote("a", OPTION_2)]);

        assert!(poll.options[OPTION_1].votes.is_empty());
        assert_eq!(poll.options[OPTION_2].votes, vec![vote("a", OPTION_2)]);
    }

    #[test]
    fn test_apply_votes_same_snapshot_twice_counts_once() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![vote("a", OPTION_1), vote("a", OPTION_2)]);
        poll.apply_votes(vec![vote("a", OPTION_1), vote("a", OPTION_2)]);

        assert_eq!(poll.options[OPTION_1].votes, vec![vote("a", OPTION_1)]);
        assert_eq!(poll.options[OPTION_2].votes, vec![vote("a", OPTION_2)]);
    }

    #[test]
    fn test_apply_votes_leaves_other_participants_in_order() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![vote("a", OPTION_1)]);
        poll.apply_votes(vec![vote("b", OPTION_1)]);
        poll.apply_votes(vec![vote("c", OPTION_1)]);
        // `b` re-votes: `a` and `c` keep their order, `b` moves to the end.
        poll.apply_votes(vec![vote("b", OPTION_1), vote("b", OPTION_2)]);

        assert_eq!(
            poll.options[OPTION_1].votes,
            vec![
                vote("a", OPTION_1),
                vote("c", OPTION_1),
                vote("b", OPTION_1)
            ]
        );
        assert_eq!(poll.options[OPTION_2].votes, vec![vote("b", OPTION_2)]);
    }

    #[test]
    fn test_apply_votes_partial_retraction_removes_absent_option() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![vote("a", OPTION_1), vote("a", OPTION_3)]);
        poll.apply_votes(vec![vote("a", OPTION_3)]);

        assert!(poll.options[OPTION_1].votes.is_empty());
        assert!(poll.options[OPTION_2].votes.is_empty());
        assert_eq!(poll.options[OPTION_3].votes, vec![vote("a", OPTION_3)]);
        assert!(poll.options[OPTION_4].votes.is_empty());
    }

    #[test]
    fn test_apply_votes_duplicate_pair_collapses() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![
            vote("a", OPTION_1),
            vote("a", OPTION_1),
            vote("a", OPTION_2),
        ]);

        assert_eq!(poll.options[OPTION_1].votes, vec![vote("a", OPTION_1)]);
        assert_eq!(poll.options[OPTION_2].votes, vec![vote("a", OPTION_2)]);
    }

    #[test]
    fn test_apply_votes_empty_snapshot_is_no_op() {
        let mut poll = Poll::from_payload(&archive("PollCreate.plist")).unwrap();
        poll.apply_votes(vec![vote("a", OPTION_1)]);
        poll.apply_votes(Vec::new());

        // Names no participant, so it cannot retract (GitHub issue #823).
        assert_eq!(poll.options[OPTION_1].votes, vec![vote("a", OPTION_1)]);
    }
}