Skip to main content

ios_core/services/power_assertion/
mod.rs

1//! Power assertion service client.
2//!
3//! Service: `com.apple.mobile.assertion_agent`
4
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
6
7pub const SERVICE_NAME: &str = "com.apple.mobile.assertion_agent";
8
9service_error!(PowerAssertionError);
10
11pub struct PowerAssertionClient<S> {
12    stream: S,
13}
14
15impl<S: AsyncRead + AsyncWrite + Unpin> PowerAssertionClient<S> {
16    pub fn new(stream: S) -> Self {
17        Self { stream }
18    }
19
20    pub async fn create_assertion(
21        &mut self,
22        assertion_type: &str,
23        name: &str,
24        timeout_seconds: f64,
25        details: Option<&str>,
26    ) -> Result<plist::Dictionary, PowerAssertionError> {
27        let mut request = plist::Dictionary::from_iter([
28            (
29                "CommandKey".to_string(),
30                plist::Value::String("CommandCreateAssertion".into()),
31            ),
32            (
33                "AssertionTypeKey".to_string(),
34                plist::Value::String(assertion_type.to_string()),
35            ),
36            (
37                "AssertionNameKey".to_string(),
38                plist::Value::String(name.to_string()),
39            ),
40            (
41                "AssertionTimeoutKey".to_string(),
42                plist::Value::Real(timeout_seconds),
43            ),
44        ]);
45        if let Some(details) = details {
46            request.insert(
47                "AssertionDetailKey".to_string(),
48                plist::Value::String(details.to_string()),
49            );
50        }
51
52        send_plist(&mut self.stream, &plist::Value::Dictionary(request)).await?;
53        let response = recv_plist(&mut self.stream).await?;
54        if let Some(error) = response.get("Error").and_then(plist::Value::as_string) {
55            return Err(PowerAssertionError::Protocol(error.to_string()));
56        }
57        Ok(response)
58    }
59}
60
61async fn send_plist<S: AsyncWrite + Unpin>(
62    stream: &mut S,
63    value: &plist::Value,
64) -> Result<(), PowerAssertionError> {
65    let mut buf = Vec::new();
66    plist::to_writer_xml(&mut buf, value)?;
67    stream.write_all(&(buf.len() as u32).to_be_bytes()).await?;
68    stream.write_all(&buf).await?;
69    stream.flush().await?;
70    Ok(())
71}
72
73async fn recv_plist<S: AsyncRead + Unpin>(
74    stream: &mut S,
75) -> Result<plist::Dictionary, PowerAssertionError> {
76    let mut len_buf = [0u8; 4];
77    stream.read_exact(&mut len_buf).await?;
78    let len = u32::from_be_bytes(len_buf) as usize;
79    const MAX_PLIST_SIZE: usize = 1024 * 1024;
80    if len > MAX_PLIST_SIZE {
81        return Err(PowerAssertionError::Protocol(format!(
82            "plist length {len} exceeds max {MAX_PLIST_SIZE}"
83        )));
84    }
85    let mut buf = vec![0u8; len];
86    stream.read_exact(&mut buf).await?;
87    Ok(plist::from_bytes(&buf)?)
88}
89
90#[cfg(test)]
91mod tests {
92    use crate::test_util::MockStream;
93
94    use super::*;
95
96    #[tokio::test]
97    async fn create_assertion_sends_expected_payload() {
98        let response = plist::Value::Dictionary(plist::Dictionary::new());
99        let mut stream = MockStream::with_response(response);
100        let mut client = PowerAssertionClient::new(&mut stream);
101
102        client
103            .create_assertion("PreventUserIdleSystemSleep", "ios-cli", 30.0, Some("test"))
104            .await
105            .unwrap();
106
107        let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
108        let payload = &stream.written[4..4 + len];
109        let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
110        assert_eq!(
111            dict.get("CommandKey").and_then(plist::Value::as_string),
112            Some("CommandCreateAssertion")
113        );
114        assert_eq!(
115            dict.get("AssertionTypeKey")
116                .and_then(plist::Value::as_string),
117            Some("PreventUserIdleSystemSleep")
118        );
119        assert_eq!(
120            dict.get("AssertionNameKey")
121                .and_then(plist::Value::as_string),
122            Some("ios-cli")
123        );
124        assert_eq!(
125            dict.get("AssertionDetailKey")
126                .and_then(plist::Value::as_string),
127            Some("test")
128        );
129    }
130}