1use serde::{Deserialize, Serialize};
9use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10
11pub const SERVICE_NAME: &str = "com.apple.mobile.diagnostics_relay";
12
13#[derive(Debug, thiserror::Error)]
14pub enum DiagnosticsError {
15 #[error("IO error: {0}")]
16 Io(#[from] std::io::Error),
17 #[error("plist error: {0}")]
18 Plist(String),
19 #[error("mobilegestalt deprecated: {0}")]
20 Deprecated(String),
21 #[error("protocol error: {0}")]
22 Protocol(String),
23}
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "PascalCase")]
27pub struct BatteryDiagnostics {
28 #[serde(default)]
29 pub instant_amperage: Option<i64>,
30 #[serde(default)]
31 pub temperature: Option<i64>,
32 #[serde(default)]
33 pub voltage: Option<i64>,
34 #[serde(default)]
35 pub is_charging: Option<bool>,
36 #[serde(default)]
37 pub current_capacity: Option<i64>,
38 #[serde(default)]
39 pub design_capacity: Option<u64>,
40 #[serde(default)]
41 pub nominal_charge_capacity: Option<u64>,
42 #[serde(default)]
43 pub absolute_capacity: Option<u64>,
44 #[serde(default)]
45 pub apple_raw_current_capacity: Option<u64>,
46 #[serde(default)]
47 pub apple_raw_max_capacity: Option<u64>,
48 #[serde(default)]
49 pub cycle_count: Option<u64>,
50 #[serde(default)]
51 pub at_critical_level: Option<bool>,
52 #[serde(default)]
53 pub at_warn_level: Option<bool>,
54}
55
56#[derive(Serialize)]
57#[serde(rename_all = "PascalCase")]
58struct MobileGestaltRequest<'a> {
59 request: &'static str,
60 #[serde(rename = "MobileGestaltKeys")]
61 keys: Vec<&'a str>,
62}
63
64#[derive(Serialize)]
65#[serde(rename_all = "PascalCase")]
66struct IoRegistryRequest<'a> {
67 request: &'static str,
68 #[serde(skip_serializing_if = "Option::is_none")]
69 entry_class: Option<&'a str>,
70 #[serde(skip_serializing_if = "Option::is_none")]
71 entry_name: Option<&'a str>,
72 #[serde(skip_serializing_if = "Option::is_none")]
73 current_plane: Option<&'a str>,
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub struct IoRegistryQuery<'a> {
78 pub entry_class: Option<&'a str>,
79 pub entry_name: Option<&'a str>,
80 pub current_plane: Option<&'a str>,
81}
82
83impl<'a> IoRegistryQuery<'a> {
84 pub fn by_class(entry_class: &'a str) -> Self {
85 Self {
86 entry_class: Some(entry_class),
87 entry_name: None,
88 current_plane: None,
89 }
90 }
91}
92
93pub async fn query_mobile_gestalt<S>(
95 stream: &mut S,
96 keys: &[&str],
97) -> Result<plist::Value, DiagnosticsError>
98where
99 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
100{
101 send_plist(
102 stream,
103 &MobileGestaltRequest {
104 request: "MobileGestalt",
105 keys: keys.to_vec(),
106 },
107 )
108 .await?;
109
110 let response = recv_response_dict(stream).await?;
111 let diagnostics = extract_diagnostics_payload(&response)?;
112
113 extract_mobile_gestalt_payload(diagnostics)
114}
115
116pub async fn query_all_values<S>(stream: &mut S) -> Result<plist::Value, DiagnosticsError>
118where
119 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
120{
121 #[derive(Serialize)]
122 #[serde(rename_all = "PascalCase")]
123 struct AllRequest {
124 request: &'static str,
125 }
126
127 send_plist(stream, &AllRequest { request: "All" }).await?;
128 let response = recv_response_dict(stream).await?;
129 extract_diagnostics_payload(&response)
130}
131
132pub async fn query_battery<S>(stream: &mut S) -> Result<BatteryDiagnostics, DiagnosticsError>
134where
135 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
136{
137 let io_registry = query_ioregistry(stream, "IOPMPowerSource").await?;
138 plist::from_value(&io_registry).map_err(|e| DiagnosticsError::Plist(e.to_string()))
139}
140
141pub async fn query_ioregistry<S>(
143 stream: &mut S,
144 entry_class: &str,
145) -> Result<plist::Value, DiagnosticsError>
146where
147 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
148{
149 query_ioregistry_with(stream, IoRegistryQuery::by_class(entry_class)).await
150}
151
152pub async fn query_ioregistry_with<S>(
153 stream: &mut S,
154 query: IoRegistryQuery<'_>,
155) -> Result<plist::Value, DiagnosticsError>
156where
157 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
158{
159 if query.entry_class.is_none() && query.entry_name.is_none() {
160 return Err(DiagnosticsError::Protocol(
161 "IORegistry query requires EntryClass or EntryName".into(),
162 ));
163 }
164
165 send_plist(
166 stream,
167 &IoRegistryRequest {
168 request: "IORegistry",
169 entry_class: query.entry_class,
170 entry_name: query.entry_name,
171 current_plane: query.current_plane,
172 },
173 )
174 .await?;
175
176 let response = recv_response_dict(stream).await?;
177 let diagnostics = extract_diagnostics_payload(&response)?;
178 let dict = diagnostics.into_dictionary().ok_or_else(|| {
179 DiagnosticsError::Protocol("diagnostics payload was not a dictionary".into())
180 })?;
181 let io_registry = dict
182 .get("IORegistry")
183 .cloned()
184 .ok_or_else(|| DiagnosticsError::Protocol("diagnostics missing IORegistry".into()))?;
185 Ok(io_registry)
186}
187
188pub async fn reboot<S>(stream: &mut S) -> Result<(), DiagnosticsError>
190where
191 S: AsyncRead + AsyncWrite + Unpin + ?Sized,
192{
193 #[derive(Serialize)]
194 #[serde(rename_all = "PascalCase")]
195 struct Request {
196 request: &'static str,
197 }
198 send_plist(stream, &Request { request: "Restart" }).await?;
199 recv_plist_raw(stream).await?;
200 Ok(())
201}
202
203async fn send_plist<S, T>(stream: &mut S, value: &T) -> Result<(), DiagnosticsError>
206where
207 S: AsyncWrite + Unpin + ?Sized,
208 T: Serialize,
209{
210 let mut buf = Vec::new();
211 plist::to_writer_xml(&mut buf, value).map_err(|e| DiagnosticsError::Plist(e.to_string()))?;
212 stream.write_all(&(buf.len() as u32).to_be_bytes()).await?;
213 stream.write_all(&buf).await?;
214 stream.flush().await?;
215 Ok(())
216}
217
218async fn recv_plist_raw<S>(stream: &mut S) -> Result<Vec<u8>, DiagnosticsError>
219where
220 S: AsyncRead + Unpin + ?Sized,
221{
222 let mut len_buf = [0u8; 4];
223 stream.read_exact(&mut len_buf).await?;
224 let len = u32::from_be_bytes(len_buf) as usize;
225 const MAX_PLIST_SIZE: usize = 4 * 1024 * 1024;
227 if len > MAX_PLIST_SIZE {
228 return Err(DiagnosticsError::Protocol(format!(
229 "plist length {len} exceeds max {MAX_PLIST_SIZE}"
230 )));
231 }
232 let mut buf = vec![0u8; len];
233 stream.read_exact(&mut buf).await?;
234 Ok(buf)
235}
236
237async fn recv_response_dict<S>(stream: &mut S) -> Result<plist::Dictionary, DiagnosticsError>
238where
239 S: AsyncRead + Unpin + ?Sized,
240{
241 let data = recv_plist_raw(stream).await?;
242 let value: plist::Value =
243 plist::from_bytes(&data).map_err(|e| DiagnosticsError::Plist(e.to_string()))?;
244 value.into_dictionary().ok_or_else(|| {
245 DiagnosticsError::Protocol("diagnostics response payload was not a dictionary".into())
246 })
247}
248
249fn extract_diagnostics_payload(
250 response: &plist::Dictionary,
251) -> Result<plist::Value, DiagnosticsError> {
252 response
253 .get("Diagnostics")
254 .cloned()
255 .ok_or_else(|| missing_diagnostics_error(response))
256}
257
258fn missing_diagnostics_error(response: &plist::Dictionary) -> DiagnosticsError {
259 let status = response
260 .get("Status")
261 .and_then(plist::Value::as_string)
262 .map(|value| format!(" (Status={value})"))
263 .unwrap_or_default();
264 let rendered = render_plist_value(&plist::Value::Dictionary(response.clone()));
265 DiagnosticsError::Protocol(format!(
266 "diagnostics response missing Diagnostics{status}: {rendered}"
267 ))
268}
269
270fn render_plist_value(value: &plist::Value) -> String {
271 serde_json::to_string(&plist_to_json(value))
272 .unwrap_or_else(|_| "<failed to render plist value>".to_string())
273}
274
275fn plist_to_json(value: &plist::Value) -> serde_json::Value {
276 match value {
277 plist::Value::Array(items) => {
278 serde_json::Value::Array(items.iter().map(plist_to_json).collect())
279 }
280 plist::Value::Boolean(value) => serde_json::Value::Bool(*value),
281 plist::Value::Data(bytes) => {
282 serde_json::Value::Array(bytes.iter().copied().map(serde_json::Value::from).collect())
283 }
284 plist::Value::Date(value) => serde_json::Value::String(value.to_xml_format()),
285 plist::Value::Dictionary(dict) => serde_json::Value::Object(
286 dict.iter()
287 .map(|(key, value)| (key.clone(), plist_to_json(value)))
288 .collect(),
289 ),
290 plist::Value::Integer(value) => value
291 .as_signed()
292 .map(serde_json::Value::from)
293 .or_else(|| value.as_unsigned().map(serde_json::Value::from))
294 .unwrap_or(serde_json::Value::Null),
295 plist::Value::Real(value) => serde_json::Value::from(*value),
296 plist::Value::String(value) => serde_json::Value::String(value.clone()),
297 plist::Value::Uid(value) => serde_json::Value::from(value.get()),
298 _ => serde_json::Value::Null,
299 }
300}
301
302fn extract_mobile_gestalt_payload(
303 diagnostics: plist::Value,
304) -> Result<plist::Value, DiagnosticsError> {
305 let Some(dict) = diagnostics.as_dictionary() else {
306 return Ok(diagnostics);
307 };
308
309 let Some(mobile_gestalt) = dict.get("MobileGestalt") else {
310 return Ok(diagnostics);
311 };
312
313 let mut inner = mobile_gestalt.as_dictionary().cloned().ok_or_else(|| {
314 DiagnosticsError::Protocol("MobileGestalt payload was not a dictionary".into())
315 })?;
316
317 if let Some(status) = inner.get("Status").and_then(|value| value.as_string()) {
318 match status {
319 "Success" => {
320 inner.remove("Status");
321 }
322 "MobileGestaltDeprecated" => {
323 return Err(DiagnosticsError::Deprecated(
324 "diagnostics relay reports MobileGestaltDeprecated on this OS".into(),
325 ));
326 }
327 other => {
328 return Err(DiagnosticsError::Protocol(format!(
329 "unexpected MobileGestalt status: {other}"
330 )));
331 }
332 }
333 }
334
335 Ok(plist::Value::Dictionary(inner))
336}
337
338#[cfg(test)]
339mod tests {
340 use std::pin::Pin;
341 use std::task::{Context, Poll};
342
343 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
344
345 use super::*;
346
347 #[derive(Default)]
348 struct MockStream {
349 read_buf: Vec<u8>,
350 written: Vec<u8>,
351 read_pos: usize,
352 }
353
354 impl MockStream {
355 fn with_response(value: plist::Value) -> Self {
356 let mut payload = Vec::new();
357 plist::to_writer_xml(&mut payload, &value).unwrap();
358 let mut read_buf = Vec::new();
359 read_buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
360 read_buf.extend_from_slice(&payload);
361 Self {
362 read_buf,
363 written: Vec::new(),
364 read_pos: 0,
365 }
366 }
367 }
368
369 impl AsyncRead for MockStream {
370 fn poll_read(
371 mut self: Pin<&mut Self>,
372 _cx: &mut Context<'_>,
373 buf: &mut ReadBuf<'_>,
374 ) -> Poll<std::io::Result<()>> {
375 let remaining = self.read_buf.len().saturating_sub(self.read_pos);
376 if remaining == 0 {
377 return Poll::Ready(Err(std::io::Error::new(
378 std::io::ErrorKind::UnexpectedEof,
379 "no more test data",
380 )));
381 }
382 let to_copy = remaining.min(buf.remaining());
383 let start = self.read_pos;
384 let end = start + to_copy;
385 buf.put_slice(&self.read_buf[start..end]);
386 self.read_pos = end;
387 Poll::Ready(Ok(()))
388 }
389 }
390
391 impl AsyncWrite for MockStream {
392 fn poll_write(
393 mut self: Pin<&mut Self>,
394 _cx: &mut Context<'_>,
395 buf: &[u8],
396 ) -> Poll<std::io::Result<usize>> {
397 self.written.extend_from_slice(buf);
398 Poll::Ready(Ok(buf.len()))
399 }
400
401 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
402 Poll::Ready(Ok(()))
403 }
404
405 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
406 Poll::Ready(Ok(()))
407 }
408 }
409
410 #[tokio::test]
411 async fn query_mobile_gestalt_uses_mobile_gestalt_keys_field() {
412 let mut stream =
413 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
414 "Diagnostics".to_string(),
415 plist::Value::Dictionary(plist::Dictionary::new()),
416 )])));
417
418 let _ = query_mobile_gestalt(&mut stream, &["ProductVersion"])
419 .await
420 .unwrap();
421
422 let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
423 let payload = &stream.written[4..4 + len];
424 let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
425 assert_eq!(dict["Request"].as_string(), Some("MobileGestalt"));
426 let keys = dict["MobileGestaltKeys"].as_array().unwrap();
427 assert_eq!(keys.len(), 1);
428 assert_eq!(keys[0].as_string(), Some("ProductVersion"));
429 assert!(!dict.contains_key("Keys"));
430 }
431
432 #[tokio::test]
433 async fn query_mobile_gestalt_returns_diagnostics_payload() {
434 let mut stream =
435 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
436 "Diagnostics".to_string(),
437 plist::Value::Dictionary(plist::Dictionary::from_iter([(
438 "ProductVersion".to_string(),
439 plist::Value::String("26.0".into()),
440 )])),
441 )])));
442
443 let value = query_mobile_gestalt(&mut stream, &["ProductVersion"])
444 .await
445 .unwrap();
446 let dict = value.into_dictionary().unwrap();
447 assert_eq!(dict["ProductVersion"].as_string(), Some("26.0"));
448 }
449
450 #[tokio::test]
451 async fn query_mobile_gestalt_strips_nested_success_status() {
452 let mut stream =
453 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
454 "Diagnostics".to_string(),
455 plist::Value::Dictionary(plist::Dictionary::from_iter([(
456 "MobileGestalt".to_string(),
457 plist::Value::Dictionary(plist::Dictionary::from_iter([
458 ("Status".to_string(), plist::Value::String("Success".into())),
459 (
460 "ProductVersion".to_string(),
461 plist::Value::String("26.0".into()),
462 ),
463 ])),
464 )])),
465 )])));
466
467 let value = query_mobile_gestalt(&mut stream, &["ProductVersion"])
468 .await
469 .unwrap();
470 let dict = value.into_dictionary().unwrap();
471 assert_eq!(dict["ProductVersion"].as_string(), Some("26.0"));
472 assert!(!dict.contains_key("Status"));
473 }
474
475 #[tokio::test]
476 async fn query_mobile_gestalt_returns_deprecated_error() {
477 let mut stream =
478 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
479 "Diagnostics".to_string(),
480 plist::Value::Dictionary(plist::Dictionary::from_iter([(
481 "MobileGestalt".to_string(),
482 plist::Value::Dictionary(plist::Dictionary::from_iter([(
483 "Status".to_string(),
484 plist::Value::String("MobileGestaltDeprecated".into()),
485 )])),
486 )])),
487 )])));
488
489 let err = query_mobile_gestalt(&mut stream, &["ProductVersion"])
490 .await
491 .unwrap_err();
492 assert!(matches!(err, DiagnosticsError::Deprecated(_)));
493 assert!(err.to_string().contains("deprecated"));
494 }
495
496 #[tokio::test]
497 async fn query_all_values_sends_all_request_and_returns_diagnostics_payload() {
498 let mut stream =
499 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
500 "Diagnostics".to_string(),
501 plist::Value::Dictionary(plist::Dictionary::from_iter([(
502 "GasGauge".to_string(),
503 plist::Value::Dictionary(plist::Dictionary::from_iter([(
504 "CycleCount".to_string(),
505 plist::Value::Integer(315.into()),
506 )])),
507 )])),
508 )])));
509
510 let value = query_all_values(&mut stream).await.unwrap();
511 let dict = value.into_dictionary().unwrap();
512 assert!(dict.contains_key("GasGauge"));
513
514 let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
515 let payload = &stream.written[4..4 + len];
516 let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
517 assert_eq!(dict["Request"].as_string(), Some("All"));
518 }
519
520 #[tokio::test]
521 async fn query_battery_sends_ioregistry_request_for_power_source() {
522 let mut stream =
523 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
524 "Diagnostics".to_string(),
525 plist::Value::Dictionary(plist::Dictionary::from_iter([(
526 "IORegistry".to_string(),
527 plist::Value::Dictionary(plist::Dictionary::new()),
528 )])),
529 )])));
530
531 let _ = query_battery(&mut stream).await.unwrap();
532
533 let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
534 let payload = &stream.written[4..4 + len];
535 let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
536 assert_eq!(dict["Request"].as_string(), Some("IORegistry"));
537 assert_eq!(dict["EntryClass"].as_string(), Some("IOPMPowerSource"));
538 }
539
540 #[tokio::test]
541 async fn query_battery_extracts_ioregistry_payload() {
542 let mut stream =
543 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
544 "Diagnostics".to_string(),
545 plist::Value::Dictionary(plist::Dictionary::from_iter([(
546 "IORegistry".to_string(),
547 plist::Value::Dictionary(plist::Dictionary::from_iter([
548 (
549 "CurrentCapacity".to_string(),
550 plist::Value::Integer(82.into()),
551 ),
552 ("IsCharging".to_string(), plist::Value::Boolean(true)),
553 ("CycleCount".to_string(), plist::Value::Integer(315.into())),
554 ])),
555 )])),
556 )])));
557
558 let battery = query_battery(&mut stream).await.unwrap();
559 assert_eq!(battery.current_capacity, Some(82));
560 assert_eq!(battery.is_charging, Some(true));
561 assert_eq!(battery.cycle_count, Some(315));
562 }
563
564 #[tokio::test]
565 async fn query_ioregistry_returns_raw_ioregistry_payload() {
566 let mut stream =
567 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
568 "Diagnostics".to_string(),
569 plist::Value::Dictionary(plist::Dictionary::from_iter([(
570 "IORegistry".to_string(),
571 plist::Value::Dictionary(plist::Dictionary::from_iter([(
572 "ProductName".to_string(),
573 plist::Value::String("iPhone".into()),
574 )])),
575 )])),
576 )])));
577
578 let value = query_ioregistry(&mut stream, "IOPlatformExpertDevice")
579 .await
580 .unwrap();
581 let dict = value.into_dictionary().unwrap();
582 assert_eq!(dict["ProductName"].as_string(), Some("iPhone"));
583
584 let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
585 let payload = &stream.written[4..4 + len];
586 let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
587 assert_eq!(dict["Request"].as_string(), Some("IORegistry"));
588 assert_eq!(
589 dict["EntryClass"].as_string(),
590 Some("IOPlatformExpertDevice")
591 );
592 assert!(!dict.contains_key("EntryName"));
593 assert!(!dict.contains_key("CurrentPlane"));
594 }
595
596 #[tokio::test]
597 async fn query_ioregistry_with_name_and_plane_encodes_optional_fields() {
598 let mut stream =
599 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([(
600 "Diagnostics".to_string(),
601 plist::Value::Dictionary(plist::Dictionary::from_iter([(
602 "IORegistry".to_string(),
603 plist::Value::Dictionary(plist::Dictionary::new()),
604 )])),
605 )])));
606
607 let _ = query_ioregistry_with(
608 &mut stream,
609 IoRegistryQuery {
610 entry_class: None,
611 entry_name: Some("device-tree"),
612 current_plane: Some("IODeviceTree"),
613 },
614 )
615 .await
616 .unwrap();
617
618 let len = u32::from_be_bytes(stream.written[..4].try_into().unwrap()) as usize;
619 let payload = &stream.written[4..4 + len];
620 let dict: plist::Dictionary = plist::from_bytes(payload).unwrap();
621 assert_eq!(dict["Request"].as_string(), Some("IORegistry"));
622 assert_eq!(dict["EntryName"].as_string(), Some("device-tree"));
623 assert_eq!(dict["CurrentPlane"].as_string(), Some("IODeviceTree"));
624 assert!(!dict.contains_key("EntryClass"));
625 }
626
627 #[tokio::test]
628 async fn query_ioregistry_reports_status_when_diagnostics_are_missing() {
629 let mut stream =
630 MockStream::with_response(plist::Value::Dictionary(plist::Dictionary::from_iter([
631 (
632 "Status".to_string(),
633 plist::Value::String("LookupFailed".into()),
634 ),
635 (
636 "Error".to_string(),
637 plist::Value::String("Entry not found".into()),
638 ),
639 ])));
640
641 let err = query_ioregistry_with(
642 &mut stream,
643 IoRegistryQuery {
644 entry_class: Some("IO80211Interface"),
645 entry_name: None,
646 current_plane: None,
647 },
648 )
649 .await
650 .unwrap_err();
651
652 assert!(matches!(err, DiagnosticsError::Protocol(_)));
653 let message = err.to_string();
654 assert!(message.contains("LookupFailed"));
655 assert!(message.contains("Entry not found"));
656 assert!(message.contains("Diagnostics"));
657 }
658}