Handle ERROR_MORE_DATA when scheduling read for Window named pipes

Author: toksdotdev

src/sys/windows/named_pipe.rs

@@ -7,8 +7,8 @@ use std::sync::{Arc, Mutex};
 use std::{fmt, mem, slice};
 
 use windows_sys::Win32::Foundation::{
-    ERROR_BROKEN_PIPE, ERROR_IO_INCOMPLETE, ERROR_IO_PENDING, ERROR_NO_DATA, ERROR_PIPE_CONNECTED,
-    ERROR_PIPE_LISTENING, HANDLE, INVALID_HANDLE_VALUE,
+    ERROR_BROKEN_PIPE, ERROR_IO_INCOMPLETE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_NO_DATA,
+    ERROR_PIPE_CONNECTED, ERROR_PIPE_LISTENING, HANDLE, INVALID_HANDLE_VALUE,
 };
 use windows_sys::Win32::Storage::FileSystem::{
     ReadFile, WriteFile, FILE_FLAG_FIRST_PIPE_INSTANCE, FILE_FLAG_OVERLAPPED, PIPE_ACCESS_DUPLEX,
@@ -705,10 +705,13 @@ impl Inner {
     /// Schedules a read to happen in the background, executing an overlapped
     /// operation.
     ///
-    /// This function returns `true` if a normal error happens or if the read
-    /// is scheduled in the background. If the pipe is no longer connected
-    /// (ERROR_PIPE_LISTENING) then `false` is returned and no read is
-    /// scheduled.
+    /// This function returns `true` if either of the following conditions are met:
+    /// * A normal error happens
+    /// * The read is scheduled in the background
+    /// * Data is already available to be read (ERROR_MORE_DATA)
+    ///
+    /// If the pipe is no longer connected (ERROR_PIPE_LISTENING) then `false` is
+    /// returned and no read is scheduled.
     fn schedule_read(me: &Arc<Inner>, io: &mut Io, events: Option<&mut Vec<Event>>) -> bool {
         // Check to see if a read is already scheduled/completed
         match io.read {
@@ -736,6 +739,20 @@ impl Inner {
             // we just need to wait for a connect.
             Err(ref e) if e.raw_os_error() == Some(ERROR_PIPE_LISTENING as i32) => false,
 
+            // If ERROR_MORE_DATA is returned, it means the slice of unused capacity of the
+            // buffer provided is less than the amount of data available to be read. So
+            // prioritize draining the buffer before scheduling a new read.
+            //
+            // Return `true` to indicate that an overlapped read was scheduled "successfully",
+            // without actually scheduling it. Instead, update `io.read` to `State::Ok(buf, 0)`
+            // to ensure that the next `std::io::Read::read` call is presented still with the
+            // unread data to read from.
+            Err(ref e) if e.raw_os_error() == Some(ERROR_MORE_DATA as i32) => {
+                io.read = State::Ok(buf, 0);
+                mem::forget(me.clone());
+                true
+            }
+
             // If some other error happened, though, we're now readable to give
             // out the error.
             Err(e) => {
@@ -898,6 +915,13 @@ fn read_done(status: &OVERLAPPED_ENTRY, events: Option<&mut Vec<Event>>) {
                 buf.set_len(status.bytes_transferred() as usize);
                 io.read = State::Ok(buf, 0);
             }
+            // This is non-fatal. The buffer was simply too small for the entire message.
+            // Deliver the bytes we got, and if the caller wants to read the rest of the
+            // message, they can initiate another read.
+            Err(e) if e.raw_os_error() == Some(ERROR_MORE_DATA as i32) => {
+                buf.set_len(status.bytes_transferred() as usize);
+                io.read = State::Ok(buf, 0);
+            }
             Err(e) => {
                 debug_assert_eq!(status.bytes_transferred(), 0);
                 io.read = State::Err(e);

tests/win_named_pipe.rs

@@ -10,6 +10,9 @@ use mio::windows::NamedPipe;
 use mio::{Events, Interest, Poll, Token};
 use windows_sys::Win32::{Foundation::ERROR_NO_DATA, Storage::FileSystem::FILE_FLAG_OVERLAPPED};
 
+mod util;
+use util::{expect_events, ExpectEvent};
+
 fn _assert_kinds() {
     fn _assert_send<T: Send>() {}
     fn _assert_sync<T: Sync>() {}
@@ -347,3 +350,101 @@ fn reregister_deregister_different_poll() {
         io::ErrorKind::AlreadyExists,
     );
 }
+
+#[test]
+fn read_message_larger_than_internal_buffer() {
+    let (mut server, mut client) = pipe();
+    let mut poll = t!(Poll::new());
+    t!(poll.registry().register(
+        &mut server,
+        Token(0),
+        Interest::READABLE | Interest::WRITABLE,
+    ));
+    t!(poll.registry().register(
+        &mut client,
+        Token(1),
+        Interest::READABLE | Interest::WRITABLE,
+    ));
+    let mut events = Events::with_capacity(128);
+    t!(poll.poll(&mut events, None));
+
+    // Send message larger than the IPC kernel buffer (4096 bytes)
+    let expected_msg = vec![0x5u8; 8192];
+    assert_eq!(t!(client.write(&expected_msg)), 8192);
+
+    expect_events(
+        &mut poll,
+        &mut events,
+        vec![ExpectEvent::new(Token(0), Interest::READABLE)],
+    );
+
+    let mut buf = [0u8; 4000];
+    let mut actual_msg = Vec::new();
+
+    loop {
+        match server.read(&mut buf) {
+            Ok(0) => break,
+            Ok(n) => {
+                actual_msg.extend_from_slice(&buf[..n]);
+                if actual_msg.len() >= expected_msg.len() {
+                    break;
+                }
+            }
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                t!(poll.poll(&mut events, Some(Duration::from_secs(1))));
+            }
+            Err(e) => panic!("error reading message: {e}"),
+        }
+    }
+
+    assert_eq!(expected_msg, actual_msg);
+}
+
+#[test]
+fn read_with_small_buffer_provided() {
+    let (mut server, mut client) = pipe();
+    let mut poll = t!(Poll::new());
+    t!(poll.registry().register(
+        &mut server,
+        Token(0),
+        Interest::READABLE | Interest::WRITABLE,
+    ));
+    t!(poll.registry().register(
+        &mut client,
+        Token(1),
+        Interest::READABLE | Interest::WRITABLE,
+    ));
+
+    let mut events = Events::with_capacity(128);
+    t!(poll.poll(&mut events, None));
+
+    let expected_msg = vec![1u8; 10000];
+    assert_eq!(t!(client.write(&expected_msg)), 10000);
+
+    expect_events(
+        &mut poll,
+        &mut events,
+        vec![ExpectEvent::new(Token(0), Interest::READABLE)],
+    );
+
+    let mut buf = [0u8; 128];
+    let mut actual_msg = Vec::new();
+
+    loop {
+        match server.read(&mut buf) {
+            Ok(0) => break,
+            Ok(n) => {
+                actual_msg.extend_from_slice(&buf[..n]);
+                if actual_msg.len() >= expected_msg.len() {
+                    break;
+                }
+            }
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                t!(poll.poll(&mut events, Some(Duration::from_millis(100))));
+            }
+            Err(e) => panic!("error reading message: {e}"),
+        }
+    }
+
+    assert_eq!(actual_msg, expected_msg);
+}