Skip to main content

ios_core/services/instruments/
notifications.rs

1use tokio::io::{AsyncRead, AsyncWrite};
2
3use crate::services::dtx::codec::{DtxConnection, DtxError};
4use crate::services::dtx::primitive_enc::archived_object;
5use crate::services::dtx::types::{DtxMessage, DtxPayload, NSObject};
6
7#[derive(Debug, Clone, PartialEq)]
8pub struct NotificationEvent {
9    pub selector: String,
10    pub payload: NSObject,
11    pub channel_code: i32,
12}
13
14pub struct NotificationClient<S> {
15    conn: DtxConnection<S>,
16}
17
18impl<S: AsyncRead + AsyncWrite + Unpin + Send> NotificationClient<S> {
19    pub async fn connect(stream: S) -> Result<Self, DtxError> {
20        let mut conn = DtxConnection::new(stream);
21        let channel_code = conn
22            .request_channel(super::MOBILE_NOTIFICATIONS_SVC)
23            .await?;
24
25        conn.method_call(
26            channel_code,
27            "setApplicationStateNotificationsEnabled:",
28            &[archived_object(
29                crate::proto::nskeyedarchiver_encode::archive_bool(true),
30            )],
31        )
32        .await?;
33        conn.method_call(
34            channel_code,
35            "setMemoryNotificationsEnabled:",
36            &[archived_object(
37                crate::proto::nskeyedarchiver_encode::archive_bool(true),
38            )],
39        )
40        .await?;
41
42        Ok(Self { conn })
43    }
44
45    pub async fn next_notification(&mut self) -> Result<NotificationEvent, DtxError> {
46        loop {
47            let msg = self.conn.recv().await?;
48            if msg.expects_reply {
49                self.conn.send_ack(&msg).await?;
50            }
51            if let Some(event) = parse_notification_message(&msg) {
52                return Ok(event);
53            }
54        }
55    }
56}
57
58fn parse_notification_message(msg: &DtxMessage) -> Option<NotificationEvent> {
59    let (selector, args) = match &msg.payload {
60        DtxPayload::MethodInvocation { selector, args } => (selector, args),
61        _ => return None,
62    };
63
64    if selector != "applicationStateNotification:" && selector != "memoryNotification:" {
65        return None;
66    }
67
68    let payload = args.first().cloned().unwrap_or(NSObject::Null);
69    Some(NotificationEvent {
70        selector: selector.clone(),
71        payload,
72        channel_code: msg.channel_code,
73    })
74}
75
76#[cfg(test)]
77mod tests {
78    use indexmap::IndexMap;
79
80    use super::*;
81
82    #[test]
83    fn parses_application_state_notification_payload() {
84        let msg = DtxMessage {
85            identifier: 11,
86            conversation_idx: 0,
87            channel_code: 7,
88            expects_reply: false,
89            payload: DtxPayload::MethodInvocation {
90                selector: "applicationStateNotification:".into(),
91                args: vec![NSObject::Dict(IndexMap::from_iter([
92                    (
93                        "ApplicationBundleIdentifier".into(),
94                        NSObject::String("com.apple.Preferences".into()),
95                    ),
96                    ("State".into(), NSObject::Int(8)),
97                ]))],
98            },
99        };
100
101        let event = parse_notification_message(&msg).expect("notification");
102        assert_eq!(event.selector, "applicationStateNotification:");
103        assert_eq!(event.channel_code, 7);
104        match event.payload {
105            NSObject::Dict(payload) => {
106                assert_eq!(
107                    payload.get("ApplicationBundleIdentifier"),
108                    Some(&NSObject::String("com.apple.Preferences".into()))
109                );
110                assert_eq!(payload.get("State"), Some(&NSObject::Int(8)));
111            }
112            other => panic!("unexpected payload: {other:?}"),
113        }
114    }
115
116    #[test]
117    fn ignores_non_notification_messages() {
118        let msg = DtxMessage {
119            identifier: 12,
120            conversation_idx: 0,
121            channel_code: 3,
122            expects_reply: false,
123            payload: DtxPayload::MethodInvocation {
124                selector: "runningProcesses".into(),
125                args: vec![],
126            },
127        };
128
129        assert!(parse_notification_message(&msg).is_none());
130    }
131}