From 4bd44d10b7508ba78ce3a379ae6c4c28773b736a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Thu, 21 Jul 2022 11:23:13 +0300
Subject: [PATCH] getpwent/getgrent: advance iterator only on success

---
 src/PackedGroup.zig | 33 +++++++++++++++++++++++----------
 src/PackedUser.zig  | 22 +++++++++++++++++-----
 src/libnss.zig      |  6 ++++--
 3 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/src/PackedGroup.zig b/src/PackedGroup.zig
index ebe8404..a59f4ee 100644
--- a/src/PackedGroup.zig
+++ b/src/PackedGroup.zig
@@ -1,15 +1,16 @@
 const std = @import("std");
 
+const mem = std.mem;
+const assert = std.debug.assert;
+const Allocator = mem.Allocator;
+const ArrayList = std.ArrayList;
+const BufSet = std.BufSet;
+
 const pad = @import("padding.zig");
 const validate = @import("validate.zig");
 const compress = @import("compress.zig");
 const InvalidRecord = validate.InvalidRecord;
 
-const mem = std.mem;
-const Allocator = mem.Allocator;
-const ArrayList = std.ArrayList;
-const BufSet = std.BufSet;
-
 const PackedGroup = @This();
 
 pub const GroupStored = struct {
@@ -65,18 +66,30 @@ fn validateUtf8(s: []const u8) InvalidRecord!void {
 
 pub const Iterator = struct {
     section: []const u8,
-    nextStart: usize = 0,
+    next_start: usize = 0,
     idx: u32 = 0,
     total: u32,
+    advance_by: usize = 0,
 
     pub fn next(it: *Iterator) ?PackedGroup {
-        if (it.idx == it.total) return null;
+        const gr = it.peek() orelse return null;
+        it.advance();
+        return gr;
+    }
 
-        const entry = fromBytes(it.section[it.nextStart..]);
-        it.nextStart += entry.end;
-        it.idx += 1;
+    pub fn peek(it: *Iterator) ?PackedGroup {
+        if (it.idx == it.total) return null;
+        const entry = fromBytes(it.section[it.next_start..]);
+        it.advance_by = entry.end;
         return entry.group;
     }
+
+    pub fn advance(it: *Iterator) void {
+        assert(it.advance_by > 0);
+        it.idx += 1;
+        it.next_start += it.advance_by;
+        it.advance_by = 0;
+    }
 };
 
 pub fn iterator(section: []const u8, total: u32) Iterator {
diff --git a/src/PackedUser.zig b/src/PackedUser.zig
index 72c1ebb..c5b8a42 100644
--- a/src/PackedUser.zig
+++ b/src/PackedUser.zig
@@ -105,19 +105,31 @@ pub fn fromBytes(bytes: []const u8) Entry {
 
 pub const Iterator = struct {
     section: []const u8,
-    nextStart: usize = 0,
+    next_start: usize = 0,
     shell_reader: ShellReader,
     idx: u32 = 0,
     total: u32,
+    advance_by: usize = 0,
 
     pub fn next(it: *Iterator) ?PackedUser {
-        if (it.idx == it.total) return null;
+        const u = it.peek() orelse return null;
+        it.advance();
+        return u;
+    }
 
-        const entry = PackedUser.fromBytes(it.section[it.nextStart..]);
-        it.nextStart += entry.end;
-        it.idx += 1;
+    pub fn peek(it: *Iterator) ?PackedUser {
+        if (it.idx == it.total) return null;
+        const entry = fromBytes(it.section[it.next_start..]);
+        it.advance_by = entry.end;
         return entry.user;
     }
+
+    pub fn advance(it: *Iterator) void {
+        assert(it.advance_by > 0);
+        it.idx += 1;
+        it.next_start += it.advance_by;
+        it.advance_by = 0;
+    }
 };
 
 pub fn iterator(section: []const u8, total: u32, shell_reader: ShellReader) Iterator {
diff --git a/src/libnss.zig b/src/libnss.zig
index 786c79d..14bbbfb 100644
--- a/src/libnss.zig
+++ b/src/libnss.zig
@@ -370,7 +370,7 @@ fn getgrent_r(
     defer state.getgrent_iterator_mu.unlock();
 
     if (state.getgrent_iterator) |*it| {
-        const group = it.next() orelse {
+        const group = it.peek() orelse {
             errnop.* = 0;
             return c.NSS_STATUS_NOTFOUND;
         };
@@ -381,6 +381,7 @@ fn getgrent_r(
             state.file.db.packCGroup(&group, buffer[0..buflen]);
 
         if (cgroup1) |cgroup| {
+            it.advance();
             result.* = cgroup;
             return c.NSS_STATUS_SUCCESS;
         } else |err| switch (err) {
@@ -422,7 +423,7 @@ fn getpwent_r(
     defer state.getpwent_iterator_mu.unlock();
 
     if (state.getpwent_iterator) |*it| {
-        const user = it.next() orelse {
+        const user = it.peek() orelse {
             errnop.* = 0;
             return c.NSS_STATUS_NOTFOUND;
         };
@@ -435,6 +436,7 @@ fn getpwent_r(
             },
         };
 
+        it.advance();
         result.* = cuser;
         return c.NSS_STATUS_SUCCESS;
     } else {