From 529aae6fc368818bf64d21aab217826a6c0745fc Mon Sep 17 00:00:00 2001
From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Date: Mon, 18 Nov 2024 18:58:56 +0100
Subject: [PATCH] wasi/fs: Improve stopping condition for <ReadDir as
 Iterator>::next

When upgrading [Zed](https://github.com/zed-industries/zed/pull/19349) to Rust 1.82 I've encountered a test failure in our test suite. Specifically, one of our extension tests started hanging. I've tracked it down to a call to std::fs::remove_dir_all not returning when an extension is compiled with Rust 1.82
Our extension system uses WASM components, thus I've looked at the diff between 1.81 and 1.82 with respect to WASI and found 736f773844e7ebf05ccb827c17b7ad9eb28aa295

As it turned out, calling remove_dir_all from extension returned io::ErrorKind::NotFound in 1.81;
the underlying issue is that the ReadDir iterator never actually terminates iteration, however since it loops around, with 1.81 we'd come across an entry second time and fail to remove it, since it would've been removed previously.
With 1.82 and 736f773844e7ebf05ccb827c17b7ad9eb28aa295 it is no longer the case, thus we're seeing the hang.

This commit makes ReadDir::next adhere to readdir contract, namely it will no longer call readdir once the returned # of bytes is smaller than the size of a passed-in buffer.
Previously we'd only terminate the loop if readdir returned 0.
---
 library/std/src/sys/pal/wasi/fs.rs | 31 ++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/library/std/src/sys/pal/wasi/fs.rs b/library/std/src/sys/pal/wasi/fs.rs
index 3296c762cca2..b519a0fe6a3a 100644
--- a/library/std/src/sys/pal/wasi/fs.rs
+++ b/library/std/src/sys/pal/wasi/fs.rs
@@ -172,20 +172,22 @@ impl Iterator for ReadDir {
             let offset = if self.offset == self.cap {
                 let cookie = self.cookie.take()?;
                 match self.inner.dir.fd.readdir(&mut self.buf, cookie) {
-                    Ok(bytes) => self.cap = bytes,
+                    Ok(bytes) => {
+                        // No more entries if we read less than buffer size
+                        if bytes < self.buf.len() {
+                            self.cookie = None;
+                            if bytes == 0 {
+                                return None;
+                            }
+                        } else {
+                            self.cookie = Some(cookie);
+                        }
+                        self.cap = self.buf.len();
+                        self.offset = 0;
+                        0
+                    }
                     Err(e) => return Some(Err(e)),
                 }
-                self.offset = 0;
-                self.cookie = Some(cookie);
-
-                // If we didn't actually read anything, this is in theory the
-                // end of the directory.
-                if self.cap == 0 {
-                    self.cookie = None;
-                    return None;
-                }
-
-                0
             } else {
                 self.offset
             };
@@ -197,7 +199,6 @@ impl Iterator for ReadDir {
             // where we last left off.
             let dirent_size = mem::size_of::<wasi::Dirent>();
             if data.len() < dirent_size {
-                assert!(self.cookie.is_some());
                 assert!(self.buf.len() >= dirent_size);
                 self.offset = self.cap;
                 continue;
@@ -218,7 +219,9 @@ impl Iterator for ReadDir {
                 self.offset = self.cap;
                 continue;
             }
-            self.cookie = Some(dirent.d_next);
+            self.cookie.as_mut().map(|cookie| {
+                *cookie = dirent.d_next;
+            });
             self.offset = offset + dirent_size + dirent.d_namlen as usize;
 
             let name = &data[..(dirent.d_namlen as usize)];