Skip to main content

minicbor_io/
async_writer.rs

1use crate::Error;
2use futures_io::AsyncWrite;
3use futures_util::AsyncWriteExt;
4use minicbor::Encode;
5use std::io;
6
7/// Wraps an [`AsyncWrite`] and writes length-delimited CBOR values.
8///
9/// *Requires cargo feature* `"async-io"`.
10#[derive(Debug)]
11pub struct AsyncWriter<W> {
12    writer: W,
13    buffer: Vec<u8>,
14    max_len: usize,
15    state: State
16}
17
18/// Write state.
19#[derive(Debug)]
20enum State {
21    /// Nothing is written at the moment.
22    None,
23    /// Writing buffer from offset.
24    WriteFrom(usize)
25}
26
27impl<W> AsyncWriter<W> {
28    /// Create a new writer with a max. buffer size of 512KiB.
29    pub fn new(writer: W) -> Self {
30        Self::with_buffer(writer, Vec::new())
31    }
32
33    /// Create a new writer with a max. buffer size of 512KiB.
34    pub fn with_buffer(writer: W, buffer: Vec<u8>) -> Self {
35        Self { writer, buffer, max_len: 512 * 1024, state: State::None }
36    }
37
38    /// Set the max. buffer size in bytes.
39    ///
40    /// If length values greater than this are encoded, an
41    /// [`Error::InvalidLen`] will be returned.
42    pub fn set_max_len(&mut self, val: u32) {
43        self.max_len = val as usize
44    }
45
46    /// Get a reference to the inner buffer.
47    pub fn buffer(&self) -> &Vec<u8> {
48        &self.buffer
49    }
50
51    /// Get a mutable reference to the inner buffer.
52    pub fn buffer_mut(&mut self) -> &mut Vec<u8> {
53        &mut self.buffer
54    }
55
56    /// Get a reference to the inner writer.
57    pub fn writer(&self) -> &W {
58        &self.writer
59    }
60
61    /// Get a mutable reference to the inner writer.
62    pub fn writer_mut(&mut self) -> &mut W {
63        &mut self.writer
64    }
65
66    /// Deconstruct this writer into the inner writer and the buffer.
67    pub fn into_parts(self) -> (W, Vec<u8>) {
68        (self.writer, self.buffer)
69    }
70}
71
72impl<W: AsyncWrite + Unpin> AsyncWriter<W> {
73    /// Encode and write a CBOR value and return its size in bytes.
74    ///
75    /// The value will be preceded by a `u32` (4 bytes in network byte order),
76    /// denoting the length of bytes constituting the serialised value.
77    ///
78    /// # Cancellation
79    ///
80    /// If the future returned by `AsyncWriter::write` is dropped while still
81    /// pending, subsequent calls to `AsyncWriter::write` will discard any
82    /// buffered data and instead encode, buffer and commit the new value.
83    /// Cancelling a future thus cancels the transfer. However, it is also
84    /// possible to resume the transfer by calling [`AsyncWriter::sync`]
85    /// after cancellation, which is normally called implicitly by this method.
86    pub async fn write<T: Encode<()>>(&mut self, val: T) -> Result<usize, Error> {
87        self.write_with(val, &mut ()).await
88    }
89
90    /// Like [`AsyncWriter::write`] but accepting a user provided encoding context.
91    pub async fn write_with<C, T: Encode<C>>(&mut self, val: T, ctx: &mut C) -> Result<usize, Error> {
92        self.buffer.resize(4, 0u8);
93        minicbor::encode_with(val, &mut self.buffer, ctx)?;
94        if self.buffer.len() - 4 > self.max_len {
95            return Err(Error::InvalidLen)
96        }
97        let prefix = (self.buffer.len() as u32 - 4).to_be_bytes();
98        self.buffer[.. 4].copy_from_slice(&prefix);
99        self.state = State::WriteFrom(0);
100
101        self.sync().await?;
102
103        Ok(self.buffer.len() - 4)
104    }
105
106    /// Commit any buffered data to the inner `AsyncWrite`.
107    ///
108    /// This method is implicitly called by [`AsyncWriter::write`]. The only
109    /// reason to call it explicitly is to resume the write operation started
110    /// by a previously unfinished, i.e. cancelled, `AsyncWriter::write` call.
111    pub async fn sync(&mut self) -> Result<(), Error> {
112        loop {
113            match self.state {
114                State::None => {
115                    return Ok(())
116                }
117                State::WriteFrom(o) if o >= self.buffer.len() => {
118                    self.state = State::None;
119                    return Ok(())
120                }
121                State::WriteFrom(ref mut o) => {
122                    let n = self.writer.write(&self.buffer[*o ..]).await?;
123                    if n == 0 {
124                        return Err(Error::Io(io::ErrorKind::WriteZero.into()))
125                    }
126                    *o += n
127                }
128            }
129        }
130    }
131
132    /// Flush the inner `AsyncWrite`.
133    pub async fn flush(&mut self) -> Result<(), Error> {
134        self.writer.flush().await?;
135        Ok(())
136    }
137}