async_ws/frame/frame_head/
decode.rs

1use crate::frame::decode::FrameDecodeError;
2use crate::frame::{FrameHead, FrameHeadParseError};
3use futures::prelude::*;
4use std::io;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8#[derive(Debug)]
9pub struct FrameHeadDecode<T: AsyncRead + Unpin> {
10    transport: Option<T>,
11    state: FrameHeadDecodeState,
12}
13
14impl<T: AsyncRead + Unpin> FrameHeadDecode<T> {
15    pub fn checkpoint(self) -> Option<(T, FrameHeadDecodeState)> {
16        let (transport, state) = (self.transport, self.state);
17        transport.map(|transport| (transport, state))
18    }
19}
20
21#[derive(Debug)]
22pub struct FrameHeadDecodeState {
23    buffer: [u8; 14],
24    buffer_len: usize,
25}
26
27impl FrameHeadDecodeState {
28    pub fn new() -> Self {
29        Self {
30            buffer: [0u8; 14],
31            buffer_len: 0,
32        }
33    }
34    pub fn restore<T: AsyncRead + Unpin>(self, transport: T) -> FrameHeadDecode<T> {
35        FrameHeadDecode {
36            transport: Some(transport),
37            state: self,
38        }
39    }
40    pub fn poll<T: AsyncRead + Unpin>(
41        &mut self,
42        transport: &mut T,
43        cx: &mut Context<'_>,
44    ) -> Poll<Result<FrameHead, FrameDecodeError>> {
45        loop {
46            let min = match FrameHead::parse(&self.buffer[0..self.buffer_len]) {
47                Ok(info) => return Poll::Ready(Ok(info)),
48                Err(FrameHeadParseError::Incomplete(min)) => min,
49                Err(err) => return Poll::Ready(Err(err.into())),
50            };
51            let buffer_len = self.buffer_len;
52            let read_window = &mut self.buffer[buffer_len..min];
53            match Pin::new(&mut *transport).poll_read(cx, read_window) {
54                Poll::Ready(Ok(0)) => {
55                    return Poll::Ready(Err(FrameDecodeError::Io(
56                        io::ErrorKind::UnexpectedEof.into(),
57                    )))
58                }
59                Poll::Ready(Ok(n)) => self.buffer_len += n,
60                Poll::Ready(Err(err)) => return Poll::Ready(Err(err.into())),
61                Poll::Pending => return Poll::Pending,
62            }
63        }
64    }
65}
66
67impl<T: AsyncRead + Unpin> Future for FrameHeadDecode<T> {
68    type Output = Result<(T, FrameHead), FrameDecodeError>;
69
70    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
71        let mut transport = self.transport.take().unwrap();
72        match self.state.poll(&mut transport, cx) {
73            Poll::Ready(Ok(frame_head)) => Poll::Ready(Ok((transport, frame_head))),
74            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
75            Poll::Pending => {
76                self.transport = Some(transport);
77                Poll::Pending
78            }
79        }
80    }
81}