falco_plugin/strings/
cstring_writer.rs

1use memchr::memchr;
2use std::ffi::CString;
3use std::fmt::{Debug, Formatter};
4use std::io::Write;
5
6/// # A helper that enables writing into CStrings
7///
8/// This type implements [`Write`] and yields a [`CString`] at the end,
9/// which is useful for generating string data to be shared with the Falco
10/// plugin framework.
11///
12/// The [`Write`] implementation returns an error whenever the data to be written
13/// contains a NUL byte.
14///
15/// Example:
16/// ```
17/// use std::ffi::CString;
18/// use falco_plugin::strings::CStringWriter;
19/// use std::io::Write;
20/// let mut writer = CStringWriter::default();
21///
22/// write!(writer, "Hello, world, five={}", 5)?;
23///
24/// let output: CString = writer.into_cstring();
25/// # Result::<(), std::io::Error>::Ok(())
26/// ```
27#[derive(Default)]
28pub struct CStringWriter(Vec<u8>);
29
30impl Debug for CStringWriter {
31    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
32        f.debug_tuple("CStringWriter")
33            .field(&String::from_utf8_lossy(self.0.as_slice()))
34            .finish()
35    }
36}
37
38impl Write for CStringWriter {
39    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
40        if memchr(0, buf).is_some() {
41            Err(std::io::Error::new(
42                std::io::ErrorKind::InvalidData,
43                "NUL in data",
44            ))
45        } else {
46            self.0.write(buf)
47        }
48    }
49
50    fn flush(&mut self) -> std::io::Result<()> {
51        self.0.flush()
52    }
53}
54
55impl CStringWriter {
56    /// # Finalize the writer object and return a [`CString`]
57    ///
58    /// This method consumes the CStringWriter and returns a CString
59    /// containing all the written data
60    pub fn into_cstring(mut self) -> CString {
61        self.0.push(0);
62
63        // SAFETY: we disallow embedded NULs on write and add the trailing NUL just above
64        //         so the vector contains exactly one NUL, at the end
65        unsafe { CString::from_vec_with_nul_unchecked(self.0) }
66    }
67
68    /// # Finalize the writer object and store the output in a [`CString`]
69    ///
70    /// This method consumes the CStringWriter, but instead of returning
71    /// a CString, it stores the output in an existing CString (replacing
72    /// any previous content).
73    pub fn store(self, target: &mut CString) {
74        let mut s = self.into_cstring();
75        std::mem::swap(&mut s, target)
76    }
77}
78
79/// # Extension trait to enable [`Write`] on [`CString`]
80///
81/// It receives a closure that takes a [`CStringWriter`] and stores the write
82/// result in the instance (replacing any previous content).
83///
84/// If the written data contains NUL bytes, an error is returned.
85///
86/// # Example:
87///
88/// ```
89/// use std::ffi::CString;
90/// use std::io::Write;
91/// use falco_plugin::strings::WriteIntoCString;
92/// let mut buf = CString::default();
93///
94/// buf.write_into(|w| write!(w, "hello")).unwrap();
95///
96/// assert_eq!(buf.as_c_str(), c"hello");
97/// ```
98pub trait WriteIntoCString {
99    /// Write into a [`std::ffi::CString`] using [`std::io::Write`]
100    fn write_into<F>(&mut self, func: F) -> std::io::Result<()>
101    where
102        F: FnOnce(&mut CStringWriter) -> std::io::Result<()>;
103}
104
105impl WriteIntoCString for CString {
106    fn write_into<F>(&mut self, func: F) -> std::io::Result<()>
107    where
108        F: FnOnce(&mut CStringWriter) -> std::io::Result<()>,
109    {
110        let mut w = CStringWriter::default();
111        func(&mut w)?;
112        w.store(self);
113        Ok(())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_valid_store() {
123        let mut buf = CString::default();
124
125        let mut writer = CStringWriter::default();
126        write!(writer, "hello").unwrap();
127        #[allow(clippy::write_literal)]
128        write!(writer, ", {}", "world").unwrap();
129        writer.flush().unwrap();
130
131        writer.store(&mut buf);
132
133        assert_eq!(buf.as_c_str(), c"hello, world");
134    }
135
136    #[test]
137    fn test_invalid_store() {
138        let mut writer = CStringWriter::default();
139        write!(writer, "hell\0o").unwrap_err();
140    }
141
142    #[test]
143    fn test_valid_write_into() {
144        let mut buf = CString::default();
145
146        buf.write_into(|w| write!(w, "hello")).unwrap();
147
148        assert_eq!(buf.as_c_str(), c"hello");
149    }
150
151    #[test]
152    fn test_invalid_write_into() {
153        let mut buf = CString::default();
154
155        buf.write_into(|w| write!(w, "hell\0o")).unwrap_err();
156    }
157}
falco_plugin/strings/cstring_writer.rs

falco_plugin/strings/
cstring_writer.rs
