Skip to main content

regain_core/
accessory.rs

1//! Shared client for the newline-JSON accessory workers in `regain-device`.
2//! Native camera workers keep their separate transport/recovery supervisor. COM
3//! imports can append a finite binary body under the same process request guard.
4use crate::process::ProcessGuard;
5use anyhow::{Context, Result};
6use serde::Deserialize;
7use serde::de::DeserializeOwned;
8use serde_json::Value;
9use std::{fmt, future::Future, pin::Pin, time::Duration};
10use tokio::{
11    io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
12    process::{Child, ChildStdin, ChildStdout},
13};
14
15pub const MAX_REQUEST_BYTES: usize = 4096;
16pub const MAX_RESPONSE_BYTES: usize = 1024 * 1024;
17
18/// Distinguish a framed worker error from transport loss. A worker error may
19/// still follow a dispatched USB command; neither permits automatic replay.
20#[derive(Debug)]
21pub enum AccessoryError {
22    InvalidRequest,
23    Disconnected,
24    Transport,
25    Timeout,
26    CommandFailed(String),
27}
28impl fmt::Display for AccessoryError {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        match self {
31            Self::InvalidRequest => f.write_str("Invalid or oversized accessory request"),
32            Self::Disconnected => f.write_str("Accessory worker is disconnected"),
33            Self::Transport => {
34                f.write_str("Accessory worker response failed; command was not retried")
35            }
36            Self::Timeout => f.write_str("Accessory worker timed out; command was not retried"),
37            Self::CommandFailed(message) => f.write_str(message),
38        }
39    }
40}
41impl std::error::Error for AccessoryError {}
42
43pub struct AccessoryWorker {
44    pub child: Child,
45    input: Option<ChildStdin>,
46    output: BufReader<ChildStdout>,
47    _guard: ProcessGuard,
48}
49/// A finite, owned worker command. Diagnostic text is never returned to callers.
50pub struct Collected<T> {
51    pub value: T,
52    pub diagnostics_present: bool,
53}
54impl AccessoryWorker {
55    /// Adopt a child with piped stdin/stdout. Windows ownership also attaches a
56    /// kill-on-close job so a dead parent cannot strand an exclusive device owner.
57    pub fn new(child: Child) -> Result<Self> {
58        Self::adopt(child, false)
59    }
60    /// COM imports own the private Regain worker, not vendor helpers it launches.
61    pub fn new_independent(child: Child) -> Result<Self> {
62        Self::adopt(child, true)
63    }
64    fn adopt(mut child: Child, independent: bool) -> Result<Self> {
65        let pipes = (|| {
66            let input = child.stdin.take().context("Accessory stdin unavailable")?;
67            let output = child
68                .stdout
69                .take()
70                .context("Accessory stdout unavailable")?;
71            let pid = child.id().context("Accessory process exited")?;
72            let guard = if independent {
73                ProcessGuard::attach_independent(pid)?
74            } else {
75                ProcessGuard::attach(pid)?
76            };
77            Ok::<_, anyhow::Error>((input, output, guard))
78        })();
79        match pipes {
80            Ok((input, output, guard)) => Ok(Self {
81                child,
82                input: Some(input),
83                output: BufReader::new(output),
84                _guard: guard,
85            }),
86            Err(error) => {
87                let _ = child.start_kill();
88                Err(error)
89            }
90        }
91    }
92    pub async fn request(&mut self, request: Value) -> Result<Value> {
93        self.request_with_timeout(request, Duration::from_secs(10))
94            .await
95    }
96    /// Release a one-shot worker's stdin startup barrier only after ownership is
97    /// attached, then collect bounded JSON and reap it. Consuming self makes
98    /// timeout/cancellation terminal. Neither stdout nor diagnostics are logged.
99    pub async fn collect_json_with_timeout<T: DeserializeOwned>(
100        mut self,
101        request: Value,
102        deadline: Duration,
103    ) -> Result<Collected<T>> {
104        let bytes = encode_request(request)?;
105        let operation = async {
106            let input = self.input.as_mut().ok_or(AccessoryError::Disconnected)?;
107            input.write_all(&bytes).await?;
108            input.flush().await?;
109            self.input.take();
110            let diagnostics = self.child.stderr.take();
111            let stdout = async {
112                let mut bytes = Vec::new();
113                (&mut self.output)
114                    .take((MAX_RESPONSE_BYTES + 1) as u64)
115                    .read_to_end(&mut bytes)
116                    .await?;
117                if bytes.len() > MAX_RESPONSE_BYTES {
118                    return Err(AccessoryError::Transport.into());
119                }
120                serde_json::from_slice::<T>(&bytes).map_err(|_| AccessoryError::Transport.into())
121            };
122            let stderr = async {
123                let Some(mut stream) = diagnostics else {
124                    return Ok::<_, anyhow::Error>(false);
125                };
126                let mut total = 0usize;
127                let mut buffer = [0; 4096];
128                loop {
129                    let count = stream.read(&mut buffer).await?;
130                    if count == 0 {
131                        return Ok(total != 0);
132                    }
133                    total += count;
134                    if total > MAX_RESPONSE_BYTES {
135                        return Err(AccessoryError::Transport.into());
136                    }
137                }
138            };
139            let (value, diagnostics_present, status) = tokio::try_join!(stdout, stderr, async {
140                Ok::<_, anyhow::Error>(self.child.wait().await?)
141            })?;
142            if !status.success() {
143                return Err(AccessoryError::Transport.into());
144            }
145            Ok(Collected {
146                value,
147                diagnostics_present,
148            })
149        };
150        tokio::time::timeout(deadline, operation)
151            .await
152            .unwrap_or_else(|_| Err(AccessoryError::Timeout.into()))
153    }
154    pub async fn request_with_timeout(
155        &mut self,
156        request: Value,
157        deadline: Duration,
158    ) -> Result<Value> {
159        self.request_typed_with_timeout(request, deadline).await
160    }
161    /// Deserialize the inner reply directly from its frame. Typed protocols
162    /// must reject duplicate fields before conversion to a generic JSON map.
163    pub async fn request_typed_with_timeout<T: DeserializeOwned>(
164        &mut self,
165        request: Value,
166        deadline: Duration,
167    ) -> Result<T> {
168        let bytes = encode_request(request)?;
169        if self.input.is_none() {
170            return Err(AccessoryError::Disconnected.into());
171        }
172        // This guard lives inside the awaited operation. Timeout OR cancellation
173        // closes stdin and retires the process immediately, preventing a delayed
174        // response from being consumed as the next command's acknowledgement.
175        let operation = async {
176            let mut guard = PendingRequest {
177                child: &mut self.child,
178                input: &mut self.input,
179                armed: true,
180            };
181            guard
182                .input
183                .as_mut()
184                .unwrap()
185                .write_all(&bytes)
186                .await
187                .map_err(|_| AccessoryError::Transport)?;
188            let bytes = read_response(&mut self.output).await?;
189            let result = decode_typed_response(&bytes);
190            if result.is_ok() || matches!(result, Err(AccessoryError::CommandFailed(_))) {
191                guard.armed = false;
192            }
193            result
194        };
195        tokio::time::timeout(deadline, operation)
196            .await
197            .unwrap_or(Err(AccessoryError::Timeout))
198            .map_err(Into::into)
199    }
200    /// One acknowledged frame followed by a caller-defined finite binary body.
201    /// The same process/cancellation guard spans both; a rejected header, partial
202    /// body, failed allocation, timeout or dropped future retires the stream.
203    /// The callback must consume and validate the entire body before succeeding.
204    pub async fn request_typed_with_stream<T, R, F>(
205        &mut self,
206        request: Value,
207        deadline: Duration,
208        body: F,
209    ) -> Result<R>
210    where
211        T: DeserializeOwned,
212        F: for<'a> FnOnce(
213            T,
214            &'a mut BufReader<ChildStdout>,
215        ) -> Pin<Box<dyn Future<Output = Result<R>> + Send + 'a>>,
216    {
217        let bytes = encode_request(request)?;
218        if self.input.is_none() {
219            return Err(AccessoryError::Disconnected.into());
220        }
221        let operation = async {
222            let mut guard = PendingRequest {
223                child: &mut self.child,
224                input: &mut self.input,
225                armed: true,
226            };
227            guard
228                .input
229                .as_mut()
230                .unwrap()
231                .write_all(&bytes)
232                .await
233                .map_err(|_| AccessoryError::Transport)?;
234            let bytes = read_response(&mut self.output).await?;
235            let header = match decode_typed_response(&bytes) {
236                Ok(header) => header,
237                Err(error) => {
238                    if matches!(error, AccessoryError::CommandFailed(_)) {
239                        guard.armed = false;
240                    }
241                    return Err(error.into());
242                }
243            };
244            let result = body(header, &mut self.output).await;
245            if result.is_ok() {
246                guard.armed = false;
247            }
248            result
249        };
250        tokio::time::timeout(deadline, operation)
251            .await
252            .unwrap_or_else(|_| Err(AccessoryError::Timeout.into()))
253    }
254    /// Retire a stream whose outstanding request may have been cancelled.
255    pub fn reset(&mut self) {
256        self.input.take();
257        let _ = self.child.start_kill();
258    }
259    pub async fn close(mut self) {
260        self.input.take();
261        if !matches!(
262            tokio::time::timeout(Duration::from_secs(5), self.child.wait()).await,
263            Ok(Ok(_))
264        ) {
265            let _ = self.child.kill().await;
266        }
267    }
268}
269impl Drop for AccessoryWorker {
270    fn drop(&mut self) {
271        self.reset();
272    }
273}
274struct PendingRequest<'a> {
275    child: &'a mut Child,
276    input: &'a mut Option<ChildStdin>,
277    armed: bool,
278}
279impl Drop for PendingRequest<'_> {
280    fn drop(&mut self) {
281        if self.armed {
282            self.input.take();
283            let _ = self.child.start_kill();
284        }
285    }
286}
287fn encode_request(request: Value) -> Result<Vec<u8>, AccessoryError> {
288    if !request.is_object() {
289        return Err(AccessoryError::InvalidRequest);
290    }
291    let mut bytes = serde_json::to_vec(&request).map_err(|_| AccessoryError::InvalidRequest)?;
292    if bytes.len() >= MAX_REQUEST_BYTES {
293        return Err(AccessoryError::InvalidRequest);
294    }
295    bytes.push(b'\n');
296    Ok(bytes)
297}
298async fn read_response(input: &mut (impl AsyncBufRead + Unpin)) -> Result<Vec<u8>, AccessoryError> {
299    let mut line = Vec::new();
300    loop {
301        let buffer = input
302            .fill_buf()
303            .await
304            .map_err(|_| AccessoryError::Transport)?;
305        if buffer.is_empty() {
306            return Err(AccessoryError::Transport);
307        }
308        let end = buffer.iter().position(|byte| *byte == b'\n').map(|n| n + 1);
309        let length = end.unwrap_or(buffer.len());
310        if length > MAX_RESPONSE_BYTES - line.len() {
311            return Err(AccessoryError::Transport);
312        }
313        line.extend_from_slice(&buffer[..length]);
314        input.consume(length);
315        if end.is_some() {
316            return Ok(line);
317        }
318    }
319}
320#[derive(Deserialize)]
321#[serde(bound(deserialize = "T: Deserialize<'de>"))]
322struct Reply<T> {
323    ok: bool,
324    #[serde(default, deserialize_with = "present_result")]
325    result: Option<T>,
326    error: Option<String>,
327}
328fn present_result<'de, T: Deserialize<'de>, D: serde::Deserializer<'de>>(
329    deserializer: D,
330) -> Result<Option<T>, D::Error> {
331    T::deserialize(deserializer).map(Some)
332}
333#[cfg(test)]
334fn decode_response(bytes: &[u8]) -> Result<Value, AccessoryError> {
335    decode_typed_response(bytes)
336}
337fn decode_typed_response<T: DeserializeOwned>(bytes: &[u8]) -> Result<T, AccessoryError> {
338    let reply: Reply<T> = serde_json::from_slice(bytes).map_err(|_| AccessoryError::Transport)?;
339    if reply.ok {
340        reply.result.ok_or(AccessoryError::Transport)
341    } else {
342        Err(AccessoryError::CommandFailed(
343            reply
344                .error
345                .unwrap_or_else(|| "Accessory command failed".into()),
346        ))
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353    use serde_json::json;
354
355    #[test]
356    fn typed_reply_rejects_duplicate_inner_fields_and_keeps_null_values() {
357        #[derive(Deserialize)]
358        struct Typed {
359            id: u64,
360        }
361        assert_eq!(
362            decode_typed_response::<Typed>(br#"{"ok":true,"result":{"id":1}}"#)
363                .unwrap()
364                .id,
365            1
366        );
367        assert!(matches!(
368            decode_typed_response::<Typed>(br#"{"ok":true,"result":{"id":1,"id":2}}"#),
369            Err(AccessoryError::Transport)
370        ));
371        assert_eq!(
372            decode_response(br#"{"ok":true,"result":null}"#).unwrap(),
373            Value::Null
374        );
375        assert!(decode_response(br#"{"ok":true}"#).is_err());
376    }
377
378    #[tokio::test]
379    async fn framing_bounds_partial_eof_and_multiple_responses() {
380        let mut input = &b"{\"ok\":true,\"result\":12}\n{\"ok\":true,\"result\":13}\n"[..];
381        assert_eq!(
382            decode_response(&read_response(&mut input).await.unwrap()).unwrap(),
383            12
384        );
385        assert_eq!(
386            decode_response(&read_response(&mut input).await.unwrap()).unwrap(),
387            13
388        );
389        assert!(read_response(&mut input).await.is_err());
390        assert!(read_response(&mut &b"{\"ok\":true}"[..]).await.is_err());
391        let large = vec![b'x'; MAX_RESPONSE_BYTES + 1];
392        assert!(read_response(&mut &large[..]).await.is_err());
393    }
394    #[test]
395    fn request_bounds_and_worker_errors() {
396        assert!(encode_request(json!([])).is_err());
397        assert!(encode_request(json!({"command":"x".repeat(MAX_REQUEST_BYTES)})).is_err());
398        assert_eq!(
399            encode_request(json!({"command":"status"})).unwrap().last(),
400            Some(&b'\n')
401        );
402        assert!(
403            matches!(decode_response(br#"{"ok":false,"error":"busy"}"#), Err(AccessoryError::CommandFailed(message)) if message == "busy")
404        );
405        for bytes in [
406            br#"{"ok":"true"}"#.as_slice(),
407            br#"{"ok":true,"ok":false}"#,
408            b"bad",
409        ] {
410            assert!(matches!(
411                decode_response(bytes),
412                Err(AccessoryError::Transport)
413            ));
414        }
415    }
416}