turbonss/src/user.zig

const std = @import("std");
const pad = @import("padding.zig");

const assert = std.debug.assert;
const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const cast = std.math.cast;

pub const PackedUser = packed struct {
    uid: u32,
    gid: u32,
    shell_here: bool,
    shell_len_or_idx: u6,
    home_len: u6,
    name_is_a_suffix: bool,
    name_len: u5,
    gecos_len: u10,
    padding: u3,

    // blobLength returns the length of the blob storing string values.
    pub fn blobLength(self: *const PackedUser) usize {
        var result: usize = self.realHomeLen();
        if (!self.name_is_a_suffix) {
            result += self.realNameLen();
        }
        result += self.realGecosLen();
        if (self.shell_here) {
            result += self.realShellLen();
        }

        return result;
    }

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

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

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

    pub fn realGecosLen(self: *const PackedUser) usize {
        return self.gecos_len;
    }
};
const PackedUserAlignmentBits = 3;

pub const User = struct {
    uid: u32,
    gid: u32,
    name: []const u8,
    gecos: []const u8,
    home: []const u8,
    shell: []const u8,
};

// UserWriter accepts a naive User struct and returns a PackedUser
pub const UserWriter = struct {
    // shellIndexFnType is a signature for a function that accepts a shell
    // string and returns it's index in the global shell section. Passing a
    // function makes tests easier, and removes the Shell dependency of this
    // module.
    const shellIndexFnType = fn ([]const u8) ?u6;

    appendTo: *ArrayList(u8),
    shellIndexFn: shellIndexFnType,

    pub fn init(
        appendTo: *ArrayList(u8),
        shellIndexFn: shellIndexFnType,
    ) UserWriter {
        return UserWriter{
            .appendTo = appendTo,
            .shellIndexFn = shellIndexFn,
        };
    }

    pub fn downCast(comptime T: type, n: u64) error{InvalidRecord}!T {
        return std.math.cast(T, n) catch |err| switch (err) {
            error.Overflow => {
                return error.InvalidRecord;
            },
        };
    }

    pub fn validateUtf8(s: []const u8) error{InvalidRecord}!void {
        if (!std.unicode.utf8ValidateSlice(s)) {
            return error.InvalidRecord;
        }
    }

    // FIXME(motiejus) record valiation should return a separate type. For User
    // case, it should be length-bound slices and utf8-codepoints instead of
    // strings.
    //
    // zig does not have error contexts
    // (https://github.com/ziglang/zig/issues/2647) and length-limited slices.
    // (It does have bounded_array, but that preallocates the maximum length,
    // which is not great for User records). So I am using those excuses to
    // do the validation here. I may move it once I learn the language better.
    const appendUserErr = error{InvalidRecord} || Allocator.Error;
    pub fn appendUser(self: *UserWriter, user: User) appendUserErr!void {
        const home_len = try downCast(u6, user.home.len - 1);
        const name_len = try downCast(u5, user.name.len - 1);
        const shell_len = try downCast(u6, user.shell.len - 1);
        const gecos_len = try downCast(u10, user.gecos.len);

        try validateUtf8(user.home);
        try validateUtf8(user.name);
        try validateUtf8(user.shell);
        try validateUtf8(user.gecos);

        var puser = PackedUser{
            .uid = user.uid,
            .gid = user.gid,
            .shell_here = self.shellIndexFn(user.shell) == null,
            .shell_len_or_idx = self.shellIndexFn(user.shell) orelse shell_len,
            .home_len = home_len,
            .name_is_a_suffix = std.mem.endsWith(u8, user.home, user.name),
            .name_len = name_len,
            .gecos_len = gecos_len,
            .padding = 0,
        };

        try self.appendTo.appendSlice(std.mem.asBytes(&puser));
        try self.appendTo.appendSlice(user.home);
        if (!puser.name_is_a_suffix) {
            try self.appendTo.appendSlice(user.name);
        }
        try self.appendTo.appendSlice(user.gecos);
        if (puser.shell_here) {
            try self.appendTo.appendSlice(user.shell);
        }

        try self.appendTo.appendNTimes(0, pad.roundUpPadding(
            u64,
            PackedUserAlignmentBits,
            self.appendTo.items.len,
        ));
    }
};

