datagram_connections/
client.rs

1use crate::host_to_client::HostToClientCommands;
2use crate::{
3    ClientPhase, ClientToHostChallengeCommand, ClientToHostCommands, ClientToHostPacket,
4    ConnectCommand, ConnectResponse, DatagramConnectionsError, HostToClientPacketHeader,
5    InChallengeCommand, Nonce, PacketHeader,
6};
7use datagram::{DatagramDecoder, DatagramEncoder};
8use flood_rs::in_stream::InOctetStream;
9use flood_rs::out_stream::OutOctetStream;
10use flood_rs::{ReadOctetStream, WriteOctetStream};
11use hexify::format_hex;
12use log::{info, trace};
13use secure_random::SecureRandom;
14use std::io;
15
16impl DatagramEncoder for Client {
17    fn encode(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
18        let mut out_stream = OutOctetStream::new();
19
20        let client_to_server_cmd = self
21            .send(data)
22            .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
23
24        client_to_server_cmd.to_stream(&mut out_stream)?;
25        out_stream.write(data)?;
26
27        Ok(out_stream.octets())
28    }
29}
30
31impl DatagramDecoder for Client {
32    fn decode(&mut self, buffer: &[u8]) -> io::Result<Vec<u8>> {
33        self.decode(buffer)
34            .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
35    }
36}
37
38pub struct Client {
39    phase: ClientPhase,
40}
41
42impl Client {
43    pub fn new(mut random: Box<dyn SecureRandom>) -> Self {
44        let phase = ClientPhase::Challenge(Nonce(random.get_random_u64()));
45        Self { phase }
46    }
47
48    pub fn on_challenge(
49        &mut self,
50        cmd: InChallengeCommand,
51    ) -> Result<(), DatagramConnectionsError> {
52        match self.phase {
53            ClientPhase::Challenge(nonce) => {
54                if cmd.nonce != nonce {
55                    return Err(DatagramConnectionsError::WrongNonceInChallenge);
56                }
57                self.phase = ClientPhase::Connecting(nonce, cmd.incoming_server_challenge);
58                Ok(())
59            }
60            _ => Err(DatagramConnectionsError::ReceivedChallengeInWrongPhase),
61        }
62    }
63
64    pub fn on_connect(&mut self, cmd: ConnectResponse) -> Result<(), DatagramConnectionsError> {
65        match self.phase {
66            ClientPhase::Connecting(nonce, _) => {
67                if cmd.nonce != nonce {
68                    return Err(DatagramConnectionsError::WrongNonceWhileConnecting);
69                }
70                info!(
71                    "udp_connections: on_connect connected {}",
72                    cmd.connection_id
73                );
74                self.phase = ClientPhase::Connected(cmd.connection_id);
75                Ok(())
76            }
77            _ => Err(DatagramConnectionsError::ReceiveConnectInWrongPhase),
78        }
79    }
80
81    pub fn on_packet(
82        &mut self,
83        cmd: HostToClientPacketHeader,
84        in_stream: &mut InOctetStream,
85    ) -> Result<Vec<u8>, DatagramConnectionsError> {
86        match self.phase {
87            ClientPhase::Connected(expected_connection_id) => {
88                if cmd.0.connection_id != expected_connection_id {
89                    return Err(DatagramConnectionsError::WrongConnectionId);
90                }
91                let mut target_buffer = vec![0u8; cmd.0.size as usize];
92                in_stream
93                    .read(&mut target_buffer)
94                    .map_err(DatagramConnectionsError::IoError)?;
95                trace!(
96                    "receive packet of size: {} target:{}  {}",
97                    cmd.0.size,
98                    target_buffer.len(),
99                    format_hex(target_buffer.as_slice())
100                );
101                Ok(target_buffer)
102            }
103            _ => Err(DatagramConnectionsError::ReceivedPacketInWrongPhase),
104        }
105    }
106
107    pub fn send_challenge(
108        &mut self,
109    ) -> Result<ClientToHostChallengeCommand, DatagramConnectionsError> {
110        match self.phase {
111            ClientPhase::Challenge(nonce) => Ok(ClientToHostChallengeCommand { nonce }),
112            _ => Err(DatagramConnectionsError::SendChallengeInWrongPhase),
113        }
114    }
115
116    pub fn send_connect_request(&mut self) -> Result<ConnectCommand, DatagramConnectionsError> {
117        match self.phase {
118            ClientPhase::Connecting(nonce, server_challenge) => Ok(ConnectCommand {
119                nonce,
120                server_challenge,
121            }),
122            _ => Err(DatagramConnectionsError::SendConnectRequestInWrongPhase),
123        }
124    }
125
126    pub fn send_packet(
127        &mut self,
128        data: &[u8],
129    ) -> Result<ClientToHostPacket, DatagramConnectionsError> {
130        match self.phase {
131            ClientPhase::Connected(connection_id) => {
132                trace!("send packet: {}", format_hex(data));
133                Ok(ClientToHostPacket {
134                    header: PacketHeader {
135                        connection_id,
136                        size: data.len() as u16,
137                    },
138                    payload: data.to_vec(),
139                })
140            }
141            _ => Err(DatagramConnectionsError::SendPacketInWrongPhase),
142        }
143    }
144
145    pub fn send(&mut self, data: &[u8]) -> Result<ClientToHostCommands, DatagramConnectionsError> {
146        trace!("send: phase: {}", self.phase);
147        match self.phase {
148            ClientPhase::Challenge(_) => {
149                let challenge = self.send_challenge()?;
150                Ok(ClientToHostCommands::ChallengeType(challenge))
151            }
152
153            ClientPhase::Connecting(_, _) => {
154                let connect_request = self.send_connect_request()?;
155                Ok(ClientToHostCommands::ConnectType(connect_request))
156            }
157
158            ClientPhase::Connected(_) => {
159                let packet = self.send_packet(data)?;
160                trace!("sending datagram {:?}", packet);
161                Ok(ClientToHostCommands::PacketType(packet))
162            }
163        }
164    }
165
166    pub fn decode(&mut self, buffer: &[u8]) -> Result<Vec<u8>, DatagramConnectionsError> {
167        let mut in_stream = InOctetStream::new(buffer);
168        let command = HostToClientCommands::from_stream(&mut in_stream)
169            .map_err(DatagramConnectionsError::IoError)?;
170
171        match command {
172            HostToClientCommands::ChallengeType(challenge_command) => {
173                self.on_challenge(challenge_command)?;
174                Ok(vec![])
175            }
176            HostToClientCommands::ConnectType(connect_command) => {
177                self.on_connect(connect_command)?;
178                Ok(vec![])
179            }
180            HostToClientCommands::PacketType(packet_command) => {
181                self.on_packet(packet_command, &mut in_stream)
182            }
183        }
184    }
185}
datagram_connections/client.rs

datagram_connections/
client.rs
