Skip to main content

ios_core/services/heartbeat/
mod.rs

1//! Minimal heartbeat service client.
2//!
3//! Service: `com.apple.mobile.heartbeat`
4
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
6
7pub const SERVICE_NAME: &str = "com.apple.mobile.heartbeat";
8
9service_error!(HeartbeatError);
10
11pub struct HeartbeatClient<S> {
12    stream: S,
13}
14
15impl<S: AsyncRead + AsyncWrite + Unpin> HeartbeatClient<S> {
16    pub fn new(stream: S) -> Self {
17        Self { stream }
18    }
19
20    pub async fn recv_message(&mut self) -> Result<plist::Value, HeartbeatError> {
21        recv_plist(&mut self.stream).await
22    }
23
24    pub async fn send_polo(&mut self) -> Result<(), HeartbeatError> {
25        send_plist(
26            &mut self.stream,
27            &plist::Value::Dictionary(plist::Dictionary::from_iter([(
28                "Command".to_string(),
29                plist::Value::String("Polo".into()),
30            )])),
31        )
32        .await
33    }
34
35    pub async fn ping(&mut self) -> Result<plist::Value, HeartbeatError> {
36        let message = self.recv_message().await?;
37        self.send_polo().await?;
38        Ok(message)
39    }
40}
41
42async fn send_plist<S>(stream: &mut S, value: &plist::Value) -> Result<(), HeartbeatError>
43where
44    S: AsyncWrite + Unpin,
45{
46    let mut buf = Vec::new();
47    plist::to_writer_xml(&mut buf, value)?;
48    stream.write_all(&(buf.len() as u32).to_be_bytes()).await?;
49    stream.write_all(&buf).await?;
50    stream.flush().await?;
51    Ok(())
52}
53
54async fn recv_plist<S>(stream: &mut S) -> Result<plist::Value, HeartbeatError>
55where
56    S: AsyncRead + Unpin,
57{
58    let mut len_buf = [0u8; 4];
59    stream.read_exact(&mut len_buf).await?;
60    let len = u32::from_be_bytes(len_buf) as usize;
61    const MAX_PLIST_SIZE: usize = 1024 * 1024;
62    if len > MAX_PLIST_SIZE {
63        return Err(HeartbeatError::Protocol(format!(
64            "plist length {len} exceeds max {MAX_PLIST_SIZE}"
65        )));
66    }
67    let mut buf = vec![0u8; len];
68    stream.read_exact(&mut buf).await?;
69    Ok(plist::from_bytes(&buf)?)
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[tokio::test]
77    async fn ping_reads_message_and_sends_polo() {
78        let (client_side, mut server_side) = tokio::io::duplex(4096);
79
80        tokio::spawn(async move {
81            let incoming = plist::Value::Dictionary(plist::Dictionary::from_iter([(
82                "Command".to_string(),
83                plist::Value::String("Marco".into()),
84            )]));
85            let mut buf = Vec::new();
86            plist::to_writer_xml(&mut buf, &incoming).unwrap();
87            server_side
88                .write_all(&(buf.len() as u32).to_be_bytes())
89                .await
90                .unwrap();
91            server_side.write_all(&buf).await.unwrap();
92
93            let mut len_buf = [0u8; 4];
94            server_side.read_exact(&mut len_buf).await.unwrap();
95            let len = u32::from_be_bytes(len_buf) as usize;
96            let mut payload = vec![0u8; len];
97            server_side.read_exact(&mut payload).await.unwrap();
98            let response: plist::Value = plist::from_bytes(&payload).unwrap();
99            let dict = response.into_dictionary().unwrap();
100            assert_eq!(dict["Command"].as_string(), Some("Polo"));
101        });
102
103        let mut client = HeartbeatClient::new(client_side);
104        let message = client.ping().await.unwrap();
105        let dict = message.into_dictionary().unwrap();
106        assert_eq!(dict["Command"].as_string(), Some("Marco"));
107    }
108}