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

zng_ext_setup/
sfx_client.rs
