fuchsia_fuzzctl/
writer.rs

Help
1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use anyhow::{Context as _, Result};
6use diagnostics_data::{LogsData, Severity};
7use serde_json::to_vec_pretty;
8use std::cell::RefCell;
9use std::fmt::{Debug, Display};
10use std::fs::{create_dir_all, File};
11use std::io::{self, Write};
12use std::path::{Path, PathBuf};
13use std::rc::Rc;
14use std::sync::{Arc, Mutex};
15use termion::{color, style};
16
17/// `Writer` handles formatting and delivering output from both the plugin and fuzzer.
18///
19/// Callers can use this object to configure how output should be formatted, suppressed, etc. The
20/// underlying `OutputSink` is responsible for determining where output actually goes. The plugin
21/// uses `print` and `error`, while the fuzzer uses `write_all` and `log`. Fuzzer output can be
22/// paused and resumed, with data being buffered between those calls.
23#[derive(Debug)]
24pub struct Writer<O: OutputSink> {
25    output: O,
26    muted: bool,
27    use_colors: bool,
28    file: Rc<RefCell<Option<File>>>,
29    buffered: Arc<Mutex<Option<Vec<Buffered>>>>,
30}
31
32/// `OutputSink`s takes output and writes it to some destination.
33pub trait OutputSink: Clone + Debug + 'static {
34    // Writes bytes directly to stdout.
35    fn write_all(&self, _buf: &[u8]);
36
37    // Writes a displayable message.
38    fn print<D: Display>(&self, _message: D);
39
40    // Writes an error message.
41    fn error<D: Display>(&self, _message: D);
42}
43
44#[derive(Debug)]
45enum Buffered {
46    Data(Vec<u8>),
47    Log(String),
48}
49
50impl<O: OutputSink> Writer<O> {
51    /// Creates a new `Writer`.
52    ///
53    /// This object will display data received by methods like `print` by formatting it and sending
54    /// it to the given `output` sink.
55    pub fn new(output: O) -> Self {
56        Self {
57            output,
58            muted: false,
59            use_colors: true,
60            file: Rc::new(RefCell::new(None)),
61            buffered: Arc::new(Mutex::new(None)),
62        }
63    }
64
65    /// If true, suppresses output except for plugin errors.
66    pub fn mute(&mut self, muted: bool) {
67        self.muted = muted
68    }
69
70    /// If true, embeds ANSI escape sequences in the output to add colors and styles.
71    pub fn use_colors(&mut self, use_colors: bool) {
72        self.use_colors = use_colors
73    }
74
75    /// Creates a new `Writer` that is a clone of this object, except that it duplicates its output
76    /// and writes it to a file created from `dirname` and `filename`.
77    pub fn tee<P: AsRef<Path>, S: AsRef<str>>(&self, dirname: P, filename: S) -> Result<Writer<O>> {
78        let dirname = dirname.as_ref();
79        create_dir_all(dirname)
80            .with_context(|| format!("failed to create directory: {}", dirname.display()))?;
81        let mut path = PathBuf::from(dirname);
82        path.push(filename.as_ref());
83        let file = File::options()
84            .append(true)
85            .create(true)
86            .open(&path)
87            .with_context(|| format!("failed to open file: {}", path.display()))?;
88        Ok(Self {
89            output: self.output.clone(),
90            muted: self.muted,
91            use_colors: self.use_colors,
92            file: Rc::new(RefCell::new(Some(file))),
93            buffered: Arc::clone(&self.buffered),
94        })
95    }
96
97    /// Writes a displayable message to the `OutputSink`.
98    ///
99    /// This method is used with output from the `ffx fuzz` plugin.
100    pub fn print<D: Display>(&self, message: D) {
101        if self.muted {
102            return;
103        }
104        let formatted = format!("{}{}{}", self.yellow(), message, self.reset());
105        self.output.print(formatted);
106    }
107
108    /// Like `print`, except that it also adds a newline.
109    pub fn println<D: Display>(&self, message: D) {
110        self.print(format!("{}\n", message));
111    }
112
113    /// Writes a displayable error to the `OutputSink`.
114    ///
115    /// This method is used with output from the `ffx fuzz` plugin.
116    pub fn error<D: Display>(&self, message: D) {
117        let formatted =
118            format!("{}{}ERROR: {}{}\n", self.bold(), self.red(), message, self.reset());
119        self.output.error(formatted);
120    }
121
122    /// Writes bytes directly to the `OutputSink`.
123    ///
124    /// This method is used with output from the fuzzer. If `pause` is called, data passed to this
125    /// method will be buffered until `resume` is called.
126    pub fn write_all(&self, buf: &[u8]) {
127        self.write_all_to_file(buf).expect("failed to write data to file");
128        if self.muted {
129            return;
130        }
131        let mut buffered = self.buffered.lock().unwrap();
132        match buffered.as_mut() {
133            Some(buffered) => buffered.push(Buffered::Data(buf.to_vec())),
134            None => self.output.write_all(buf),
135        };
136    }
137
138    /// Writes a structured log entry to the `OutputSink`.
139    ///
140    /// This method is used with output from the fuzzer. If `pause` is called, data passed to this
141    /// method will be buffered until `resume` is called.
142    ///
143    /// The display format loosely imitates that of `ffx log` as implemented by that plugin's
144    /// `DefaultLogFormatter`.
145    pub fn log(&self, logs_data: LogsData) {
146        let mut serialized = to_vec_pretty(&logs_data).expect("failed to serialize");
147        serialized.push('\n' as u8);
148        self.write_all_to_file(&serialized).expect("failed to write log to file");
149        if self.muted {
150            return;
151        }
152        let color_and_style = match logs_data.metadata.severity {
153            Severity::Fatal => format!("{}{}", self.bold(), self.red()),
154            Severity::Error => self.red(),
155            Severity::Warn => self.yellow(),
156            _ => String::default(),
157        };
158        let severity = &format!("{}", logs_data.metadata.severity)[..1];
159        let location = match (&logs_data.metadata.file, &logs_data.metadata.line) {
160            (Some(filename), Some(line)) => format!(": [{filename}:{line}]"),
161            (Some(filename), None) => format!(": [{filename}]"),
162            _ => String::default(),
163        };
164        let formatted = format!(
165            "[{ts:05.3}][{moniker}][{tags}][{textfmt}{sev}{reset}]{loc} {textfmt}{msg}{reset}\n",
166            ts = logs_data.metadata.timestamp.into_nanos() as f64 / 1_000_000_000 as f64,
167            moniker = logs_data.moniker,
168            tags = logs_data.tags().map(|t| t.join(",")).unwrap_or(String::default()),
169            textfmt = color_and_style,
170            sev = severity,
171            reset = self.reset(),
172            loc = location,
173            msg = logs_data.msg().unwrap_or("<missing message>"),
174        );
175        let mut buffered = self.buffered.lock().unwrap();
176        match buffered.as_mut() {
177            Some(buffered) => buffered.push(Buffered::Log(formatted)),
178            None => self.output.print(formatted),
179        };
180    }
181
182    /// Pauses the display of fuzzer output.
183    ///
184    /// When this is called, any data passed to `write_all` or `log` will be buffered until `resume`
185    /// is called.
186    pub fn pause(&self) {
187        let mut buffered = self.buffered.lock().unwrap();
188        if buffered.is_none() {
189            *buffered = Some(Vec::new());
190        }
191    }
192
193    /// Resumes the display of fuzzer output.
194    ///
195    /// When this is called, any data passed to `write_all` or `log` since `pause` was invoked will
196    /// be sent to the `OutputSink`.
197    pub fn resume(&self) {
198        let buffered = match self.buffered.lock().unwrap().take() {
199            Some(buffered) => buffered,
200            None => return,
201        };
202        for message in buffered.into_iter() {
203            match message {
204                Buffered::Data(bytes) => self.output.write_all(&bytes),
205                Buffered::Log(formatted) => self.output.print(formatted),
206            }
207        }
208    }
209
210    fn write_all_to_file(&self, buf: &[u8]) -> Result<()> {
211        let mut file = self.file.borrow_mut();
212        if let Some(file) = file.as_mut() {
213            file.write_all(buf).context("failed to write to file")?;
214        }
215        Ok(())
216    }
217
218    fn bold(&self) -> String {
219        if self.use_colors {
220            style::Bold.to_string()
221        } else {
222            String::default()
223        }
224    }
225
226    fn yellow(&self) -> String {
227        if self.use_colors {
228            color::Fg(color::Yellow).to_string()
229        } else {
230            String::default()
231        }
232    }
233
234    fn red(&self) -> String {
235        if self.use_colors {
236            color::Fg(color::Red).to_string()
237        } else {
238            String::default()
239        }
240    }
241
242    fn reset(&self) -> String {
243        if self.use_colors {
244            style::Reset.to_string()
245        } else {
246            String::default()
247        }
248    }
249}
250
251impl<O: OutputSink> Clone for Writer<O> {
252    fn clone(&self) -> Self {
253        Self {
254            output: self.output.clone(),
255            muted: self.muted,
256            use_colors: self.use_colors,
257            file: Rc::clone(&self.file),
258            buffered: Arc::clone(&self.buffered),
259        }
260    }
261}
262
263/// `StdioSink` sends output to standard output and standard error.
264#[derive(Clone, Debug)]
265pub struct StdioSink {
266    /// Indicates this sink is connected to a tty, and can support raw terminal mode.
267    pub is_tty: bool,
268}
269
270impl OutputSink for StdioSink {
271    fn write_all(&self, buf: &[u8]) {
272        self.raw_write(io::stdout(), buf).expect("failed to write to stdout");
273    }
274
275    fn print<D: Display>(&self, message: D) {
276        let formatted = format!("{}", message);
277        self.raw_write(io::stdout(), formatted.as_bytes()).expect("failed to write to stdout");
278    }
279
280    fn error<D: Display>(&self, message: D) {
281        let formatted = format!("{}\n", message);
282        self.raw_write(io::stderr(), formatted.as_bytes()).expect("failed to write to stderr");
283    }
284}
285
286impl StdioSink {
287    fn raw_write<W: Write>(&self, mut w: W, buf: &[u8]) -> Result<()> {
288        if !self.is_tty {
289            w.write_all(buf)?;
290            return Ok(());
291        }
292        let bufs = buf.split_inclusive(|&c| c == '\n' as u8);
293        for buf in bufs.into_iter() {
294            let last = buf.last().unwrap();
295            if *last == '\n' as u8 {
296                // TTY may be in "raw" mode; move back to the start of the line on newline.
297                w.write_all(&buf[0..buf.len() - 1])?;
298                w.write_all("\r\n".as_bytes())?;
299            } else {
300                w.write_all(buf)?;
301            }
302        }
303        Ok(())
304    }
305}
fuchsia_fuzzctl/writer.rs

fuchsia_fuzzctl/
writer.rs
