turbonss/lib/Corpus.zig

const std = @import("std");
const mem = std.mem;
const math = std.math;
const sort = std.sort;
const unicode = std.unicode;
const Allocator = std.mem.Allocator;
const ArenaAllocator = std.heap.ArenaAllocator;
const StringHashMap = std.StringHashMap;
const MultiArrayList = std.MultiArrayList;
const ArrayListUnmanaged = std.ArrayListUnmanaged;

const User = @import("User.zig");
const Group = @import("Group.zig");

pub const Corpus = @This();

arena: ArenaAllocator,

// sorted by name, by unicode codepoint
users: MultiArrayList(User),
// sorted by gid
groups: MultiArrayList(Group),

name2user: StringHashMap(u32),
name2group: StringHashMap(u32),
group2users: []const []const u32,
user2groups: []const []const u32,

getgr_bufsize: usize,
getpw_bufsize: usize,

pub fn init(
    baseAllocator: Allocator,
    usersConst: []const User,
    groupsConst: []const Group,
) error{ OutOfMemory, InvalidUtf8, Duplicate, NotFound, TooMany }!Corpus {
    if (usersConst.len >= math.maxInt(u32)) return error.TooMany;
    if (groupsConst.len >= math.maxInt(u32)) return error.TooMany;

    var arena = ArenaAllocator.init(baseAllocator);
    var allocator = arena.allocator();
    errdefer arena.deinit();

    var groups_arr = try allocator.alloc(Group, groupsConst.len);
    var users_arr = try allocator.alloc(User, usersConst.len);
    var getgr_bufsize: usize = 0;
    for (groupsConst) |*group, i| {
        groups_arr[i] = try group.clone(allocator);
        getgr_bufsize = math.max(getgr_bufsize, group.strlenZ());
    }
    var getpw_bufsize: usize = 0;
    for (usersConst) |*user, i| {
        users_arr[i] = try user.clone(allocator);
        getpw_bufsize = math.max(getpw_bufsize, user.strlenZ());
    }

    sort.sort(User, users_arr, {}, cmpUser);
    sort.sort(Group, groups_arr, {}, cmpGroup);

    var users = MultiArrayList(User){};
    try users.ensureTotalCapacity(allocator, users_arr.len);
    for (users_arr) |user|
        users.appendAssumeCapacity(user);
    var groups = MultiArrayList(Group){};
    try groups.ensureTotalCapacity(allocator, groups_arr.len);
    for (groups_arr) |group|
        groups.appendAssumeCapacity(group);

    var name2user = StringHashMap(u32).init(allocator);
    var name2group = StringHashMap(u32).init(allocator);
    for (users.items(.name)) |name, i| {
        var res1 = try name2user.getOrPut(name);
        if (res1.found_existing)
            return error.Duplicate;
        res1.value_ptr.* = @intCast(u32, i);
    }

    for (groups.items(.name)) |name, i| {
        var res1 = try name2group.getOrPut(name);
        if (res1.found_existing)
            return error.Duplicate;
        res1.value_ptr.* = @intCast(u32, i);
    }

    var group2users = try allocator.alloc([]u32, groups.len);

    // uses baseAllocator, because it will be freed before
    // returning from this function. This keeps the arena clean.
    var user2groups = try baseAllocator.alloc(ArrayListUnmanaged(u32), users.len);
    defer baseAllocator.free(user2groups);
    mem.set(ArrayListUnmanaged(u32), user2groups, ArrayListUnmanaged(u32){});

    for (groups.items(.members)) |groupmembers, i| {
        var members = try allocator.alloc(u32, groupmembers.len);
        members.len = 0;

        for (groupmembers) |member_name| {
            if (name2user.get(member_name)) |user_idx| {
                members.len += 1;
                members[members.len - 1] = user_idx;
                try user2groups[user_idx].append(allocator, @intCast(u32, i));
            } else return error.NotFound;
        }

        group2users[i] = members;
    }

    for (group2users) |*groupusers|
        sort.sort(u32, groupusers.*, {}, comptime sort.asc(u32));

    var user2groups_final = try allocator.alloc([]const u32, users.len);
    user2groups_final.len = users.len;
    for (user2groups) |*usergroups, i| {
        sort.sort(u32, usergroups.items, {}, comptime sort.asc(u32));
        user2groups_final[i] = usergroups.toOwnedSlice(allocator);
    }

    return Corpus{
        .arena = arena,
        .users = users,
        .groups = groups,
        .name2user = name2user,
        .name2group = name2group,
        .group2users = group2users,
        .user2groups = user2groups_final,
        .getgr_bufsize = getgr_bufsize,
        .getpw_bufsize = getpw_bufsize,
    };
}

pub fn deinit(self: *Corpus) void {
    self.arena.deinit();
    self.* = undefined;
}

