1
//! Utilities
2

            
3
use std::fs::{self, File};
4
use std::io::{self, BufReader, BufWriter, Read as _, Write as _};
5
use std::str::FromStr;
6

            
7
use anyhow::anyhow;
8

            
9
use super::*;
10

            
11
define_derive_deftly! {
12
    /// Define `fn new`
13
    ///
14
    /// # Field attributes
15
    ///
16
    ///  * `#[deftly(new(arg))]`: this field should be an argument to `new`.
17
    ///    If not specified, `Default` is used.
18
    //
19
    // TODO enhance documentation and promote somewhere?
20
    // tor_netdoc::Constructor is not suitable because all the fields would have to be pub(crate)
21
    // derive_more::Constructor is not suitable because it can't Default fields
22
    New beta_deftly:
23

            
24
    ${defcond ARG fmeta(new(arg))}
25

            
26
    $impl {
27
        $/// Make a new `$tname`
28
        $tvis fn new(
29
            $(
30
                ${when ARG}
31
                $fname: $ftype,
32
            )
33
2
        ) -> Self {
34
            $tname {
35
                $(
36
                    $fname
37
                    ${if not(ARG) {
38
                        : Default::default()
39
                    }},
40
                )
41
            }
42
        }
43
    }
44
}
45

            
46
/// Command line filename argument, allowing `-` for stdin/stdout
47
///
48
/// Doesn't implement `Display`; for content error reporting prefer
49
/// [`Reading::description`].
50
//
51
// TODO move this somewhere deeper in the stack (tor-basic-utils even maybe?)
52
// and replace open-coding in eg crates/arti/src/subcommands/hsc.rs display_service_discovery_key
53
// If we do that:
54
//   - consider whether we should preserve file permissions
55
//   - see the comment about leftover `.tmp` files in `write`, below
56
//   - fix the locking problem (see the two TODO locking/hang)
57
#[derive(Debug, Clone)]
58
pub(super) enum FilenameOrStdio {
59
    /// Filename
60
    Path(String),
61
    /// `-`
62
    Stdio,
63
}
64

            
65
/// Output file currently being read from
66
///
67
/// See [`FilenameOrStdio::start_reading`].
68
#[derive(Educe)]
69
#[educe(Debug)]
70
pub(super) struct Reading {
71
    /// Actual open-file
72
    #[educe(Debug(ignore))]
73
    handle: io::BufReader<Box<dyn io::Read>>,
74
    /// Description (for error messages), already quoted
75
    description: String,
76
}
77

            
78
/// Output file currently being written to
79
///
80
/// See [`FilenameOrStdio::start_writing`].
81
#[derive(Educe)]
82
#[educe(Debug)]
83
pub(super) struct Writing {
84
    /// Actual open-file
85
    #[educe(Debug(ignore))]
86
    handle: io::BufWriter<Box<dyn io::Write>>,
87
    /// Filenames
88
    files: Option<WritingFiles>,
89
}
90

            
91
/// Filenames when writing an output file
92
#[derive(Debug)]
93
struct WritingFiles {
94
    /// The `.tmp` file
95
    tmp: String,
96
    /// The main file
97
    main: String,
98
}
99

            
100
impl FromStr for FilenameOrStdio {
101
    type Err = anyhow::Error;
102
    fn from_str(s: &str) -> Result<Self, Self::Err> {
103
        match s {
104
            "" => Err(anyhow!("empty filename")),
105
            "-" => Ok(FilenameOrStdio::Stdio),
106
            other => Ok(FilenameOrStdio::Path(other.to_owned())),
107
        }
108
    }
109
}
110

            
111
impl FilenameOrStdio {
112
    /// Write the output file, with write-to-`.tmp`-and-rename
113
    ///
114
    /// `writer` should generate the actual output.
115
    /// It shouldn't fail other than for write errors.
116
    ///
117
    /// Makes no attempt to preserve file permissions.
118
    ///
119
    /// ### `.tmp` file cleanup
120
    ///
121
    /// If this fails, we can leave a `.tmp` file lying about.  This is fine for
122
    /// the C Tor Arti consensus method plugin.
123
    ///
124
    /// If we promote this code we might want to make an effort to clean up leftover `.tmp`
125
    /// files.  Note however that it is not possible to make that cleanup reliable.
126
    /// Therefore whatever invokes us will need to be able to clean up such garbage, anyway.
127
    /// It might be better to rely entirely on that, avoiding the situation where
128
    /// a caller fails to have any cleanup functionality and very occasionally
129
    /// `.tmp` files get left over anyway (and *this* code unjustifiably gets the blame).
130
    pub(super) fn write<W>(&self, writer: W) -> Result<(), CliError>
131
    where
132
        W: FnOnce(&mut dyn io::Write) -> io::Result<()>,
133
    {
134
        let mut w = self.start_writing()?;
135
        w.append_with(writer)?;
136
        w.finish()
137
    }
138

            
139
    /// Start writing the output file, with write-to-`.tmp`-and-rename
140
    /// in the case of the file being on the file system (i.e. no stdout).
141
    ///
142
    /// Returns a [`Writing`], which implements `io::Write` -
143
    /// but usually it's better to use [`Writing::append_with`]
144
    /// since that automatically converts errors to a nice `CliError`.
145
    ///
146
    /// When the output is complete, you *must* call [`Writing::finish`]
147
    /// to install the output file.
148
    ///
149
    /// If `Writing` is dropped, we leave the `.tmp` file lying about.
150
    /// See [`FilenameOrStdio::write`] for more information.
151
    ///
152
    /// Makes no attempt to preserve file permissions.
153
4
    pub(super) fn start_writing(&self) -> Result<Writing, CliError> {
154
4
        match self {
155
            FilenameOrStdio::Stdio => {
156
                //
157
                Ok(Writing {
158
                    // TODO locking/hang
159
                    //
160
                    // If multiple `FilenameOrStdio`s referring to stdin/stdout are
161
                    // used simultaneously (rather than sequentially), this will hang.
162
                    //
163
                    // Eg, `compute-mds --mds-out - --meta-out -` will hang.
164
                    //
165
                    // Unfortunately there is no `.try_lock()`.  We could have a private
166
                    // global lock to detect this situation.  That seems overkill for
167
                    // the dirauth plugin, but ought to be done before these routines
168
                    // are promoted to general utilities.
169
                    handle: BufWriter::new(Box::new(io::stdout().lock())),
170
                    files: None,
171
                })
172
            }
173
4
            FilenameOrStdio::Path(main) => {
174
4
                let tmp = format!("{main}.tmp");
175
4
                let f = File::create(&tmp)
176
4
                    .with_context(|| format!("create {tmp:?}"))
177
4
                    .map_err(convert_output_error)?;
178
4
                Ok(Writing {
179
4
                    handle: BufWriter::new(Box::new(f)),
180
4
                    files: Some(WritingFiles {
181
4
                        tmp,
182
4
                        main: main.clone(),
183
4
                    }),
184
4
                })
185
            }
186
        }
187
4
    }
188

            
189
    /// Start reading this input file
190
    pub(super) fn start_reading(&self) -> Result<Reading, CliError> {
191
        let (handle, description);
192
        match self {
193
            FilenameOrStdio::Stdio => {
194
                // TODO locking/hang, see above
195
                handle = Box::new(io::stdin().lock()) as _;
196
                description = "<stdin>".into();
197
            }
198
            FilenameOrStdio::Path(path) => {
199
                handle = Box::new(
200
                    File::open(path)
201
                        .with_context(|| format!("open input file {path:?}"))
202
                        .map_err(CliError::OperationalError)?,
203
                ) as _;
204
                description = format!("{path:?}");
205
            }
206
        }
207
        let handle = BufReader::new(handle);
208
        Ok(Reading {
209
            handle,
210
            description,
211
        })
212
    }
213
}
214

            
215
impl io::Read for Reading {
216
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
217
        self.handle.read(buf)
218
    }
