socket.recv read amount of bytes asked

palaviv · palaviv · commit 2bdbc40a4599 · 2019-02-27T19:50:31.000+02:00
diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py
@@ -12,9 +12,13 @@
 message_b = b'bbbbb'
 
 connector.send(message_a)
-connector.close()
-recv_a = connection.recv(10)
+connection.send(message_b)
+recv_a = connection.recv(len(message_a))
+recv_b = connector.recv(len(message_b))
+assert recv_a == message_a
+assert recv_b == message_b
 
 connection.close()
+connector.close()
 listener.close()
 
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
@@ -223,10 +223,10 @@ fn socket_recv(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
     match zelf.payload {
         PyObjectPayload::Socket { ref socket } => {
-            let mut buffer = Vec::new();
-            let _temp = match socket.borrow_mut().con {
-                Some(ref mut v) => v.read_to_end(&mut buffer).unwrap(),
-                None => 0,
+            let mut buffer = vec![0u8; objint::get_value(bufsize).to_usize().unwrap()];
+            match socket.borrow_mut().con {
+                Some(ref mut v) => v.read_exact(&mut buffer).unwrap(),
+                None => return Err(vm.new_type_error("".to_string())),
             };
             Ok(vm.ctx.new_bytes(buffer))
         }
