turbonss/src/user.zig

const std = @import("std");

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

const assert = std.debug.assert;
const mem = std.mem;
const Allocator = mem.Allocator;
const ArrayList = std.ArrayList;
// Idx2ShellProto is a function prototype that, given a shell's index (in
// global shell section), will return a shell string. Matches ShellReader.get.
const Idx2ShellProto = fn (u6) []const u8;

// User is a convenient public struct for record construction and
// serialization. Iterator can help retrieve these records.
pub const User = struct {
    uid: u32,
    gid: u32,
    name: []const u8,
    gecos: []const u8,
    home: []const u8,
    shell: []const u8,
};

pub const PackedUserMut = packedUser(false);
pub const PackedUserConst = packedUser(true);

fn packedUser(immutable: bool) type {
    return struct {
        const Self = @This();

        const alignmentBits = 3;
        const shell2idxProto = fn ([]const u8) ?u6;

        const InnerSize = @divExact(@bitSizeOf(Inner), 8);
        const Inner = packed struct {
            uid: u32,
            gid: u32,
            additional_gids_offset: u29,
            shell_len_or_idx: u6,
            shell_here: bool,
            home_len: u6,
            name_is_a_suffix: bool,
            name_len: u5,
            gecos_len: u8,

            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;
            }

            // blobLength returns the length of the blob storing string values.
            fn blobLength(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;
            }
        };

        const Bytes = if (immutable) []const u8 else []u8;
        const InnerPointer = if (immutable) *const Inner else *Inner;

        // PackedUser does not allocate; it re-interprets the "bytes" blob field.
        // Both of those fields are pointers to "our representation" of that field.
        inner: InnerPointer,
        userdata: Bytes,

        pub const Entry = struct {
            user: Self,
            next: ?Bytes,
        };

        pub fn fromBytes(bytes: Bytes) Entry {
            const inner = mem.bytesAsValue(
                Inner,
                // Should use InnerSize instead of sizeOf, see
                // https://github.com/ziglang/zig/issues/10958
                bytes[0..@sizeOf(Inner)],
            );

            const startBlob = InnerSize;
            const endBlob = startBlob + inner.blobLength();

            const nextStart = pad.roundUp(usize, alignmentBits, endBlob);
            var next: ?Bytes = null;
            if (nextStart < bytes.len) {
                next = bytes[nextStart..];
            }

            return Entry{
                .user = Self{
                    .inner = inner,
                    .userdata = bytes[startBlob..endBlob],
                },
                .next = next,
            };
        }

        pub const Iterator = struct {
            section: ?Bytes,
            shellIndex: Idx2ShellProto,

            pub fn next(it: *Iterator) ?Self {
                if (it.section) |section| {
                    const entry = Self.fromBytes(section);
                    it.section = entry.next;
                    return entry.user;
                }
                return null;
            }
        };

        pub fn iterator(section: Bytes, idxFn: Idx2ShellProto) Iterator {
            return Iterator{
                .section = section,
                .shellIndex = idxFn,
            };
        }

        // packTo packs the User record and copies it to the given byte slice. The
        // slice must have at least maxRecordSize() bytes available.
        // The slice is passed as a pointer, so it can be mutated.
        const packErr = InvalidRecord || Allocator.Error;
        pub fn packTo(
            arr: *ArrayList(u8),
            user: User,
            idxFn: shell2idxProto,
        ) packErr!void {
            // function arguments are consts. We need to mutate the underlying
            // slice, so passing it via pointer instead.
            const home_len = try validate.downCast(u6, user.home.len - 1);
            const name_len = try validate.downCast(u5, user.name.len - 1);
            const shell_len = try validate.downCast(u6, user.shell.len - 1);
            const gecos_len = try validate.downCast(u8, 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,
                .additional_gids_offset = std.math.maxInt(u29),
                .shell_here = idxFn(user.shell) == null,
                .shell_len_or_idx = idxFn(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,
            };
            const innerBytes = mem.asBytes(&inner);

            // innerBytes.len is longer than InnerSize. We want to copy
            // only the InnerSize-number of bytes.
            try arr.*.appendSlice(innerBytes[0..InnerSize]);
            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 pad.arrayList(arr, alignmentBits);
        }

        // maxSize is the maximum number of records a PackedUser can take
        // (struct + userdata).
        pub fn maxSize() usize {
            comptime {
                const unpadded = InnerSize +
                    std.math.maxInt(u6) + 1 + // home
                    std.math.maxInt(u5) + 1 + // name
                    std.math.maxInt(u6) + 1 + // shell
                    std.math.maxInt(u8); // gecos
                return pad.roundUp(u64, alignmentBits, unpadded);
            }
        }

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

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

        pub fn additionalGidsOffset(self: Self) u29 {
            return self.inner.additional_gids_offset;
        }

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

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

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

        pub fn shell(self: Self, idxFn: Idx2ShellProto) []const u8 {
            if (self.inner.shell_here) {
                const shell_pos = self.inner.maybeShellStart();
                const shell_len = self.inner.shellLen();
                return self.userdata[shell_pos .. shell_pos + shell_len];
            }
            return idxFn(self.inner.shell_len_or_idx);
        }

        pub fn setAdditionalGidsOffset(self: Self, new: u29) void {
            // TODO(motiejus) how to not declare function for const PackedUser at all?
            if (immutable) {
                @compileError("this function is available only for mutable PackedUsers");
            }
            self.inner.additional_gids_offset = new;
        }
    };
}

