Skip to main content

zng_ext_setup/
sfx_client.rs

1use std::{fmt, io, path::Path};
2
3use zng_txt::{ToTxt, Txt};
4
5/// Connection with a self-extracting executable generated by `cargo zng res --tool sfx`.
6#[derive(Clone, PartialEq, Debug)]
7pub struct SfxClient {
8    sfx_args: Box<[Txt]>,
9    manifest: Box<[SfxDataInfo]>,
10}
11impl SfxClient {
12    /// Connect with the "SFX_ARGS" defined server.
13    pub fn connect_blocking() -> Result<Self, SfxError> {
14        // read args, if is running by sfx
15        let args = std::env::var("SFX_ARGS")?;
16        let mut sfx_args = vec![];
17        for arg in args.split('\n') {
18            let arg = arg.trim();
19            if arg.is_empty() {
20                continue;
21            }
22            sfx_args.push(Txt::from_str(arg));
23        }
24
25        if sfx_args.is_empty() {
26            return Err(io::Error::new(io::ErrorKind::InvalidFilename, "SFX_ARGS is empty").into());
27        }
28
29        // read manifest
30        let out = std::process::Command::new(&sfx_args[0]).env("SFX_GET_MANIFEST", "").output()?;
31        if !out.status.success() {
32            let err = String::from_utf8_lossy(&out.stderr);
33            return Err(io::Error::other(format!("SFX_GET_MANIFEST failed\ncode: {:?}\nstderr:\n{}", out.status, err)).into());
34        }
35        let stdout = match String::from_utf8(out.stdout) {
36            Ok(s) => s,
37            Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidData, e).into()),
38        };
39        let mut manifest = vec![];
40        for line in stdout.lines() {
41            let mut ok = line.is_empty();
42            if let Some((name, len)) = line.rsplit_once(':') {
43                if let Ok(len) = len.parse() {
44                    ok = true;
45                    manifest.push(SfxDataInfo {
46                        name: name.to_txt(),
47                        exact_len: Some(len),
48                    });
49                } else if len == "unknown" {
50                    ok = true;
51                    manifest.push(SfxDataInfo {
52                        name: name.to_txt(),
53                        exact_len: None,
54                    });
55                }
56            }
57            if !ok {
58                return Err(io::Error::new(io::ErrorKind::InvalidData, format!("unexpected manifest format, {line:?}")).into());
59            }
60        }
61        manifest.sort_by(|a, b| a.name.cmp(&b.name));
62        manifest.dedup_by(|a, b| a.name == b.name);
63
64        Ok(Self {
65            sfx_args: sfx_args.into_boxed_slice(),
66            manifest: manifest.into_boxed_slice(),
67        })
68    }
69
70    /// Connect with the "SFX_ARGS" defined server.
71    pub async fn connect() -> Result<Self, SfxError> {
72        zng_task::wait(Self::connect_blocking).await
73    }
74
75    /// The SFX server executable.
76    ///
77    /// This is the `sfx_args()[0]` value.
78    pub fn sfx_exe(&self) -> &Path {
79        Path::new(&self.sfx_args[0])
80    }
81
82    /// The command line arguments from the SFX executable.
83    ///
84    /// Note that the sfx process that spawned this process is still running and
85    /// replicating stdout/err from this process, impersonating the extracted executable.
86    ///
87    /// These args are from the user, the [`std::env::args`] are from the sfx tool and are usually empty.
88    pub fn sfx_args(&self) -> &[Txt] {
89        &self.sfx_args
90    }
91
92    /// Info an all data the server can provide.
93    ///
94    /// List is sorted by name.
95    pub fn manifest(&self) -> &[SfxDataInfo] {
96        &self.manifest
97    }
98
99    /// Get `name` data info it it is in the server manifest.
100    pub fn info(&self, name: &str) -> Option<&SfxDataInfo> {
101        let i = self.manifest.binary_search_by_key(&name, |n| n.name.as_str()).ok()?;
102        Some(&self.manifest[i])
103    }
104
105    /// Read raw data.
106    ///
107    /// The sfx server will automatically decompress/decode data that was marked `compress` on build.
108    pub fn read_blocking(&self, name: &str) -> Result<SfxReadBlocking, SfxError> {
109        if let Some(info) = self.info(name) {
110            let mut r = std::process::Command::new(self.sfx_exe()).env("SFX_GET_DATA", name).spawn()?;
111            r.stdin = None;
112            Ok(SfxReadBlocking(DefaultSfxReadBlocking {
113                server: r,
114                exact_len: info.exact_len,
115            }))
116        } else {
117            Err(SfxError::NotFound(Txt::from_str(name)))
118        }
119    }
120
121    /// Read raw data.
122    ///
123    /// The sfx server will automatically decompress/decode data that was marked `compress` on build.
124    pub async fn read(&self, name: &str) -> Result<SfxRead, SfxError> {
125        if let Some(info) = self.info(name) {
126            let mut cmd = std::process::Command::new(self.sfx_exe());
127            cmd.env("SFX_GET_DATA", name);
128            let exact_len = info.exact_len;
129            zng_task::wait(move || {
130                let mut r = cmd.spawn()?;
131                r.stdin = None;
132                let inner = SfxReadBlocking(DefaultSfxReadBlocking { server: r, exact_len });
133                Ok(SfxRead {
134                    inner: zng_task::io::Unblock::new(inner),
135                })
136            })
137            .await
138        } else {
139            Err(SfxError::NotFound(Txt::from_str(name)))
140        }
141    }
142}
143
144/// Represents a data entry that can be served from a [`SfxClient`].
145#[derive(Clone, PartialEq, Debug)]
146#[non_exhaustive]
147pub struct SfxDataInfo {
148    /// Name.
149    pub name: Txt,
150
151    /// Exact data length, if it is known.
152    ///
153    /// This is `None` if the data is decompressing from a custom zstd
154    /// file that does not define "contentsize".
155    pub exact_len: Option<u64>,
156}
157
158/// Read requested data.
159///
160/// See [`SfxClient::read_blocking`] for more details.
161pub struct SfxReadBlocking(DefaultSfxReadBlocking);
162impl SfxReadBlocking {
163    /// Length not yet read.
164    ///
165    /// Is `None` if the exact length is not known.
166    pub fn exact_len(&self) -> Option<u64> {
167        self.0.exact_len
168    }
169}
170struct DefaultSfxReadBlocking {
171    // default io::Read::read_to_end impl
172    server: std::process::Child,
173    exact_len: Option<u64>,
174}
175impl io::Read for DefaultSfxReadBlocking {
176    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
177        if let Some(s) = &mut self.server.stdout {
178            let n = s.read(buf)?;
179            if n == 0 {
180                self.server.stdout = None;
181                self.server.kill()?;
182                let status = self.server.wait()?;
183                if !status.success() {
184                    let mut stderr = String::new();
185                    if let Some(mut s) = self.server.stderr.take() {
186                        s.read_to_string(&mut stderr)?;
187                    }
188                    return Err(io::Error::new(
189                        io::ErrorKind::ConnectionAborted,
190                        format!("server failed\nexit code: {status:?}\nstderr:\n{stderr}"),
191                    ));
192                }
193                if let Some(l) = self.exact_len
194                    && l > 0
195                {
196                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "stdout ended before exact length"));
197                }
198            } else if let Some(l) = &mut self.exact_len {
199                let n = n as u64;
200                if *l >= n {
201                    *l -= n;
202                } else {
203                    self.server.stdout = None;
204                    let _ = self.server.kill();
205                    return Err(io::Error::other("stdout longer than expected exact length"));
206                }
207            }
208            Ok(n)
209        } else {
210            Ok(0)
211        }
212    }
213}
214impl io::Read for SfxReadBlocking {
215    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
216        self.0.read(buf)
217    }
218
219    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
220        if let Some(l) = self.0.exact_len {
221            let l = l.min(usize::MAX as u64) as usize;
222            buf.try_reserve(l)?;
223        }
224        self.0.read_to_end(buf)
225    }
226
227    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
228        if let Some(l) = self.0.exact_len {
229            let l = l.min(usize::MAX as u64) as usize;
230            buf.try_reserve(l)?;
231        }
232        self.0.read_to_string(buf)
233    }
234}
235impl Drop for SfxReadBlocking {
236    fn drop(&mut self) {
237        if let Err(e) = self.0.server.kill() {
238            tracing::error!("cannot kill server on drop, {e}");
239        }
240    }
241}
242
243/// Read requested data.
244///
245/// See [`SfxClient::read`] for more details.
246pub struct SfxRead {
247    inner: zng_task::io::Unblock<SfxReadBlocking>,
248}
249impl SfxRead {
250    /// Length not yet read.
251    ///
252    /// Is `None` if the exact length is not known.
253    pub async fn exact_len(&mut self) -> Option<u64> {
254        self.inner.get_mut().await.0.exact_len
255    }
256
257    /// Convert to [`SfxReadBlocking`] that reads the rest of the API.
258    pub async fn into_blocking(self) -> SfxReadBlocking {
259        self.inner.into_inner().await
260    }
261}
262impl zng_task::io::AsyncRead for SfxRead {
263    fn poll_read(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &mut [u8]) -> std::task::Poll<io::Result<usize>> {
264        std::pin::pin!(&mut self.get_mut().inner).poll_read(cx, buf)
265    }
266}
267
268/// Error connecting or requesting data using [`SfxClient`].
269#[derive(Debug)]
270#[non_exhaustive]
271pub enum SfxError {
272    /// Connect error, "SFX_ARGS" not set or invalid.
273    Var(std::env::VarError),
274    /// IO error attempting to read data.
275    Io(io::Error),
276    /// Data name not found.
277    NotFound(Txt),
278}
279impl SfxError {
280    /// If this error indicates that "SFX_ARGS" env var is not set.
281    ///
282    /// This can be ok when the current executable is both a setup and the main app.
283    pub fn is_no_sfx(&self) -> bool {
284        matches!(self, Self::Var(std::env::VarError::NotPresent))
285    }
286}
287impl From<std::env::VarError> for SfxError {
288    fn from(var: std::env::VarError) -> Self {
289        Self::Var(var)
290    }
291}
292impl From<io::Error> for SfxError {
293    fn from(e: io::Error) -> Self {
294        Self::Io(e)
295    }
296}
297impl fmt::Display for SfxError {
298    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
299        match self {
300            SfxError::Var(e) => fmt::Display::fmt(e, f),
301            SfxError::Io(e) => fmt::Display::fmt(e, f),
302            SfxError::NotFound(name) => write!(f, "no {name:?} data on the server catalog"),
303        }
304    }
305}
306impl std::error::Error for SfxError {
307    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
308        match self {
309            SfxError::Var(e) => Some(e),
310            SfxError::Io(e) => Some(e),
311            SfxError::NotFound(_) => None,
312        }
313    }
314}