219
}
220

            
221
impl io::BufRead for Reading {
222
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
223
        self.handle.fill_buf()
224
    }
225
    #[allow(clippy::semicolon_if_nothing_returned)] // more consistent without the ;
226
    fn consume(&mut self, n: usize) {
227
        self.handle.consume(n)
228
    }
229
}
230

            
231
impl Reading {
232
    /// Reads the whole file into memory as a `String`
233
    ///
234
    /// Non-UTF-8 input is treated as an operational error, just like an io error.
235
    pub(crate) fn read_entire_string(mut self) -> Result<String, CliError> {
236
        let mut s = String::new();
237
        self.handle
238
            .read_to_string(&mut s)
239
            .map_err(self.handle_read_error())?;
240
        Ok(s)
241
    }
242

            
243
    /// Provide a human-readable description of what this is (eg for error reporting)
244
    pub(crate) fn description(&self) -> &str {
245
        &self.description
246
    }
247

            
248
    /// Returns an error handler for IO errors from this input file
249
    pub(crate) fn handle_read_error(&self) -> impl FnOnce(io::Error) -> CliError {
250
        let m = format!("error reading {}", self.description);
251
        move |e| CliError::OperationalError(anyhow::Error::from(e).context(m))
252
    }
253
}
254

            
255
impl io::Write for Writing {
256
20
    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
257
20
        self.handle.write(b)
258
20
    }
