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    /// Get if environment "SFX_ARGS" is set.
13    ///
14    /// Note that `connect` may still fail if the environment is not properly set.
15    pub fn can_connect() -> bool {
16        match std::env::var("SFX_ARGS") {
17            Ok(_) => true,
18            Err(e) => !SfxError::from(e).is_no_sfx(),
19        }
20    }
21
22    /// Connect with the "SFX_ARGS" defined server.
23    pub fn connect_blocking() -> Result<Self, SfxError> {
24        // read args, if is running by sfx
25        let args = std::env::var("SFX_ARGS")?;
26        let mut sfx_args = vec![];
27        for arg in args.split('\n') {
28            let arg = arg.trim();
29            if arg.is_empty() {
30                continue;
31            }
32            sfx_args.push(Txt::from_str(arg));
33        }
34
35        if sfx_args.is_empty() {
36            return Err(io::Error::new(io::ErrorKind::InvalidFilename, "SFX_ARGS is empty").into());
37        }
38
39        #[cfg(feature = "test_util")]
40        if sfx_args[0] == test_util::FAKE_EXE {
41            return Ok(Self {
42                sfx_args: sfx_args.into_boxed_slice(),
43                manifest: test_util::test_manifest(),
44            });
45        }
46
47        // read manifest
48        let out = std::process::Command::new(&sfx_args[0]).env("SFX_GET_MANIFEST", "").output()?;
49        if !out.status.success() {
50            let err = String::from_utf8_lossy(&out.stderr);
51            return Err(io::Error::other(format!("SFX_GET_MANIFEST failed\ncode: {:?}\nstderr:\n{}", out.status, err)).into());
52        }
53        let stdout = match String::from_utf8(out.stdout) {
54            Ok(s) => s,
55            Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidData, e).into()),
56        };
57        let mut manifest = vec![];
58        for line in stdout.lines() {
59            let mut ok = line.is_empty();
60            if let Some((name, len)) = line.rsplit_once(':') {
61                if let Ok(len) = len.parse() {
62                    ok = true;
63                    manifest.push(SfxDataInfo {
64                        name: name.to_txt(),
65                        exact_len: Some(len),
66                    });
67                } else if len == "unknown" {
68                    ok = true;
69                    manifest.push(SfxDataInfo {
70                        name: name.to_txt(),
71                        exact_len: None,
72                    });
73                }
74            }
75            if !ok {
76                return Err(io::Error::new(io::ErrorKind::InvalidData, format!("unexpected manifest format, {line:?}")).into());
77            }
78        }
79        manifest.sort_by(|a, b| a.name.cmp(&b.name));
80        manifest.dedup_by(|a, b| a.name == b.name);
81
82        Ok(Self {
83            sfx_args: sfx_args.into_boxed_slice(),
84            manifest: manifest.into_boxed_slice(),
85        })
86    }
87
88    /// Connect with the "SFX_ARGS" defined server.
89    pub async fn connect() -> Result<Self, SfxError> {
90        zng_task::wait(Self::connect_blocking).await
91    }
92
93    /// The SFX server executable.
94    ///
95    /// This is the `sfx_args()[0]` value.
96    pub fn sfx_exe(&self) -> &Path {
97        Path::new(&self.sfx_args[0])
98    }
99
100    /// The command line arguments from the SFX executable.
101    ///
102    /// Note that the sfx process that spawned this process is still running and
103    /// replicating stdout/err from this process, impersonating the extracted executable.
104    ///
105    /// These args are from the user, the [`std::env::args`] are from the sfx tool and are usually empty.
106    pub fn sfx_args(&self) -> &[Txt] {
107        &self.sfx_args
108    }
109
110    /// Info an all data the server can provide.
111    ///
112    /// List is sorted by name.
113    pub fn manifest(&self) -> &[SfxDataInfo] {
114        &self.manifest
115    }
116
117    /// Get `name` data info it it is in the server manifest.
118    pub fn info(&self, name: &str) -> Option<&SfxDataInfo> {
119        let i = self.manifest.binary_search_by_key(&name, |n| n.name.as_str()).ok()?;
120        Some(&self.manifest[i])
121    }
122
123    /// Read raw data.
124    ///
125    /// The sfx server will automatically decompress/decode data that was marked `compress` on build.
126    pub fn read_blocking(&self, name: &str) -> Result<SfxReadBlocking, SfxError> {
127        if let Some(info) = self.info(name) {
128            #[cfg(feature = "test_util")]
129            if self.sfx_args[0] == test_util::FAKE_EXE {
130                let data = test_util::req_data(name);
131                return Ok(SfxReadBlocking(SfxReadBlockingImpl::Test(std::io::Cursor::new(data))));
132            }
133
134            let mut r = std::process::Command::new(self.sfx_exe()).env("SFX_GET_DATA", name).spawn()?;
135            r.stdin = None;
136            Ok(SfxReadBlocking(SfxReadBlockingImpl::Default(DefaultSfxReadBlocking {
137                server: r,
138                exact_len: info.exact_len,
139            })))
140        } else {
141            Err(SfxError::NotFound(Txt::from_str(name)))
142        }
143    }
144
145    /// Read raw data.
146    ///
147    /// The sfx server will automatically decompress/decode data that was marked `compress` on build.
148    pub async fn read(&self, name: &str) -> Result<SfxRead, SfxError> {
149        if let Some(info) = self.info(name) {
150            #[cfg(feature = "test_util")]
151            if self.sfx_args[0] == test_util::FAKE_EXE {
152                let data = test_util::req_data(name);
153                let inner = SfxReadBlocking(SfxReadBlockingImpl::Test(std::io::Cursor::new(data)));
154                return Ok(SfxRead {
155                    inner: zng_task::io::Unblock::new(inner),
156                });
157            }
158
159            let mut cmd = std::process::Command::new(self.sfx_exe());
160            cmd.env("SFX_GET_DATA", name);
161            let exact_len = info.exact_len;
162            zng_task::wait(move || {
163                let mut r = cmd.spawn()?;
164                r.stdin = None;
165                let inner = SfxReadBlocking(SfxReadBlockingImpl::Default(DefaultSfxReadBlocking { server: r, exact_len }));
166                Ok(SfxRead {
167                    inner: zng_task::io::Unblock::new(inner),
168                })
169            })
170            .await
171        } else {
172            Err(SfxError::NotFound(Txt::from_str(name)))
173        }
174    }
175}
176#[cfg(feature = "test_util")]
177mod test_util {
178    use super::*;
179
180    pub const FAKE_EXE: &str = "test:/fake-sfx";
181
182    impl SfxClient {
183        /// Call this function before [`SfxClient::connect_blocking`] to run
184        /// in setup mode with fake data.
185        ///
186        /// The test [`SfxClient::sfx_exe`] is `"test:/fake-sfx"`, the [`SfxClient::sfx_args`] is
187        /// the test path with `test_args` appended.
188        ///
189        /// The [`SfxClient::read`] will serve the `test_data`.
190        ///
191        /// # Panics
192        ///
193        /// Panics if called after any `SfxClient` operation that uses data or if called twice.
194        pub fn setup_test(test_args: Vec<Txt>, mut test_data: Vec<(Txt, Vec<u8>)>) {
195            test_data.sort_by(|a, b| a.0.cmp(&b.0));
196            test_data.dedup_by(|a, b| a.0 == b.0);
197
198            zng_unique_id::lazy_static_init(&TEST_DATA, Some(test_data)).unwrap_or_else(|_| panic!("SfxClient already used"));
199
200            let mut args = test_args;
201            args.insert(0, FAKE_EXE.into());
202            unsafe {
203                std::env::set_var("SFX_ARGS", args.join("\n"));
204            }
205        }
206    }
207
208    type TestData = Vec<(Txt, Vec<u8>)>;
209    zng_unique_id::lazy_static! {
210        static ref TEST_DATA: Option<TestData> = None;
211    }
212
213    pub(super) fn test_manifest() -> Box<[SfxDataInfo]> {
214        if let Some(data) = &*TEST_DATA {
215            data.iter()
216                .map(|(n, b)| SfxDataInfo {
217                    name: n.clone(),
218                    exact_len: Some(b.len() as u64),
219                })
220                .collect()
221        } else {
222            Box::new([])
223        }
224    }
225
226    pub(super) fn req_data(name: &str) -> &'static [u8] {
227        if let Some(data) = &*TEST_DATA {
228            for (n, entry) in data {
229                if n == name {
230                    return entry;
231                }
232            }
233        }
234        panic!()
235    }
236}
237
238/// Represents a data entry that can be served from a [`SfxClient`].
239#[derive(Clone, PartialEq, Debug)]
240#[non_exhaustive]
241pub struct SfxDataInfo {
242    /// Name.
243    pub name: Txt,
244
245    /// Exact data length, if it is known.
246    ///
247    /// This is `None` if the data is decompressing from a custom zstd
248    /// file that does not define "contentsize".
249    pub exact_len: Option<u64>,
250}
251
252/// Read requested data.
253///
254/// See [`SfxClient::read_blocking`] for more details.
255pub struct SfxReadBlocking(SfxReadBlockingImpl);
256enum SfxReadBlockingImpl {
257    Default(DefaultSfxReadBlocking),
258    #[cfg(feature = "test_util")]
259    Test(std::io::Cursor<&'static [u8]>),
260}
261impl SfxReadBlocking {
262    /// Length not yet read.
263    ///
264    /// Is `None` if the exact length is not known.
265    pub fn exact_len(&self) -> Option<u64> {
266        match &self.0 {
267            SfxReadBlockingImpl::Default(b) => b.exact_len,
268            #[cfg(feature = "test_util")]
269            SfxReadBlockingImpl::Test(c) => Some(c.get_ref().len() as u64 - c.position()),
270        }
271    }
272}
273struct DefaultSfxReadBlocking {
274    // default io::Read::read_to_end impl
275    server: std::process::Child,
276    exact_len: Option<u64>,
277}
278impl io::Read for DefaultSfxReadBlocking {
279    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
280        if let Some(s) = &mut self.server.stdout {
281            let n = s.read(buf)?;
282            if n == 0 {
283                self.server.stdout = None;
284                self.server.kill()?;
285                let status = self.server.wait()?;
286                if !status.success() {
287                    let mut stderr = String::new();
288                    if let Some(mut s) = self.server.stderr.take() {
289                        s.read_to_string(&mut stderr)?;
290                    }
291                    return Err(io::Error::new(
292                        io::ErrorKind::ConnectionAborted,
293                        format!("server failed\nexit code: {status:?}\nstderr:\n{stderr}"),
294                    ));
295                }
296                if let Some(l) = self.exact_len
297                    && l > 0
298                {
299                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "stdout ended before exact length"));
300                }
301            } else if let Some(l) = &mut self.exact_len {
302                let n = n as u64;
303                if *l >= n {
304                    *l -= n;
305                } else {
306                    self.server.stdout = None;
307                    let _ = self.server.kill();
308                    return Err(io::Error::other("stdout longer than expected exact length"));
309                }
310            }
311            Ok(n)
312        } else {
313            Ok(0)
314        }
315    }
316}
317impl io::Read for SfxReadBlocking {
318    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
319        match &mut self.0 {
320            SfxReadBlockingImpl::Default(r) => r.read(buf),
321            #[cfg(feature = "test_util")]
322            SfxReadBlockingImpl::Test(c) => c.read(buf),
323        }
324    }
325
326    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
327        if let Some(l) = self.exact_len() {
328            let l = l.min(usize::MAX as u64) as usize;
329            buf.try_reserve(l)?;
330        }
331        match &mut self.0 {
332            SfxReadBlockingImpl::Default(r) => r.read_to_end(buf),
333            #[cfg(feature = "test_util")]
334            SfxReadBlockingImpl::Test(c) => c.read_to_end(buf),
335        }
336    }
337
338    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
339        if let Some(l) = self.exact_len() {
340            let l = l.min(usize::MAX as u64) as usize;
341            buf.try_reserve(l)?;
342        }
343        match &mut self.0 {
344            SfxReadBlockingImpl::Default(r) => r.read_to_string(buf),
345            #[cfg(feature = "test_util")]
346            SfxReadBlockingImpl::Test(c) => c.read_to_string(buf),
347        }
348    }
349}
350impl Drop for SfxReadBlocking {
351    fn drop(&mut self) {
352        if let SfxReadBlockingImpl::Default(d) = &mut self.0
353            && let Err(e) = d.server.kill()
354        {
355            tracing::error!("cannot kill server on drop, {e}");
356        }
357    }
358}
359
360/// Read requested data.
361///
362/// See [`SfxClient::read`] for more details.
363pub struct SfxRead {
364    inner: zng_task::io::Unblock<SfxReadBlocking>,
365}
366impl SfxRead {
367    /// Length not yet read.
368    ///
369    /// Is `None` if the exact length is not known.
370    pub async fn exact_len(&mut self) -> Option<u64> {
371        self.inner.get_mut().await.exact_len()
372    }
373
374    /// Convert to [`SfxReadBlocking`] that reads the rest of the API.
375    pub async fn into_blocking(self) -> SfxReadBlocking {
376        self.inner.into_inner().await
377    }
378}
379impl zng_task::io::AsyncRead for SfxRead {
380    fn poll_read(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &mut [u8]) -> std::task::Poll<io::Result<usize>> {
381        std::pin::pin!(&mut self.get_mut().inner).poll_read(cx, buf)
382    }
383}
384
385/// Error connecting or requesting data using [`SfxClient`].
386#[derive(Debug)]
387#[non_exhaustive]
388pub enum SfxError {
389    /// Connect error, "SFX_ARGS" not set or invalid.
390    Var(std::env::VarError),
391    /// IO error attempting to read data.
392    Io(io::Error),
393    /// Data name not found.
394    NotFound(Txt),
395}
396impl SfxError {
397    /// If this error indicates that "SFX_ARGS" env var is not set.
398    ///
399    /// This can be ok when the current executable is both a setup and the main app.
400    pub fn is_no_sfx(&self) -> bool {
401        matches!(self, Self::Var(std::env::VarError::NotPresent))
402    }
403}
404impl From<std::env::VarError> for SfxError {
405    fn from(var: std::env::VarError) -> Self {
406        Self::Var(var)
407    }
408}
409impl From<io::Error> for SfxError {
410    fn from(e: io::Error) -> Self {
411        Self::Io(e)
412    }
413}
414impl fmt::Display for SfxError {
415    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
416        match self {
417            SfxError::Var(e) => fmt::Display::fmt(e, f),
418            SfxError::Io(e) => fmt::Display::fmt(e, f),
419            SfxError::NotFound(name) => write!(f, "no {name:?} data on the server catalog"),
420        }
421    }
422}
423impl std::error::Error for SfxError {
424    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
425        match self {
426            SfxError::Var(e) => Some(e),
427            SfxError::Io(e) => Some(e),
428            SfxError::NotFound(_) => None,
429        }
430    }
431}
zng_ext_setup/sfx_client.rs

zng_ext_setup/
sfx_client.rs
