summaryrefslogtreecommitdiff
path: root/crates/arti-dirauth/src/authority_plugin/utils.rs
blob: bdc60c7799a345e92280d0d1aebc5902ae797cdf (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
//! Utilities

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

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

use super::CliError;

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

/// Output file currently being written to
///
/// See [`FilenameOrStdio::start_writing`].
pub(super) struct Writing {
    /// Actual open-file
    handle: Box<dyn io::Write>,
    /// Filenames
    files: Option<WritingFiles>,
}

/// Filenames when writing an output file
struct WritingFiles {
    /// The `.tmp` file
    tmp: String,
    /// The main file
    main: String,
}

impl FromStr for FilenameOrStdio {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "" => Err(anyhow!("empty filename")),
            "-" => Ok(FilenameOrStdio::Stdio),
            other => Ok(FilenameOrStdio::Path(other.to_owned())),
        }
    }
}

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

    /// Start writing the output file, with write-to-`.tmp`-and-rename
    /// in the case of the file being on the file system (i.e. no stdout).
    ///
    /// Returns a [`Writing`], which implements `io::Write` -
    /// but usually it's better to use [`Writing::append_with`]
    /// since that automatically converts errors to a nice `CliError`.
    ///
    /// When the output is complete, you *must* call [`Writing::finish`]
    /// to install the output file.
    ///
    /// If `Writing` is dropped, we leave the `.tmp` file lying about.
    /// See [`FilenameOrStdio::write`] for more information.
    ///
    /// Makes no attempt to preserve file permissions.
    pub(super) fn start_writing(&self) -> Result<Writing, CliError> {
        match self {
            FilenameOrStdio::Stdio => {
                //
                Ok(Writing {
                    handle: Box::new(io::stdout().lock()),
                    files: None,
                })
            }
            FilenameOrStdio::Path(main) => {
                let tmp = format!("{main}.tmp");
                let f = File::create(&tmp)
                    .with_context(|| format!("create {tmp:?}"))
                    .map_err(convert_output_error)?;
                let f = BufWriter::new(f);
                Ok(Writing {
                    handle: Box::new(f),
                    files: Some(WritingFiles {
                        tmp,
                        main: main.clone(),
                    }),
                })
            }
        }
    }
}

impl io::Write for Writing {
    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
        self.handle.write(b)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.handle.flush()
    }
}

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

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

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

        Ok(())
    }

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

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