Просмотр исходного кода

Don't allow 8 bytes for length of payload

Vecna 3 лет назад
Родитель
Сommit
e1f4305733
2 измененных файлов с 8 добавлено и 8 удалено
  1. 4 4
      src/client_net.rs
  2. 4 4
      src/server_net.rs

+ 4 - 4
src/client_net.rs

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

+ 4 - 4
src/server_net.rs

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