turbonss/src/PackedUser.zig

const std = @import("std");
const assert = std.debug.assert;
const mem = std.mem;
const math = std.math;
const Allocator = mem.Allocator;
const ArrayListAligned = std.ArrayListAligned;
const StringHashMap = std.StringHashMap;
const fieldInfo = std.meta.fieldInfo;

const validate = @import("validate.zig");
const compress = @import("compress.zig");
const ShellReader = @import("shell.zig").ShellReader;
const InvalidRecord = validate.InvalidRecord;
const User = @import("User.zig");

const PackedUser = @This();
pub const alignment_bits = 3;

const Inner = packed struct {
    uid: u32,
    gid: u32,
    shell_len_or_idx: u8,
    shell_here: bool,
    name_is_a_suffix: bool,
    home_len: u6,
    name_len: u5,
    gecos_len: u11,

    fn homeLen(self: *const Inner) usize {
        return @as(u32, self.home_len) + 1;
    }

    fn nameStart(self: *const Inner) usize {
        const name_len = self.nameLen();
        if (self.name_is_a_suffix) {
            return self.homeLen() - name_len;
        } else return self.homeLen();
    }

    fn nameLen(self: *const Inner) usize {
        return @as(u32, self.name_len) + 1;
    }

    fn gecosStart(self: *const Inner) usize {
        if (self.name_is_a_suffix) {
            return self.homeLen();
        } else return self.homeLen() + self.nameLen();
    }

    fn gecosLen(self: *const Inner) usize {
        return self.gecos_len;
    }

    fn maybeShellStart(self: *const Inner) usize {
        assert(self.shell_here);
        return self.gecosStart() + self.gecosLen();
    }

    fn shellLen(self: *const Inner) usize {
        return @as(u32, self.shell_len_or_idx) + 1;
    }

    // stringLength returns the length of the blob storing string values.
    fn stringLength(self: *const Inner) usize {
        var result: usize = self.homeLen() + self.gecosLen();
        if (!self.name_is_a_suffix)
            result += self.nameLen();
        if (self.shell_here)
            result += self.shellLen();
        return result;
    }
};

// PackedUser does not allocate; it re-interprets the var_payload
// field. Both of those fields are pointers to "our representation" of
// that field.
inner: *align(8) const Inner,
var_payload: []const u8,
additional_gids_offset: u64,

pub const Entry = struct {
    user: PackedUser,
    end: usize,
};

pub fn fromBytes(blob: []align(8) const u8) error{Overflow}!Entry {
    const start_var_payload = @bitSizeOf(Inner) / 8;
    const inner = @as(*align(8) const Inner, @ptrCast(blob[0..start_var_payload]));
    const end_strings = start_var_payload + inner.stringLength();
    const gids_offset = try compress.uvarint(blob[end_strings..]);
    const end_payload = end_strings + gids_offset.bytes_read;

    return Entry{
        .user = PackedUser{
            .inner = inner,
            .var_payload = blob[start_var_payload..end_payload],
            .additional_gids_offset = gids_offset.value,
        },
        .end = mem.alignForward(usize, end_payload, 8),
    };
}

pub const Iterator = struct {
    section: []align(8) const u8,
    next_start: usize = 0,
    shell_reader: ShellReader,
    idx: u32 = 0,
    total: u32,
    advanced_by: usize = 0,

    pub fn next(it: *Iterator) error{Overflow}!?PackedUser {
        if (it.idx == it.total) return null;
        const entry = try fromBytes(@alignCast(it.section[it.next_start..]));
        it.idx += 1;
        it.next_start += entry.end;
        it.advanced_by = entry.end;
        return entry.user;
    }

    pub fn rollback(it: *Iterator) void {
        assert(it.advanced_by > 0);
        it.idx -= 1;
        it.next_start -= it.advanced_by;
        it.advanced_by = 0;
    }
};

pub fn iterator(section: []align(8) const u8, total: u32, shell_reader: ShellReader) Iterator {
    return Iterator{
        .section = section,
        .total = total,
        .shell_reader = shell_reader,
    };
}

// packTo packs the User record and copies it to the given arraylist.
pub fn packTo(
    arr: *ArrayListAligned(u8, 8),
    user: User,
    additional_gids_offset: u64,
    idxFn: StringHashMap(u8),
) error{ InvalidRecord, OutOfMemory }!void {
    std.debug.assert(arr.items.len & 7 == 0);
    // function arguments are consts. We need to mutate the underlying
    // slice, so passing it via pointer instead.
    const home_len = try validate.downCast(fieldInfo(Inner, .home_len).type, user.home.len - 1);
    const name_len = try validate.downCast(fieldInfo(Inner, .name_len).type, user.name.len - 1);
    const shell_len = try validate.downCast(fieldInfo(Inner, .shell_len_or_idx).type, user.shell.len - 1);
    const gecos_len = try validate.downCast(fieldInfo(Inner, .gecos_len).type, user.gecos.len);

    try validate.utf8(user.home);
    try validate.utf8(user.name);
    try validate.utf8(user.shell);
    try validate.utf8(user.gecos);

    const inner = Inner{
        .uid = user.uid,
        .gid = user.gid,
        .shell_here = idxFn.get(user.shell) == null,
        .shell_len_or_idx = idxFn.get(user.shell) orelse shell_len,
        .home_len = home_len,
        .name_is_a_suffix = mem.endsWith(u8, user.home, user.name),
        .name_len = name_len,
        .gecos_len = gecos_len,
    };
    try arr.*.appendSlice(mem.asBytes(&inner)[0 .. @bitSizeOf(Inner) / 8]);
    try arr.*.appendSlice(user.home);

    if (!inner.name_is_a_suffix)
        try arr.*.appendSlice(user.name);
    try arr.*.appendSlice(user.gecos);
    if (inner.shell_here)
        try arr.*.appendSlice(user.shell);
    try compress.appendUvarint(arr, additional_gids_offset);
}

