sheave_core/readers/
encryption_algorithm.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 futures::ready;
11use tokio::io::{
12    AsyncRead,
13    ReadBuf
14};
15use crate::handshake::EncryptionAlgorithm;
16
17#[doc(hidden)]
18#[derive(Debug)]
19pub struct EncryptionAlgorithmReader<'a, R: AsyncRead> {
20    reader: Pin<&'a mut R>
21}
22
23#[doc(hidden)]
24impl<R: AsyncRead> Future for EncryptionAlgorithmReader<'_, R> {
25    type Output = IOResult<EncryptionAlgorithm>;
26
27    fn poll(mut self: Pin<&mut Self>, cx: &mut FutureContext<'_>) -> Poll<Self::Output> {
28        let mut encryption_algorithm_byte: [u8; 1] = [0; 1];
29        let mut buf = ReadBuf::new(&mut encryption_algorithm_byte);
30        ready!(self.reader.as_mut().poll_read(cx, &mut buf))?;
31        Poll::Ready(Ok(u8::from_be_bytes(encryption_algorithm_byte).into()))
32    }
33}
34
35/// Reads one byte to indicate the encryption algorithm from streams.
36///
37/// # Examples
38///
39/// ```rust
40/// use std::{
41///     io::Result as IOResult,
42///     pin::pin
43/// };
44/// use sheave_core::{
45///     handshake::EncryptionAlgorithm::*,
46///     readers::read_encryption_algorithm
47/// };
48///
49/// #[tokio::main]
50/// async fn main() -> IOResult<()> {
51///     let reader: [u8; 1] = [3; 1];
52///     let result = read_encryption_algorithm(pin!(reader.as_slice())).await?;
53///     assert_eq!(NotEncrypted, result);
54///     Ok(())
55/// }
56/// ```
57pub fn read_encryption_algorithm<R: AsyncRead>(reader: Pin<&mut R>) -> EncryptionAlgorithmReader<'_, R> {
58    EncryptionAlgorithmReader { reader }
59}
60
61#[cfg(test)]
62mod tests {
63    use std::pin::pin;
64    use crate::handshake::EncryptionAlgorithm::*;
65    use super::*;
66
67    #[tokio::test]
68    async fn read_not_encrypted() {
69        let reader: [u8; 1] = [3];
70        let result = read_encryption_algorithm(pin!(reader.as_slice())).await;
71        assert!(result.is_ok());
72        assert_eq!(NotEncrypted, result.unwrap())
73    }
74
75    #[tokio::test]
76    async fn read_diffie_hellman() {
77        let reader: [u8; 1] = [6];
78        let result = read_encryption_algorithm(pin!(reader.as_slice())).await;
79        assert!(result.is_ok());
80        assert_eq!(DiffieHellman, result.unwrap())
81    }
82
83    #[tokio::test]
84    async fn read_xtea() {
85        let reader: [u8; 1] = [8];
86        let result = read_encryption_algorithm(pin!(reader.as_slice())).await;
87        assert!(result.is_ok());
88        assert_eq!(Xtea, result.unwrap())
89    }
90
91    #[tokio::test]
92    async fn read_blowfish() {
93        let reader: [u8; 1] = [9];
94        let result = read_encryption_algorithm(pin!(reader.as_slice())).await;
95        assert!(result.is_ok());
96        assert_eq!(Blowfish, result.unwrap())
97    }
98
99    #[tokio::test]
100    async fn read_other() {
101        let reader: [u8; 1] = [0];
102        let result = read_encryption_algorithm(pin!(reader.as_slice())).await;
103        assert!(result.is_ok());
104        assert_eq!(Other, result.unwrap())
105    }
106}