const testing = std.testing;

test "PackedUser internal and external alignment" {
    // External padding (alignmentBits) must be higher or equal to
    // the "internal" PackedUser alignment. By aligning PackedUser we are also
    // working around https://github.com/ziglang/zig/issues/10958 ; PackedUser
    // cannot be converted from/to [@bitSizeOf(PackedUser)/8]u8;
    // asBytes/bytesAsValue use @sizeOf, which is larger. Now we are putting no
    // more than 1, but it probably could be higher.
    try testing.expectEqual(8, @sizeOf(PackedUserConst.Inner) * 8 - @bitSizeOf(PackedUserConst.Inner));
}

fn testShellIndex(shell: []const u8) ?u6 {
    if (mem.eql(u8, shell, "/bin/bash")) {
        return 0;
    } else if (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();

    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 = std.math.maxInt(u32),
        .name = "Name" ** 8,
        .gecos = "Gecos" ** 51,
        .home = "Home" ** 16,
        .shell = "She.LllL" ** 8,
    }, User{
        .uid = 1002,
        .gid = 1002,
        .name = "svc-bar",
        .gecos = "",
        .home = "/",
        .shell = "/",
    } };
    for (users) |user| {
        try PackedUserConst.packTo(&buf, user, testShellIndex);
    }

    var i: u29 = 0;
    {
        var it = PackedUserConst.iterator(buf.items, testShell);
        while (it.next()) |user| : (i += 1) {
            try testing.expectEqual(users[i].uid, user.uid());
            try testing.expectEqual(users[i].gid, user.gid());
            try testing.expectEqual(
                @as(u29, std.math.maxInt(u29)),
                user.additionalGidsOffset(),
            );
            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(testShell));
        }
        try testing.expectEqual(users.len, i);
    }

    {
        var it = PackedUserMut.iterator(buf.items, testShell);
        i = 0;
        while (it.next()) |user| : (i += 1) {
            user.setAdditionalGidsOffset(i);
        }
        try testing.expectEqual(users.len, i);
    }

    {
        var it = PackedUserConst.iterator(buf.items, testShell);
        i = 0;
        while (it.next()) |user| : (i += 1) {
            try testing.expectEqual(users[i].gid, user.gid());
            try testing.expectEqual(i, user.additionalGidsOffset());
            try testing.expectEqualStrings(users[i].name, user.name());
        }
        try testing.expectEqual(users.len, i);
    }
}

test "PackedUser.maxSize()" {
    // TODO(motiejus) try using a slice that points to an array in stack.
    // As of writing, I am getting a stack smashing error.
    var buf = try ArrayList(u8).initCapacity(testing.allocator, PackedUserConst.maxSize());
    defer buf.deinit();

    const largeUser = User{
        .uid = std.math.maxInt(u32),
        .gid = std.math.maxInt(u32),
        .name = "Name" ** 8, // 32
        .gecos = "Gecos" ** 51, // 255
        .home = "Home" ** 16, // 64
        .shell = "She.LllL" ** 8, // 64
    };
    try PackedUserConst.packTo(&buf, largeUser, testShellIndex);
    try testing.expectEqual(PackedUserConst.maxSize(), buf.items.len);
}