send_it/
writer.rs

1use std::io::Write;
2use crate::Segment;
3
4
5/// A writer for sending several segments over a stream using variable length encoding
6/// Data is written in little-endian if the feature "big-endian" is not enabled
7/// # Example
8/// ```
9/// use send_it::writer::VarWriter;
10///
11/// let mut sender = VarWriter::new();
12///
13/// sender.add_string("Hello");
14/// sender.add_string("World");
15///
16/// let mut buffer = Vec::new();
17/// sender.send(&mut buffer).unwrap();
18/// ```
19pub struct VarWriter {
20    data: Vec<Segment>,
21}
22
23impl VarWriter {
24    /// Create a new VarWriter
25    pub fn new() -> VarWriter {
26        VarWriter {
27            data: Vec::new(),
28        }
29    }
30
31    /// Add a segment to the writer
32    pub fn add(&mut self, segment: Segment) {
33        self.data.push(segment);
34    }
35
36    /// Add a string to the writer
37    /// # Example
38    /// ```
39    /// use send_it::writer::VarWriter;
40    ///
41    /// let mut sender = VarWriter::new();
42    ///
43    /// sender.add_string("Hello");
44    /// ```
45    pub fn add_string<S: Into<String>>(&mut self, string: S) {
46        self.add(Segment::from(string.into()))
47    }
48
49    /// Add raw data to the writer
50    /// # Example
51    /// ```
52    /// use send_it::writer::VarWriter;
53    ///
54    /// let mut sender = VarWriter::new();
55    ///
56    /// sender.add_raw(&[0x48, 0x65, 0x6C, 0x6C, 0x6F]);
57    /// ```
58    pub fn add_raw(&mut self, raw: &[u8]) {
59        self.data.push(Segment::from(raw));
60    }
61
62    /// Encodes the data and sends it over the stream.
63    /// * The data is cleared after sending.
64    /// # Example
65    /// ```
66    /// use send_it::writer::VarWriter;
67    ///
68    /// let mut sender = VarWriter::new();
69    ///
70    /// sender.add_string("Hello");
71    /// sender.add_string("World");
72    ///
73    /// let mut buffer = Vec::new();
74    /// sender.send(&mut buffer).unwrap();
75    /// ```
76    pub fn send<W: Write>(&mut self, stream: &mut W) -> std::io::Result<()> {
77        self.send_without_clearing(stream)?;
78
79        // Clear the internal data after sending
80        self.clear();
81
82        Ok(())
83    }
84
85    /// Encodes the data and sends it over the stream.
86    /// * The data is not cleared after sending.
87    /// # Example
88    /// ```
89    /// use send_it::writer::VarWriter;
90    ///
91    /// let mut sender = VarWriter::new();
92    ///
93    /// sender.add_string("Hello");
94    /// sender.add_string("World");
95    ///
96    /// let mut buffer = Vec::new();
97    /// sender.send_without_clearing(&mut buffer).unwrap();
98    /// ```
99    pub fn send_without_clearing<W: Write>(&mut self, stream: &mut W) -> std::io::Result<()> {
100        let total_size: usize = self.data.iter().map(|segment| segment.len() + 4).sum();
101
102        // Write the total size as varint
103        self.write_varint(stream, total_size)?;
104
105        // Write each segment's size and the segment itself
106        for segment in &self.data {
107            self.write_u32(stream, segment.len() as u32)?;
108            // write the segment
109            stream.write_all(segment.as_ref())?;
110        }
111
112        Ok(())
113    }
114    
115    fn write_varint<W: Write>(&self, writer: &mut W, mut value: usize) -> std::io::Result<()> {
116        loop {
117            let mut byte = (value & 0x7F) as u8;
118            value >>= 7;
119            if value != 0 {
120                byte |= 0x80;
121            }
122            writer.write_all(&[byte])?;
123            if value == 0 {
124                break;
125            }
126        }
127        Ok(())
128    }
129
130    #[cfg(not(feature = "big-endian"))]
131    fn write_u32<W: Write>(&self, writer: &mut W, value: u32) -> std::io::Result<()> {
132        // writes little-endian
133        writer.write_all(&[
134            (value & 0xFF) as u8,
135            ((value >> 8) & 0xFF) as u8,
136            ((value >> 16) & 0xFF) as u8,
137            ((value >> 24) & 0xFF) as u8,
138        ])?;
139        Ok(())
140    }
141
142    #[cfg(feature = "big-endian")]
143    fn write_u32<W: Write>(&self, writer: &mut W, value: u32) -> std::io::Result<()> {
144        // writes big-endian
145        writer.write_all(&[
146            ((value >> 24) & 0xFF) as u8,
147            ((value >> 16) & 0xFF) as u8,
148            ((value >> 8) & 0xFF) as u8,
149            (value & 0xFF) as u8,
150        ])?;
151        Ok(())
152    }
153
154    /// Removes all segments from the writer
155    pub fn clear(&mut self) {
156        self.data.clear();
157    }
158}
159
160impl Write for VarWriter {
161    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
162        let vec = buf.to_vec();
163        let size = vec.len();
164        self.data.push(Segment::from(vec));
165        Ok(size)
166    }
167
168    fn flush(&mut self) -> std::io::Result<()> {
169        // nothing to do
170        Ok(())
171    }
172}
173
174impl Default for VarWriter {
175    fn default() -> Self {
176        Self::new()
177    }
178}
send_it/writer.rs

send_it/
writer.rs
