1 files changed, 220 insertions, 0 deletions
diff --git a/src/pwcache.c b/src/pwcache.c
new file mode 100644
index 0000000..af8c237
--- /dev/null
+++ b/src/pwcache.c
@@ -0,0 +1,220 @@
+// Copyright © Tavian Barnes <tavianator@tavianator.com>
+// SPDX-License-Identifier: 0BSD
+
+#include "prelude.h"
+#include "pwcache.h"
+#include "alloc.h"
+#include "trie.h"
+#include <errno.h>
+#include <grp.h>
+#include <pwd.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+/** Represents cache hits for negative results. */
+static void *MISSING = &MISSING;
+
+/** Callback type for bfs_getent(). */
+typedef void *bfs_getent_fn(const void *key, void *ptr, size_t bufsize);
+
+/** Shared scaffolding for get{pw,gr}{nam,?id}_r(). */
+static void *bfs_getent(bfs_getent_fn *fn, const void *key, struct trie_leaf *leaf, struct varena *varena) {
+	if (leaf->value) {
+		errno = 0;
+		return leaf->value == MISSING ? NULL : leaf->value;
+	}
+
+	// _SC_GET{PW,GR}_R_SIZE_MAX tend to be fairly large (~1K).  That's okay
+	// for temporary allocations, but for these long-lived ones, let's start
+	// with a smaller buffer.
+	size_t bufsize = 128;
+	void *ptr = varena_alloc(varena, bufsize);
+	if (!ptr) {
+		return NULL;
+	}
+
+	while (true) {
+		void *ret = fn(key, ptr, bufsize);
+		if (ret) {
+			leaf->value = ret;
+			return ret;
+		} else if (errno == 0) {
+			leaf->value = MISSING;
+			break;
+		} else if (errno == ERANGE) {
+			void *next = varena_grow(varena, ptr, &bufsize);
+			if (!next) {
+				break;
+			}
+			ptr = next;
+		} else {
+			break;
+		}
+	}
+
+	varena_free(varena, ptr, bufsize);
+	return NULL;
+}
+
+/**
+ * An arena-allocated struct passwd.
+ */
+struct bfs_passwd {
+	struct passwd pwd;
+	char buf[];
+};
+
+struct bfs_users {
+	/** bfs_passwd arena. */
+	struct varena varena;
+	/** A map from usernames to entries. */
+	struct trie by_name;
+	/** A map from UIDs to entries. */
+	struct trie by_uid;
+};
+
+struct bfs_users *bfs_users_new(void) {
+	struct bfs_users *users = ALLOC(struct bfs_users);
+	if (!users) {
+		return NULL;
+	}
+
+	VARENA_INIT(&users->varena, struct bfs_passwd, buf);
+	trie_init(&users->by_name);
+	trie_init(&users->by_uid);
+	return users;
+}
+
+/** bfs_getent() callback for getpwnam_r(). */
+static void *bfs_getpwnam_impl(const void *key, void *ptr, size_t bufsize) {
+	struct bfs_passwd *storage = ptr;
+
+	struct passwd *ret = NULL;
+	errno = getpwnam_r(key, &storage->pwd, storage->buf, bufsize, &ret);
+	return ret;
+}
+
+const struct passwd *bfs_getpwnam(struct bfs_users *users, const char *name) {
+	struct trie_leaf *leaf = trie_insert_str(&users->by_name, name);
+	if (!leaf) {
+		return NULL;
+	}
+
+	return bfs_getent(bfs_getpwnam_impl, name, leaf, &users->varena);
+}
+
+/** bfs_getent() callback for getpwuid_r(). */
+static void *bfs_getpwuid_impl(const void *key, void *ptr, size_t bufsize) {
+	const uid_t *uid = key;
+	struct bfs_passwd *storage = ptr;
+
+	struct passwd *ret = NULL;
+	errno = getpwuid_r(*uid, &storage->pwd, storage->buf, bufsize, &ret);
+	return ret;
+}
+
+const struct passwd *bfs_getpwuid(struct bfs_users *users, uid_t uid) {
+	struct trie_leaf *leaf = trie_insert_mem(&users->by_uid, &uid, sizeof(uid));
+	if (!leaf) {
+		return NULL;
+	}
+
+	return bfs_getent(bfs_getpwuid_impl, &uid, leaf, &users->varena);
+}
+
+void bfs_users_flush(struct bfs_users *users) {
+	trie_clear(&users->by_uid);
+	trie_clear(&users->by_name);
+	varena_clear(&users->varena);
+}
+
+void bfs_users_free(struct bfs_users *users) {
+	if (users) {
+		trie_destroy(&users->by_uid);
+		trie_destroy(&users->by_name);
+		varena_destroy(&users->varena);
+		free(users);
+	}
+}
+
+/**
+ * An arena-allocated struct group.
+ */
+struct bfs_group {
+	struct group grp;
+	char buf[];
+};
+
+struct bfs_groups {
+	/** bfs_group arena. */
+	struct varena varena;
+	/** A map from group names to entries. */
+	struct trie by_name;
+	/** A map from GIDs to entries. */
+	struct trie by_gid;
+};
+
+struct bfs_groups *bfs_groups_new(void) {
+	struct bfs_groups *groups = ALLOC(struct bfs_groups);
+	if (!groups) {
+		return NULL;
+	}
+
+	VARENA_INIT(&groups->varena, struct bfs_group, buf);
+	trie_init(&groups->by_name);
+	trie_init(&groups->by_gid);
+	return groups;
+}
+
+/** bfs_getent() callback for getgrnam_r(). */
+static void *bfs_getgrnam_impl(const void *key, void *ptr, size_t bufsize) {
+	struct bfs_group *storage = ptr;
+
+	struct group *ret = NULL;
+	errno = getgrnam_r(key, &storage->grp, storage->buf, bufsize, &ret);
+	return ret;
+}
+
+const struct group *bfs_getgrnam(struct bfs_groups *groups, const char *name) {
+	struct trie_leaf *leaf = trie_insert_str(&groups->by_name, name);
+	if (!leaf) {
+		return NULL;
+	}
+
+	return bfs_getent(bfs_getgrnam_impl, name, leaf, &groups->varena);
+}
+
+/** bfs_getent() callback for getgrgid_r(). */
+static void *bfs_getgrgid_impl(const void *key, void *ptr, size_t bufsize) {
+	const gid_t *gid = key;
+	struct bfs_group *storage = ptr;
+
+	struct group *ret = NULL;
+	errno = getgrgid_r(*gid, &storage->grp, storage->buf, bufsize, &ret);
+	return ret;
+}
+
+const struct group *bfs_getgrgid(struct bfs_groups *groups, gid_t gid) {
+	struct trie_leaf *leaf = trie_insert_mem(&groups->by_gid, &gid, sizeof(gid));
+	if (!leaf) {
+		return NULL;
+	}
+
+	return bfs_getent(bfs_getgrgid_impl, &gid, leaf, &groups->varena);
+}
+
+void bfs_groups_flush(struct bfs_groups *groups) {
+	trie_clear(&groups->by_gid);
+	trie_clear(&groups->by_name);
+	varena_clear(&groups->varena);
+}
+
+void bfs_groups_free(struct bfs_groups *groups) {
+	if (groups) {
+		trie_destroy(&groups->by_gid);
+		trie_destroy(&groups->by_name);
+		varena_destroy(&groups->varena);
+		free(groups);
+	}
+}