edge_ws/
io.rs

1use core::cmp::min;
2
3use embedded_io_async::{self, Read, ReadExactError, Write};
4
5use super::*;
6
7pub type Error<E> = super::Error<E>;
8
9impl<E> Error<E>
10where
11    E: embedded_io_async::Error,
12{
13    pub fn erase(&self) -> Error<embedded_io_async::ErrorKind> {
14        match self {
15            Self::Incomplete(size) => Error::Incomplete(*size),
16            Self::Invalid => Error::Invalid,
17            Self::BufferOverflow => Error::BufferOverflow,
18            Self::InvalidLen => Error::InvalidLen,
19            Self::Io(e) => Error::Io(e.kind()),
20        }
21    }
22}
23
24impl<E> From<ReadExactError<E>> for Error<E> {
25    fn from(e: ReadExactError<E>) -> Self {
26        match e {
27            ReadExactError::UnexpectedEof => Error::Invalid,
28            ReadExactError::Other(e) => Error::Io(e),
29        }
30    }
31}
32
33impl FrameHeader {
34    pub async fn recv<R>(mut read: R) -> Result<Self, Error<R::Error>>
35    where
36        R: Read,
37    {
38        let mut header_buf = [0; FrameHeader::MAX_LEN];
39        let mut read_offset = 0;
40        let mut read_end = FrameHeader::MIN_LEN;
41
42        loop {
43            read.read_exact(&mut header_buf[read_offset..read_end])
44                .await
45                .map_err(Error::from)?;
46
47            match FrameHeader::deserialize(&header_buf[..read_end]) {
48                Ok((header, _)) => return Ok(header),
49                Err(Error::Incomplete(more)) => {
50                    read_offset = read_end;
51                    read_end += more;
52                }
53                Err(e) => return Err(e.recast()),
54            }
55        }
56    }
57
58    pub async fn send<W>(&self, mut write: W) -> Result<(), Error<W::Error>>
59    where
60        W: Write,
61    {
62        let mut header_buf = [0; FrameHeader::MAX_LEN];
63        let header_len = unwrap!(self.serialize(&mut header_buf));
64
65        write
66            .write_all(&header_buf[..header_len])
67            .await
68            .map_err(Error::Io)
69    }
70
71    pub async fn recv_payload<'a, R>(
72        &self,
73        mut read: R,
74        payload_buf: &'a mut [u8],
75    ) -> Result<&'a [u8], Error<R::Error>>
76    where
77        R: Read,
78    {
79        if (payload_buf.len() as u64) < self.payload_len {
80            Err(Error::BufferOverflow)
81        } else if self.payload_len == 0 {
82            Ok(&[])
83        } else {
84            let payload = &mut payload_buf[..self.payload_len as _];
85
86            read.read_exact(payload).await.map_err(Error::from)?;
87
88            self.mask(payload, 0);
89
90            Ok(payload)
91        }
92    }
93
94    pub async fn send_payload<'a, W>(
95        &'a self,
96        mut write: W,
97        payload: &'a [u8],
98    ) -> Result<(), Error<W::Error>>
99    where
100        W: Write,
101    {
102        let payload_buf_len = payload.len() as u64;
103
104        if payload_buf_len != self.payload_len {
105            Err(Error::InvalidLen)
106        } else if payload.is_empty() {
107            Ok(())
108        } else if self.mask_key.is_none() {
109            write.write_all(payload).await.map_err(Error::Io)
110        } else {
111            let mut buf = [0_u8; 32];
112
113            let mut offset = 0;
114
115            while offset < payload.len() {
116                let len = min(buf.len(), payload.len() - offset);
117
118                let buf = &mut buf[..len];
119
120                buf.copy_from_slice(&payload[offset..offset + len]);
121
122                self.mask(buf, offset);
123
124                write.write_all(buf).await.map_err(Error::Io)?;
125
126                offset += len;
127            }
128
129            Ok(())
130        }
131    }
132}
133
134pub async fn recv<R>(
135    mut read: R,
136    frame_data_buf: &mut [u8],
137) -> Result<(FrameType, usize), Error<R::Error>>
138where
139    R: Read,
140{
141    let header = FrameHeader::recv(&mut read).await?;
142    header.recv_payload(read, frame_data_buf).await?;
143
144    Ok((header.frame_type, header.payload_len as _))
145}
146
147pub async fn send<W>(
148    mut write: W,
149    frame_type: FrameType,
150    mask_key: Option<u32>,
151    frame_data_buf: &[u8],
152) -> Result<(), Error<W::Error>>
153where
154    W: Write,
155{
156    let header = FrameHeader {
157        frame_type,
158        payload_len: frame_data_buf.len() as _,
159        mask_key,
160    };
161
162    header.send(&mut write).await?;
163    header.send_payload(write, frame_data_buf).await
164}