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"); .expect("Failed to create TcpStream");
// send number of bytes in payload // 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 stream
.write_all(&payload_size) .write_all(&payload_size)
.await .await
@ -25,19 +25,19 @@ pub async fn send(addr: String, payload: Vec<u8>) -> Vec<u8> {
.expect("Failed to write data to stream"); .expect("Failed to write data to stream");
// get number of bytes in response // get number of bytes in response
let mut nbuf: [u8; 8] = [0; 8]; let mut nbuf: [u8; 4] = [0; 4];
stream stream
.read(&mut nbuf) .read(&mut nbuf)
.await .await
.expect("Failed to get number of bytes to read"); .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 { if n == 0 {
return vec![0; 0]; return vec![0; 0];
} }
// receive response // receive response
let mut buf = vec![0; n]; let mut buf = vec![0; n.try_into().unwrap()];
stream stream
.read(&mut buf) .read(&mut buf)
.await .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 { // tokio::spawn(async move {
loop { loop {
// get number of bytes to receive // get number of bytes to receive
let mut nbuf: [u8; 8] = [0; 8]; let mut nbuf: [u8; 4] = [0; 4];
socket socket
.read(&mut nbuf) .read(&mut nbuf)
.await .await
.expect("Failed to get number of bytes to read"); .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 { if n == 0 {
break; break;
// return; // return;
} }
let mut buf = vec![0; n]; let mut buf = vec![0; n.try_into().unwrap()];
// receive data // receive data
socket socket
@ -50,7 +50,7 @@ pub async fn listen(addr: String, fun: &mut dyn FnMut(Vec<u8>) -> Vec<u8>) {
let response = fun(buf); let response = fun(buf);
// send number of bytes in response // 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 socket
.write_all(&response_size) .write_all(&response_size)
.await .await