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

sheave_core/writers/
encryption_algorithm.rs
