ios_core/services/instruments/
tap.rs

1use bytes::Bytes;
2use tokio::io::{AsyncRead, AsyncWrite};
3
4use super::unarchive_raw_payload;
5use crate::services::dtx::codec::{DtxConnection, DtxError};
6use crate::services::dtx::primitive_enc::archived_object;
7use crate::services::dtx::types::{DtxPayload, NSObject};
8
9#[derive(Debug, Clone)]
10pub enum TapMessage {
11    Data(Bytes),
12    Plist(NSObject),
13}
14
15pub struct TapClient<S> {
16    conn: DtxConnection<S>,
17    channel_code: i32,
18}
19
20impl<S: AsyncRead + AsyncWrite + Unpin + Send> TapClient<S> {
21    pub async fn connect(
22        stream: S,
23        service_name: &str,
24        config: Vec<(String, plist::Value)>,
25    ) -> Result<Self, DtxError> {
26        let mut conn = DtxConnection::new(stream);
27        let channel_code = conn.request_channel(service_name).await?;
28        conn.method_call_async(
29            channel_code,
30            "setConfig:",
31            &[archived_object(
32                crate::proto::nskeyedarchiver_encode::archive_dict(config),
33            )],
34        )
35        .await?;
36        conn.method_call_async(channel_code, "start", &[]).await?;
37
38        let mut client = Self { conn, channel_code };
39        client.wait_for_start_message().await?;
40        Ok(client)
41    }
42
43    pub async fn next_message(&mut self) -> Result<TapMessage, DtxError> {
44        loop {
45            let msg = self.conn.recv().await?;
46            if msg.expects_reply {
47                self.conn.send_ack(&msg).await?;
48            }
49            if msg.channel_code != self.channel_code && msg.channel_code != -1 {
50                continue;
51            }
52            if let Some(message) = payload_to_tap_message(msg.payload) {
53                return Ok(message);
54            }
55        }
56    }
57
58    pub async fn stop(&mut self) -> Result<(), DtxError> {
59        self.conn
60            .method_call_async(self.channel_code, "stop", &[])
61            .await
62    }
63
64    async fn wait_for_start_message(&mut self) -> Result<(), DtxError> {
65        loop {
66            match self.next_message().await? {
67                TapMessage::Plist(_) => return Ok(()),
68                TapMessage::Data(bytes) if unarchive_raw_payload(&bytes).is_some() => return Ok(()),
69                TapMessage::Data(_) => {}
70            }
71        }
72    }
73}
74
75fn payload_to_tap_message(payload: DtxPayload) -> Option<TapMessage> {
76    match payload {
77        DtxPayload::Raw(bytes) => Some(TapMessage::Data(bytes)),
78        DtxPayload::RawWithAux { payload, aux } => aux
79            .into_iter()
80            .find_map(|value| match value {
81                NSObject::Data(bytes) => Some(TapMessage::Data(bytes)),
82                _ => None,
83            })
84            .or_else(|| (!payload.is_empty()).then_some(TapMessage::Data(payload))),
85        DtxPayload::Response(value) => Some(TapMessage::Plist(value)),
86        DtxPayload::MethodInvocation { args, .. } => Some(TapMessage::Plist(if args.len() == 1 {
87            args.into_iter().next().unwrap_or(NSObject::Null)
88        } else {
89            NSObject::Array(args)
90        })),
91        DtxPayload::Notification { object, .. } => Some(TapMessage::Plist(object)),
92        DtxPayload::Empty => None,
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use tokio::io::{duplex, AsyncWriteExt};
99
100    use super::*;
101    use crate::services::dtx::{encode_dtx, read_dtx_frame, DtxPayload, NSObject};
102
103    const MSG_RESPONSE: u32 = 3;
104    const MSG_UNKNOWN_TYPE_ONE: u32 = 1;
105
106    #[test]
107    fn payload_to_tap_message_prefers_auxiliary_data_for_raw_with_aux() {
108        let payload = DtxPayload::RawWithAux {
109            payload: bytes::Bytes::from_static(b"body"),
110            aux: vec![NSObject::Data(bytes::Bytes::from_static(b"aux"))],
111        };
112
113        let message = payload_to_tap_message(payload).expect("tap message");
114        match message {
115            TapMessage::Data(data) => assert_eq!(data.as_ref(), b"aux"),
116            other => panic!("unexpected tap message: {other:?}"),
117        }
118    }
119
120    #[test]
121    fn payload_to_tap_message_wraps_multiple_method_invocation_args_as_array() {
122        let message = payload_to_tap_message(DtxPayload::MethodInvocation {
123            selector: "event".into(),
124            args: vec![NSObject::String("a".into()), NSObject::Int(2)],
125        })
126        .expect("tap message");
127
128        match message {
129            TapMessage::Plist(NSObject::Array(values)) => {
130                assert_eq!(values, vec![NSObject::String("a".into()), NSObject::Int(2)]);
131            }
132            other => panic!("unexpected tap plist payload: {other:?}"),
133        }
134    }
135
136    #[tokio::test]
137    async fn connect_sends_config_and_start_then_waits_for_plist_event() {
138        let (client, mut server) = duplex(4096);
139        let task = tokio::spawn(async move {
140            TapClient::connect(
141                client,
142                "com.apple.instruments.server.services.exampletap",
143                vec![("interval".to_string(), plist::Value::Integer(100.into()))],
144            )
145            .await
146            .unwrap()
147        });
148
149        let channel_request = read_dtx_frame(&mut server).await.unwrap();
150        match channel_request.payload {
151            DtxPayload::MethodInvocation { selector, args } => {
152                assert_eq!(selector, "_requestChannelWithCode:identifier:");
153                assert!(matches!(
154                    args.get(1),
155                    Some(NSObject::String(name))
156                    if name == "com.apple.instruments.server.services.exampletap"
157                ));
158            }
159            other => panic!("unexpected channel request: {other:?}"),
160        }
161        server
162            .write_all(&encode_dtx(
163                channel_request.identifier,
164                1,
165                0,
166                false,
167                MSG_RESPONSE,
168                &[],
169                &[],
170            ))
171            .await
172            .unwrap();
173
174        let set_config = read_dtx_frame(&mut server).await.unwrap();
175        match set_config.payload {
176            DtxPayload::MethodInvocation { selector, args } => {
177                assert_eq!(selector, "setConfig:");
178                assert!(matches!(
179                    args.first(),
180                    Some(NSObject::Dict(dict))
181                    if dict.get("interval") == Some(&NSObject::Int(100))
182                ));
183            }
184            other => panic!("unexpected setConfig request: {other:?}"),
185        }
186
187        let start = read_dtx_frame(&mut server).await.unwrap();
188        match start.payload {
189            DtxPayload::MethodInvocation { selector, args } => {
190                assert_eq!(selector, "start");
191                assert!(args.is_empty());
192            }
193            other => panic!("unexpected start request: {other:?}"),
194        }
195
196        let start_notice = crate::proto::nskeyedarchiver_encode::archive_string("started");
197        server
198            .write_all(&encode_dtx(
199                99,
200                0,
201                -set_config.channel_code,
202                false,
203                MSG_RESPONSE,
204                &start_notice,
205                &[],
206            ))
207            .await
208            .unwrap();
209
210        let mut client = task.await.unwrap();
211        server
212            .write_all(&encode_dtx(
213                100,
214                0,
215                -start.channel_code,
216                false,
217                MSG_UNKNOWN_TYPE_ONE,
218                b"trace-bytes",
219                &[],
220            ))
221            .await
222            .unwrap();
223
224        match client.next_message().await.unwrap() {
225            TapMessage::Data(bytes) => assert_eq!(bytes.as_ref(), b"trace-bytes"),
226            other => panic!("unexpected live tap message: {other:?}"),
227        }
228    }
229
230    #[tokio::test]
231    async fn connect_accepts_raw_archived_plist_as_start_ack() {
232        let (client, mut server) = duplex(4096);
233        let task = tokio::spawn(async move {
234            TapClient::connect(
235                client,
236                "com.apple.instruments.server.services.exampletap",
237                vec![("interval".to_string(), plist::Value::Integer(100.into()))],
238            )
239            .await
240            .unwrap()
241        });
242
243        let channel_request = read_dtx_frame(&mut server).await.unwrap();
244        server
245            .write_all(&encode_dtx(
246                channel_request.identifier,
247                1,
248                0,
249                false,
250                MSG_RESPONSE,
251                &[],
252                &[],
253            ))
254            .await
255            .unwrap();
256
257        let set_config = read_dtx_frame(&mut server).await.unwrap();
258        match set_config.payload {
259            DtxPayload::MethodInvocation { selector, .. } => {
260                assert_eq!(selector, "setConfig:");
261            }
262            other => panic!("unexpected setConfig request: {other:?}"),
263        }
264
265        let start = read_dtx_frame(&mut server).await.unwrap();
266        match start.payload {
267            DtxPayload::MethodInvocation { selector, args } => {
268                assert_eq!(selector, "start");
269                assert!(args.is_empty());
270            }
271            other => panic!("unexpected start request: {other:?}"),
272        }
273
274        let start_notice = crate::proto::nskeyedarchiver_encode::archive_string("started");
275        server
276            .write_all(&encode_dtx(
277                99,
278                0,
279                -start.channel_code,
280                false,
281                MSG_UNKNOWN_TYPE_ONE,
282                &start_notice,
283                &[],
284            ))
285            .await
286            .unwrap();
287
288        let mut client = task.await.unwrap();
289        server
290            .write_all(&encode_dtx(
291                100,
292                0,
293                -start.channel_code,
294                false,
295                MSG_UNKNOWN_TYPE_ONE,
296                b"trace-bytes",
297                &[],
298            ))
299            .await
300            .unwrap();
301
302        match client.next_message().await.unwrap() {
303            TapMessage::Data(bytes) => assert_eq!(bytes.as_ref(), b"trace-bytes"),
304            other => panic!("unexpected live tap message: {other:?}"),
305        }
306    }
307}
ios_core/services/instruments/tap.rs

ios_core/services/instruments/
tap.rs
