sheave_core/writers/
basic_header.rs

1use std::{
2    future::Future,
3    io::Result as IOResult,
4    pin::Pin,
5    task::{
6        Context as FutureContext,
7        Poll
8    }
9};
10use tokio::io::AsyncWrite;
11use crate::messages::headers::BasicHeader;
12
13#[doc(hidden)]
14#[derive(Debug)]
15pub struct BasicHeaderWriter<'a, W: AsyncWrite> {
16    writer: Pin<&'a mut W>,
17    basic_header: &'a BasicHeader
18}
19
20#[doc(hidden)]
21impl<W: AsyncWrite> Future for BasicHeaderWriter<'_, W> {
22    type Output = IOResult<()>;
23
24    fn poll(mut self: Pin<&mut Self>, cx: &mut FutureContext<'_>) -> Poll<Self::Output> {
25        let message_format = self.basic_header.get_message_format() as u8;
26        let chunk_id = self.basic_header.get_chunk_id();
27        let mut basic_header_bytes: Vec<u8> = Vec::new();
28        if chunk_id >= 320 {
29            basic_header_bytes.push(message_format << 6 | 1);
30            basic_header_bytes.extend_from_slice((chunk_id - 64).to_le_bytes().as_slice());
31        } else if chunk_id >= 64 {
32            basic_header_bytes.push(message_format << 6);
33            basic_header_bytes.push((chunk_id - 64) as u8);
34        } else {
35            basic_header_bytes.push(message_format << 6 | (chunk_id as u8));
36        }
37        self.writer.as_mut().poll_write(cx, basic_header_bytes.as_slice()).map_ok(|_| ())
38    }
39}
40
41/// Writes a basic header into streams.
42///
43/// # Examples
44///
45/// ```rust
46/// use std::{
47///     io::Result as IOResult,
48///     pin::{
49///         Pin,
50///         pin
51///     }
52/// };
53/// use sheave_core::{
54///     messages::headers::{
55///         BasicHeader,
56///         MessageFormat::*
57///     },
58///     writers::write_basic_header
59/// };
60///
61/// #[tokio::main]
62/// async fn main() -> IOResult<()> {
63///     // In case of 1 byte.
64///     let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
65///     let message_format = New;
66///     let chunk_id = 2u16;
67///     let basic_header = BasicHeader::new(message_format, chunk_id);
68///     let result = write_basic_header(writer.as_mut(), &basic_header).await;
69///     assert!(result.is_ok());
70///     assert_eq!(message_format as u8, writer[0] >> 6);
71///     assert_eq!(chunk_id, (writer[0] << 2 >> 2) as u16);
72///
73///     // In case of 2 bytes.
74///     let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
75///     let message_fomrat = New;
76///     let chunk_id = 64u16;
77///     let basic_header = BasicHeader::new(message_fomrat, chunk_id);
78///     let result = write_basic_header(writer.as_mut(), &basic_header).await;
79///     assert!(result.is_ok());
80///     assert_eq!(message_format as u8, writer[0] >> 6);
81///     assert_eq!(0, writer[0] << 2 >> 2);
82///     assert_eq!(0, writer[1] as u16);
83///
84///     // In case of 3 bytes.
85///     let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
86///     let message_format = New;
87///     let chunk_id = 320u16;
88///     let basic_header = BasicHeader::new(message_format, chunk_id);
89///     let result = write_basic_header(writer.as_mut(), &basic_header).await;
90///     let mut written: [u8; 2] = [0; 2];
91///     written.copy_from_slice(&writer[1..]);
92///     let written = u16::from_le_bytes(written);
93///     assert!(result.is_ok());
94///     assert_eq!(message_format as u8, writer[0] >> 6);
95///     assert_eq!(1, writer[0] << 2 >> 2);
96///     assert_eq!(256, written);
97///     Ok(())
98/// }
99/// ```
100pub fn write_basic_header<'a, W: AsyncWrite>(writer: Pin<&'a mut W>, basic_header: &'a BasicHeader) -> BasicHeaderWriter<'a, W> {
101    BasicHeaderWriter { writer, basic_header }
102}
103
104#[cfg(test)]
105mod tests {
106    use std::{
107        cmp::max,
108        pin::{
109            Pin,
110            pin
111        }
112    };
113    use rand::random;
114    use crate::messages::headers::{
115        BasicHeader,
116        MessageFormat
117    };
118    use super::*;
119
120    #[tokio::test]
121    async fn write_one_byte() {
122        let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
123        let byte = random::<u8>();
124        let message_format: MessageFormat = (byte >> 6).into();
125        let chunk_id = (byte << 2 >> 2) as u16;
126        let basic_header = BasicHeader::new(message_format, chunk_id);
127        let result = write_basic_header(writer.as_mut(), &basic_header).await;
128        assert!(result.is_ok());
129        assert_eq!(message_format as u8, writer[0] >> 6);
130        assert_eq!(chunk_id, (writer[0] << 2 >> 2) as u16)
131    }
132
133    #[tokio::test]
134    async fn write_two_bytes() {
135        let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
136        let message_format: MessageFormat = (random::<u8>() >> 6).into();
137        let chunk_id = max(64, random::<u8>()) as u16;
138        let basic_header = BasicHeader::new(message_format, chunk_id);
139        let result = write_basic_header(writer.as_mut(), &basic_header).await;
140        assert!(result.is_ok());
141        assert_eq!(message_format as u8, writer[0] >> 6);
142        assert_eq!(0, writer[0] << 2 >> 2);
143        assert_eq!(chunk_id - 64, writer[1] as u16)
144    }
145
146    #[tokio::test]
147    async fn write_three_bytes() {
148        let mut writer: Pin<&mut Vec<u8>> = pin!(Vec::new());
149        let message_format: MessageFormat = (random::<u8>() >> 6).into();
150        let chunk_id = max(320, random::<u16>());
151        let basic_header = BasicHeader::new(message_format, chunk_id);
152        let result = write_basic_header(writer.as_mut(), &basic_header).await;
153        assert!(result.is_ok());
154        assert_eq!(message_format as u8, writer[0] >> 6);
155        assert_eq!(1, writer[0] << 2 >> 2);
156        let mut written: [u8; 2] = [0; 2];
157        written.copy_from_slice(&writer[1..]);
158        let written = u16::from_le_bytes(written);
159        assert_eq!(chunk_id - 64, written)
160    }
161}
sheave_core/writers/basic_header.rs

sheave_core/writers/
basic_header.rs
