1
//! Utilities
2

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

            
7
use anyhow::{Context as _, anyhow};
8

            
9
use super::CliError;
10

            
11
/// Command line filename argument, allowing `-` for stdin/stdout
12
//
13
// TODO DIRAUTH currently this can only be used for output file arguments,
14
// but we will implement using this for an input file argument too.
15
//
16
// TODO move this somewhere deeper in the stack (tor-basic-utils even maybe?)
17
// and replace open-coding in eg crates/arti/src/subcommands/hsc.rs display_service_discovery_key
18
// If we do that:
19
//   - consider whether we should preserve file permissions
20
//   - see the comment about leftover `.tmp` files in `write`, below
21
#[derive(Debug, Clone)]
22
pub(super) enum FilenameOrStdio {
23
    /// Filename
24
    Path(String),
25
    /// `-`
26
    Stdio,
27
}
28

            
29
/// Output file currently being written to
30
///
31
/// See [`FilenameOrStdio::start_writing`].
32
pub(super) struct Writing {
33
    /// Actual open-file
34
    handle: Box<dyn io::Write>,
35
    /// Filenames
36
    files: Option<WritingFiles>,
37
}
38

            
39
/// Filenames when writing an output file
40
struct WritingFiles {
41
    /// The `.tmp` file
42
    tmp: String,
43
    /// The main file
44
    main: String,
45
}
46

            
47
impl FromStr for FilenameOrStdio {
48
    type Err = anyhow::Error;
49
    fn from_str(s: &str) -> Result<Self, Self::Err> {
50
        match s {
51
            "" => Err(anyhow!("empty filename")),
52
            "-" => Ok(FilenameOrStdio::Stdio),
53
            other => Ok(FilenameOrStdio::Path(other.to_owned())),
54
        }
55
    }
56
}
57

            
58
impl FilenameOrStdio {
59
    /// Write the output file, with write-to-`.tmp`-and-rename
60
    ///
61
    /// `writer` should generate the actual output.
62
    /// It shouldn't fail other than for write errors.
63
    ///
64
    /// Makes no attempt to preserve file permissions.
65
    ///
66
    /// ### `.tmp` file cleanup
67
    ///
68
    /// If this fails, we can leave a `.tmp` file lying about.  This is fine for
69
    /// the C Tor Arti consensus method plugin.
70
    ///
71
    /// If we promote this code we might want to make an effort to clean up leftover `.tmp`
72
    /// files.  Note however that it is not possible to make that cleanup reliable.
73
    /// Therefore whatever invokes us will need to be able to clean up such garbage, anyway.
74
    /// It might be better to rely entirely on that, avoiding the situation where
75
    /// a caller fails to have any cleanup functionality and very occasionally
76
    /// `.tmp` files get left over anyway (and *this* code unjustifiably gets the blame).
77
    pub(super) fn write<W>(&self, writer: W) -> Result<(), CliError>
78
    where
79
        W: FnOnce(&mut dyn io::Write) -> io::Result<()>,
80
    {
81
        let mut w = self.start_writing()?;
82
        w.append_with(writer)?;
83
        w.finish()
84
    }
85

            
86
    /// Start writing the output file, with write-to-`.tmp`-and-rename
87
    /// in the case of the file being on the file system (i.e. no stdout).
88
    ///
89
    /// Returns a [`Writing`], which implements `io::Write` -
90
    /// but usually it's better to use [`Writing::append_with`]
91
    /// since that automatically converts errors to a nice `CliError`.
92
    ///
93
    /// When the output is complete, you *must* call [`Writing::finish`]
94
    /// to install the output file.
95
    ///
96
    /// If `Writing` is dropped, we leave the `.tmp` file lying about.
97
    /// See [`FilenameOrStdio::write`] for more information.
98
    ///
99
    /// Makes no attempt to preserve file permissions.
100
    pub(super) fn start_writing(&self) -> Result<Writing, CliError> {
101
        match self {
102
            FilenameOrStdio::Stdio => {
103
                //
104
                Ok(Writing {
105
                    handle: Box::new(io::stdout().lock()),
106
                    files: None,
107
                })
108
            }
109
            FilenameOrStdio::Path(main) => {
110
                let tmp = format!("{main}.tmp");
111
                let f = File::create(&tmp)
112
                    .with_context(|| format!("create {tmp:?}"))
113
                    .map_err(convert_output_error)?;
114
                let f = BufWriter::new(f);
115
                Ok(Writing {
116
                    handle: Box::new(f),
117
                    files: Some(WritingFiles {
118
                        tmp,
119
                        main: main.clone(),
120
                    }),
121
                })
122
            }
123
        }
124
    }
125
}
126

            
127
impl io::Write for Writing {
128
    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
129
        self.handle.write(b)
130
    }
131
    fn flush(&mut self) -> io::Result<()> {
132
        self.handle.flush()
133
    }
134
}
135

            
136
impl Writing {
137
    /// Write to an output file, with a `dyn io::Write`, handling errors
138
    ///
139
    /// Calls `writer(self)` but converts the error into a nice `CliError`
140
    pub(super) fn append_with<W>(&mut self, writer: W) -> Result<(), CliError>
141
    where
142
        W: FnOnce(&mut dyn io::Write) -> io::Result<()>,
143
    {
144
        writer(self).map_err(self.handle_write_error())
145
    }
146

            
147
    /// Finish writing and install the output file
148
    pub(super) fn finish(mut self) -> Result<(), CliError> {
149
        self.flush().map_err(self.handle_write_error())?;
150

            
151
        if let Some(WritingFiles { main, tmp }) = &self.files {
152
            fs::rename(tmp, main)
153
                .with_context(|| format!("install {tmp:?} as {main:?}"))
154
                .map_err(convert_output_error)?;
155
        }
156

            
157
        Ok(())
158
    }
159

            
160
    /// Helper to handle errors from `io::Write::write` and `flush`
161
    fn handle_write_error(&self) -> impl FnOnce(io::Error) -> CliError {
162
        let msg = self
163
            .files
164
            .as_ref()
165
            .map(|files| format!("write to {:?}", files.tmp))
166
            .unwrap_or("write to stdout".into());
167
        move |e| convert_output_error(anyhow::Error::from(e).context(msg))
168
    }
169
}
170

            
171
/// Helper to convert an error encountered while writing to CliError
172
fn convert_output_error(e: anyhow::Error) -> CliError {
173
    CliError::OperationalError(e.context("write output"))
174
}