Don't allow 8 bytes for length of payload
This commit is contained in:
parent
266f53cb31
commit
e1f4305733
|
@ -12,7 +12,7 @@ pub async fn send(addr: String, payload: Vec<u8>) -> Vec<u8> {
|
|||
.expect("Failed to create TcpStream");
|
||||
|
||||
// send number of bytes in payload
|
||||
let payload_size = usize::to_be_bytes(payload.len());
|
||||
let payload_size = u32::to_be_bytes(payload.len().try_into().unwrap());
|
||||
stream
|
||||
.write_all(&payload_size)
|
||||
.await
|
||||
|
@ -25,19 +25,19 @@ pub async fn send(addr: String, payload: Vec<u8>) -> Vec<u8> {
|
|||
.expect("Failed to write data to stream");
|
||||
|
||||
// get number of bytes in response
|
||||
let mut nbuf: [u8; 8] = [0; 8];
|
||||
let mut nbuf: [u8; 4] = [0; 4];
|
||||
stream
|
||||
.read(&mut nbuf)
|
||||
.await
|
||||
.expect("Failed to get number of bytes to read");
|
||||
let n = usize::from_be_bytes(nbuf);
|
||||
let n = u32::from_be_bytes(nbuf);
|
||||
|
||||
if n == 0 {
|
||||
return vec![0; 0];
|
||||
}
|
||||
|
||||
// receive response
|
||||
let mut buf = vec![0; n];
|
||||
let mut buf = vec![0; n.try_into().unwrap()];
|
||||
stream
|
||||
.read(&mut buf)
|
||||
.await
|
||||
|
|
|
@ -22,19 +22,19 @@ pub async fn listen(addr: String, fun: &mut dyn FnMut(Vec<u8>) -> Vec<u8>) {
|
|||
// tokio::spawn(async move {
|
||||
loop {
|
||||
// get number of bytes to receive
|
||||
let mut nbuf: [u8; 8] = [0; 8];
|
||||
let mut nbuf: [u8; 4] = [0; 4];
|
||||
socket
|
||||
.read(&mut nbuf)
|
||||
.await
|
||||
.expect("Failed to get number of bytes to read");
|
||||
let n = usize::from_be_bytes(nbuf);
|
||||
let n = u32::from_be_bytes(nbuf);
|
||||
|
||||
if n == 0 {
|
||||
break;
|
||||
// return;
|
||||
}
|
||||
|
||||
let mut buf = vec![0; n];
|
||||
let mut buf = vec![0; n.try_into().unwrap()];
|
||||
|
||||
// receive data
|
||||
socket
|
||||
|
@ -50,7 +50,7 @@ pub async fn listen(addr: String, fun: &mut dyn FnMut(Vec<u8>) -> Vec<u8>) {
|
|||
let response = fun(buf);
|
||||
|
||||
// send number of bytes in response
|
||||
let response_size = usize::to_be_bytes(response.len());
|
||||
let response_size = u32::to_be_bytes(response.len().try_into().unwrap());
|
||||
socket
|
||||
.write_all(&response_size)
|
||||
.await
|
||||
|
|
Loading…
Reference in New Issue