Bevy+Ratutui powered Monitoring of Pico-Strike devices
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

Improve handshake net code

+18 -5
+18 -5
src/net.rs
··· 129 129 async fn transport_handshake(tcp: &Async<TcpStream>) -> color_eyre::Result<Client> { 130 130 let (encap, handshake) = ClientBuilder::start(); 131 131 132 - tcp.write_with(|mut tcp| tcp.write_all(encap.serialize().as_bytes())) 133 - .await?; 132 + let msg = encap.serialize(); 133 + let msg = msg.as_bytes(); 134 + let msg_len = u16::try_from(msg.len()).unwrap().to_be_bytes(); 135 + 136 + tcp.write_with(|mut tcp| tcp.write_all(&msg_len)).await?; 137 + tcp.write_with(|mut tcp| tcp.write_all(msg)).await?; 134 138 135 139 let mut payload = vec![0u8; 1024]; 140 + let mut msg_len = [0u8; 2]; 136 141 137 - let read = tcp.read_with(|mut tcp| tcp.read(&mut payload)).await?; 142 + let msg_len = tcp 143 + .read_with(|mut tcp| { 144 + tcp.read_exact(&mut msg_len)?; 145 + Ok(usize::from(u16::from_be_bytes(msg_len))) 146 + }) 147 + .await?; 148 + 149 + tcp.read_with(|mut tcp| tcp.read_exact(&mut payload[..msg_len])) 150 + .await?; 138 151 139 - let transport = handshake.finish(&payload[..read], &NOISE_PSK)?; 152 + let transport = handshake.finish(&payload[..msg_len], &NOISE_PSK)?; 140 153 141 154 Ok(transport) 142 155 } 143 156 144 157 async fn recv(stream: &Async<TcpStream>) -> std::io::Result<Vec<u8>> { 145 - let mut msg_len_buf = [0_u8; 2]; 158 + let mut msg_len_buf = [0u8; 2]; 146 159 let mut msg = vec![0u8; 4096]; 147 160 148 161 stream