pub fn uid(self: PackedUser) u32 {
    return self.inner.uid;
}

pub fn gid(self: PackedUser) u32 {
    return self.inner.gid;
}

pub fn additionalGidsOffset(self: PackedUser) u64 {
    return self.additional_gids_offset;
}

pub fn home(self: PackedUser) []const u8 {
    return self.var_payload[0..self.inner.homeLen()];
}

pub fn name(self: PackedUser) []const u8 {
    const name_pos = self.inner.nameStart();
    const name_len = self.inner.nameLen();
    return self.var_payload[name_pos .. name_pos + name_len];
}

pub fn gecos(self: PackedUser) []const u8 {
    const gecos_pos = self.inner.gecosStart();
    const gecos_len = self.inner.gecosLen();
    return self.var_payload[gecos_pos .. gecos_pos + gecos_len];
}

pub fn shell(self: PackedUser, shell_reader: ShellReader) []const u8 {
    if (self.inner.shell_here) {
        const shell_pos = self.inner.maybeShellStart();
        const shell_len = self.inner.shellLen();
        return self.var_payload[shell_pos .. shell_pos + shell_len];
    }
    return shell_reader.get(self.inner.shell_len_or_idx);
}

// returns a User representation of the PackedUser.
pub fn toUser(self: *const PackedUser, shell_reader: ShellReader) User {
    return User{
        .uid = self.uid(),
        .gid = self.gid(),
        .name = self.name(),
        .gecos = self.gecos(),
        .home = self.home(),
        .shell = self.shell(shell_reader),
    };
}

pub const max_home_len = math.maxInt(fieldInfo(Inner, .home_len).type) + 1;
pub const max_name_len = math.maxInt(fieldInfo(Inner, .name_len).type) + 1;
pub const max_gecos_len = math.maxInt(fieldInfo(Inner, .gecos_len).type);

pub const max_str_len =
    math.maxInt(fieldInfo(Inner, .shell_len_or_idx).type) + 1 +
    max_home_len +
    max_name_len +
    max_gecos_len;

const testing = std.testing;

fn testShellIndex(allocator: Allocator) StringHashMap(u8) {
    var result = StringHashMap(u8).init(allocator);
    result.put("/bin/bash", 0) catch unreachable;
    result.put("/bin/zsh", 1) catch unreachable;
    return result;
}

const test_shell_reader_index: [3]u16 align(8) = .{ 0, 9, 17 };
const test_shell_reader = ShellReader{
    .blob = "/bin/bash/bin/zsh",
    .index = &test_shell_reader_index,
};

test "PackedUser pack max_user" {
    var arr = ArrayListAligned(u8, 8).init(testing.allocator);
    defer arr.deinit();

    var idx_noop = StringHashMap(u8).init(testing.allocator);
    defer idx_noop.deinit();

    try packTo(&arr, User.max_user, 0, idx_noop);
}

test "PackedUser construct section" {
    var buf = ArrayListAligned(u8, 8).init(testing.allocator);
    defer buf.deinit();

    const users = [_]User{ User{
        .uid = 1000,
        .gid = 1000,
        .name = "vidmantas",
        .gecos = "Vidmantas Kaminskas",
        .home = "/home/vidmantas",
        .shell = "/bin/bash",
    }, User{
        .uid = 1001,
        .gid = 1001,
        .name = "svc-foo",
        .gecos = "Service Account",
        .home = "/home/service1",
        .shell = "/usr/bin/nologin",
    }, User.max_user, User{
        .uid = 1002,
        .gid = 1002,
        .name = "svc-bar",
        .gecos = "",
        .home = "/",
        .shell = "/bin/zsh",
    } };
    var shellIndex = testShellIndex(testing.allocator);
    const additional_gids = math.maxInt(u64);
    defer shellIndex.deinit();
    for (users) |user| {
        try PackedUser.packTo(&buf, user, additional_gids, shellIndex);
        try buf.appendNTimes(0, mem.alignForward(usize, buf.items.len, 8) - buf.items.len);
    }

    var i: u29 = 0;
    var it1 = PackedUser.iterator(buf.items, users.len, test_shell_reader);
    while (try it1.next()) |user| : (i += 1) {
        try testing.expectEqual(users[i].uid, user.uid());
        try testing.expectEqual(users[i].gid, user.gid());
        try testing.expectEqual(user.additionalGidsOffset(), additional_gids);
        try testing.expectEqualStrings(users[i].name, user.name());
        try testing.expectEqualStrings(users[i].gecos, user.gecos());
        try testing.expectEqualStrings(users[i].home, user.home());
        try testing.expectEqualStrings(users[i].shell, user.shell(test_shell_reader));
    }
    try testing.expectEqual(users.len, i);
}