// cmpUser compares two users for sorting. By username's utf8 codepoints, ascending.
fn cmpUser(_: void, a: User, b: User) bool {
    var utf8_a = (unicode.Utf8View.init(a.name) catch unreachable).iterator();
    var utf8_b = (unicode.Utf8View.init(b.name) catch unreachable).iterator();
    while (utf8_a.nextCodepoint()) |codepoint_a| {
        if (utf8_b.nextCodepoint()) |codepoint_b| {
            if (codepoint_a == codepoint_b) {
                continue;
            } else return codepoint_a < codepoint_b;
        }

        // a is a prefix of b. It is thus shorter.
        return false;
    }
    // b is a prefix of a
    return true;
}

fn cmpGroup(_: void, a: Group, b: Group) bool {
    return a.gid < b.gid;
}

fn testUser(name: []const u8) User {
    var result = mem.zeroes(User);
    result.name = name;
    return result;
}

const testing = std.testing;
const someMembers = @import("Group.zig").someMembers;

test "users compare function" {
    const a = testUser("a");
    const b = testUser("b");
    const bb = testUser("bb");
    try testing.expect(cmpUser({}, a, b));
    try testing.expect(!cmpUser({}, b, a));
    try testing.expect(cmpUser({}, a, bb));
    try testing.expect(!cmpUser({}, bb, a));
    try testing.expect(cmpUser({}, b, bb));
    try testing.expect(!cmpUser({}, bb, b));
}

pub fn testCorpus(allocator: Allocator) !Corpus {
    const users = [_]User{ User{
        .uid = 0,
        .gid = 0,
        .name = "root",
        .gecos = "",
        .home = "/root",
        .shell = "/bin/bash",
    }, User{
        .uid = 128,
        .gid = 128,
        .name = "vidmantas",
        .gecos = "Vidmantas Kaminskas",
        .home = "/home/vidmantas",
        .shell = "/bin/bash",
    }, User{
        .uid = 1000,
        .gid = math.maxInt(u32),
        .name = "Name" ** 8,
        .gecos = "Gecos" ** 51,
        .home = "Home" ** 16,
        .shell = "She.LllL" ** 8,
    }, User{
        .uid = 100000,
        .gid = 1002,
        .name = "svc-bar",
        .gecos = "",
        .home = "/",
        .shell = "/",
    }, User{
        .uid = 65534,
        .gid = 65534,
        .name = "nobody",
        .gecos = "nobody",
        .home = "/nonexistent",
        .shell = "/usr/sbin/nologin",
    } };

    var group0 = try Group.init(allocator, 0, "root", &[_][]const u8{"root"});
    var group1 = try Group.init(allocator, 128, "vidmantas", &[_][]const u8{"vidmantas"});
    const members2 = &[_][]const u8{ "svc-bar", "Name" ** 8, "vidmantas", "root" };
    var group2 = try Group.init(allocator, 9999, "all", members2);
    const members3 = &[_][]const u8{ "svc-bar", "vidmantas" };
    var group3 = try Group.init(allocator, 100000, "service-account", members3);
    defer group0.deinit(allocator);
    defer group1.deinit(allocator);
    defer group2.deinit(allocator);
    defer group3.deinit(allocator);

    const groups = [_]Group{ group0, group1, group2, group3 };

    return try Corpus.init(allocator, users[0..], groups[0..]);
}

test "test corpus" {
    var corpus = try testCorpus(testing.allocator);
    defer corpus.deinit();

    const name_name = 0;
    const nobody = 1;
    const root = 2;
    const svc_bar = 3;
    const vidmantas = 4;

    const usernames = corpus.users.items(.name);
    try testing.expectEqualStrings(usernames[name_name], "Name" ** 8);
    try testing.expectEqualStrings(usernames[nobody], "nobody");
    try testing.expectEqualStrings(usernames[root], "root");
    try testing.expectEqualStrings(usernames[svc_bar], "svc-bar");
    try testing.expectEqualStrings(usernames[vidmantas], "vidmantas");

    const g_root = 0;
    const g_vidmantas = 1;
    const g_all = 2;
    const g_service_account = 3;

    const groupnames = corpus.groups.items(.name);
    try testing.expectEqualStrings(groupnames[g_root], "root");
    try testing.expectEqualStrings(groupnames[g_service_account], "service-account");
    try testing.expectEqualStrings(groupnames[g_vidmantas], "vidmantas");
    try testing.expectEqualStrings(groupnames[g_all], "all");

    try testing.expectEqual(corpus.name2user.get("404"), null);
    try testing.expectEqual(corpus.name2user.get("vidmantas").?, vidmantas);
    try testing.expectEqual(corpus.name2group.get("404"), null);
    try testing.expectEqual(corpus.name2group.get("vidmantas").?, g_vidmantas);

    const membersOfAll = corpus.group2users[g_all];
    try testing.expectEqual(membersOfAll[0], name_name);
    try testing.expectEqual(membersOfAll[1], root);
    try testing.expectEqual(membersOfAll[2], svc_bar);
    try testing.expectEqual(membersOfAll[3], vidmantas);

    const groupsOfVidmantas = corpus.user2groups[vidmantas];
    try testing.expectEqual(groupsOfVidmantas[0], g_vidmantas);
    try testing.expectEqual(groupsOfVidmantas[1], g_all);
    try testing.expectEqual(groupsOfVidmantas[2], g_service_account);
}