Bevy+Ratutui powered Monitoring of Pico-Strike devices
Improve handshake net code
+18
-5
+18
-5
src/net.rs
+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