Don't allow 8 bytes for length of payload

This commit is contained in:
Vecna 2023-02-05 16:26:14 -05:00
parent 266f53cb31
commit e1f4305733
2 changed files with 8 additions and 8 deletions

View File

@ -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

View File

@ -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