259
4
    fn flush(&mut self) -> io::Result<()> {
260
4
        self.handle.flush()
261
4
    }
262
}
263

            
264
impl Writing {
265
    /// Write to an output file, with a `dyn io::Write`, handling errors
266
    ///
267
    /// Calls `writer(self)` but converts the error into a nice `CliError`
268
20
    pub(super) fn append_with<W>(&mut self, writer: W) -> Result<(), CliError>
269
20
    where
270
20
        W: FnOnce(&mut dyn io::Write) -> io::Result<()>,
271
    {
272
20
        writer(self).map_err(self.handle_write_error())
273
20
    }
274

            
275
    /// Finish writing and install the output file
276
4
    pub(super) fn finish(mut self) -> Result<(), CliError> {
277
4
        self.flush().map_err(self.handle_write_error())?;
278

            
279
4
        if let Some(WritingFiles { main, tmp }) = &self.files {
280
4
            fs::rename(tmp, main)
281
4
                .with_context(|| format!("install {tmp:?} as {main:?}"))
282
4
                .map_err(convert_output_error)?;
283
        }
284

            
285
4
        Ok(())
286
4
    }
287

            
288
    /// Helper to handle errors from `io::Write::write` and `flush`
289
24
    fn handle_write_error(&self) -> impl FnOnce(io::Error) -> CliError {
290
24
        let msg = self
291
24
            .files
292
24
            .as_ref()
293
36
            .map(|files| format!("write to {:?}", files.tmp))
294
24
            .unwrap_or("write to stdout".into());
295
        move |e| convert_output_error(anyhow::Error::from(e).context(msg))
296
24
    }
297
}
298

            
299
/// Helper to convert an error encountered while writing to CliError
300
fn convert_output_error(e: anyhow::Error) -> CliError {
301
    CliError::OperationalError(e.context("write output"))
302
}
303

            
304
/// A network document, but possibly preceded by C Tor `@`-annotation(s)
305
///
306
/// Parsing adapter wrapper.
307
///
308
/// Implements `NetdocParseable`: discards any annotations,
309
/// and then parses `D`.
310
pub(crate) struct CTorAnnotated<D>(pub(crate) D);
311

            
312
impl<D: NetdocParseable> NetdocParseable for CTorAnnotated<D> {
313
    fn doctype_for_error() -> &'static str {
314
        D::doctype_for_error()
315
    }
316
    fn is_intro_item_keyword(kw: tor_netdoc::parse2::KeywordRef<'_>) -> bool {
317
        D::is_intro_item_keyword(kw)
318
    }
319
    fn is_structural_keyword(
320
        kw: tor_netdoc::parse2::KeywordRef<'_>,
321
    ) -> Option<parse2::IsStructural> {
322
        D::is_structural_keyword(kw)
323
    }
324
10
    fn from_items(
325
10
        input: &mut parse2::ItemStream<'_>,
326
10
        stop_at: tor_netdoc::stop_at!(),
327
10
    ) -> Result<Self, parse2::ErrorProblem> {
328
10
        input.with_inner_lines_mut(|lines| {
329
12
            while let Some(peeked) = lines.peek() {
330
12
                let line = lines.peeked_line(&peeked);
331
12
                if !line.starts_with('@') {
332
10
                    break;
333
2
                }
334
2
                let _: &str = lines.next().expect("just peeked");
335
            }
336
10
        });
337
10
        D::from_items(input, stop_at).map(CTorAnnotated)
338
10
    }
339
}