pub const UserReader = struct {
    const shellIndexProto = fn (u6) []const u8;

    section: []const u8,
    shellIndex: shellIndexProto,

    pub const PackedEntry = struct {
        packed_user: *PackedUser,
        section: []const u8,
    };

    pub fn init(section: []u8, shellIndex: shellIndexProto) UserReader {
        return UserReader{
            .section = section,
            .shellIndex = shellIndex,
        };
    }

    pub const Entry = struct {
        user: User,
        nextOffset: usize,
    };

    // atOffset returns a ?User in a given offset of the User section. Also,
    // the offset to the next user.
    pub fn atOffset(self: *UserReader, index: usize) ?Entry {
        if (index == self.section.len) return null;
        assert(index < self.section.len);
        const endUser = index + @sizeOf(PackedUser);
        var u = std.mem.bytesAsValue(
            PackedUser,
            self.section[index..endUser][0..@sizeOf(PackedUser)],
        );
        const startBlob = endUser;
        const endBlob = startBlob + u.blobLength();
        const section = self.section[startBlob..endBlob];

        const home = section[0..u.realHomeLen()];
        var name: []const u8 = undefined;
        var pos: usize = undefined;

        if (u.name_is_a_suffix) {
            const name_start = u.realHomeLen() - u.realNameLen();
            name = section[name_start..u.realHomeLen()];
            pos = u.realHomeLen();
        } else {
            const name_start = u.realHomeLen();
            name = section[name_start .. name_start + u.realNameLen()];
            pos = name_start + u.realNameLen();
        }
        const gecos = section[pos .. pos + u.realGecosLen()];
        pos += u.realGecosLen();

        var shell: []const u8 = undefined;
        if (u.shell_here) {
            shell = section[pos .. pos + u.realShellLen()];
        } else {
            shell = self.shellIndex(u.shell_len_or_idx);
        }

        return Entry{
            .user = User{
                .uid = u.uid,
                .gid = u.gid,
                .name = name,
                .gecos = gecos,
                .home = home,
                .shell = shell,
            },
            .nextOffset = pad.roundUp(usize, PackedUserAlignmentBits, endBlob),
        };
    }

    pub const Iterator = struct {
        ur: *UserReader,
        offset: usize = 0,

        pub fn next(it: *Iterator) ?User {
            if (it.ur.atOffset(it.offset)) |result| {
                it.offset = result.nextOffset;
                return result.user;
            }
            return null;
        }
    };

    pub fn iterator(self: *UserReader) Iterator {
        return Iterator{
            .ur = self,
            .offset = 0,
        };
    }
};

const testing = std.testing;

test "PackedUser internal and external alignment" {
    try testing.expectEqual(@bitSizeOf(PackedUser), @sizeOf(PackedUser) * 8);
}

fn testShellIndex(shell: []const u8) ?u6 {
    if (std.mem.eql(u8, shell, "/bin/bash")) {
        return 0;
    } else if (std.mem.eql(u8, shell, "/bin/zsh")) {
        return 1;
    }
    return null;
}

fn testShell(index: u6) []const u8 {
    return switch (index) {
        0 => "/bin/bash",
        1 => "/bin/zsh",
        else => unreachable,
    };
}

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

    var writer = UserWriter.init(&buf, testShellIndex);
    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{
        .uid = 0,
        .gid = 4294967295,
        .name = "n" ** 32,
        .gecos = "g" ** 1023,
        .home = "h" ** 64,
        .shell = "s" ** 64,
    } };
    for (users) |user| {
        try writer.appendUser(user);
    }

    var rd = UserReader.init(buf.items, testShell);

    var it = rd.iterator();
    var i: u32 = 0;
    while (it.next()) |user| : (i += 1) {
        try testing.expectEqual(users[i].uid, user.uid);
        try testing.expectEqual(users[i].gid, user.gid);
        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);
    }
}