Merge pull request #13725 from mathetake/fixreaddir

wasi: fixes IterableDir.nextWasi for large directory

2 files changed, 47 insertions, 1 deletions

diff --git a/lib/std/fs.zig b/lib/std/fs.zig
index 8dfd811930..3ef8e5319c 100644
--- a/lib/std/fs.zig
+++ b/lib/std/fs.zig
@@ -809,7 +809,11 @@ pub const IterableDir = struct {
                 // and we avoid the code complexity here.
                 const w = os.wasi;
                 start_over: while (true) {
-                    if (self.index >= self.end_index) {
+                    // TODO https://github.com/ziglang/zig/issues/12498
+                    _ = @sizeOf(w.dirent_t) + 1;
+                    // According to the WASI spec, the last entry might be truncated,
+                    // so we need to check if the left buffer contains the whole dirent.
+                    if (self.end_index - self.index < @sizeOf(w.dirent_t)) {
                         var bufused: usize = undefined;
                         switch (w.fd_readdir(self.dir.fd, &self.buf, self.buf.len, self.cookie, &bufused)) {
                             .SUCCESS => {},
@@ -828,6 +832,11 @@ pub const IterableDir = struct {
                     const entry = @ptrCast(*align(1) w.dirent_t, &self.buf[self.index]);
                     const entry_size = @sizeOf(w.dirent_t);
                     const name_index = self.index + entry_size;
+                    if (name_index + entry.d_namlen > self.end_index) {
+                        // This case, the name is truncated, so we need to call readdir to store the entire name.
+                        self.end_index = self.index; // Force fd_readdir in the next loop.
+                        continue :start_over;
+                    }
                     const name = mem.span(self.buf[name_index .. name_index + entry.d_namlen]);
 
                     const next_index = name_index + entry.d_namlen;
diff --git a/lib/std/fs/test.zig b/lib/std/fs/test.zig
index 00e42b6417..556f3b8459 100644
--- a/lib/std/fs/test.zig
+++ b/lib/std/fs/test.zig
@@ -186,6 +186,43 @@ test "Dir.Iterator" {
     try testing.expect(contains(&entries, .{ .name = "some_dir", .kind = .Directory }));
 }
 
+test "Dir.Iterator many entries" {
+    if (builtin.os.tag == .wasi and !builtin.link_libc) try os.initPreopensWasi(std.heap.page_allocator, "/");
+
+    var tmp_dir = tmpIterableDir(.{});
+    defer tmp_dir.cleanup();
+
+    const num = 1024;
+    var i: usize = 0;
+    var buf: [4]u8 = undefined; // Enough to store "1024".
+    while (i < num) : (i += 1) {
+        const name = try std.fmt.bufPrint(&buf, "{}", .{i});
+        const file = try tmp_dir.iterable_dir.dir.createFile(name, .{});
+        file.close();
+    }
+
+    var arena = ArenaAllocator.init(testing.allocator);
+    defer arena.deinit();
+    const allocator = arena.allocator();
+
+    var entries = std.ArrayList(IterableDir.Entry).init(allocator);
+
+    // Create iterator.
+    var iter = tmp_dir.iterable_dir.iterate();
+    while (try iter.next()) |entry| {
+        // We cannot just store `entry` as on Windows, we're re-using the name buffer
+        // which means we'll actually share the `name` pointer between entries!
+        const name = try allocator.dupe(u8, entry.name);
+        try entries.append(.{ .name = name, .kind = entry.kind });
+    }
+
+    i = 0;
+    while (i < num) : (i += 1) {
+        const name = try std.fmt.bufPrint(&buf, "{}", .{i});
+        try testing.expect(contains(&entries, .{ .name = name, .kind = .File }));
+    }
+}
+
 test "Dir.Iterator twice" {
     var tmp_dir = tmpIterableDir(.{});
     defer tmp_dir.cleanup();