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 math = std.math;
const mem = std.mem;

// TODO(motiejus) move to the struct where it's used.
const shellIndexProto = fn (u6) []const u8;
const PackedUserAlignmentBits = 3;

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

pub const PackedUser = struct {
    // TODO: use @bitSizeOf(Inner)/8:
    //const InnerSize = @divExact(@bitSizeOf(Inner), 8);
    const shellIndexFnType = fn ([]const u8) ?u6;

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

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

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

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

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

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

            return result;
        }

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

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

        pub fn maybeShellPos(self: *const Inner) usize {
            assert(self.shell_here);
            return self.gecosPos() + self.gecosLen();
        }
    };

    inner: Inner,
    userdata: []const u8,

    pub fn fromBytes(bytes: []const u8) PackedUser {
        const inner = std.mem.bytesAsValue(
            PackedUser,
            bytes[0..@sizeOf(Inner)][0..@sizeOf(Inner)],
        );

        const startBlob = InnerSize;
        const endBlob = startBlob + inner.blobLength();
        return PackedUser{
            .inner = inner,
            .userdata = bytes[startBlob..endBlob],
        };
    }

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

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

    // packTo packs the User record and copies it to the given byte slice. The
    // slice must have at least maxRecordSize() bytes available.
    pub fn packTo(buf1: *[]u8, user: User, shellIndexFn: shellIndexFnType) error{InvalidRecord}!void {
        var buf = buf1.*;
        const bufStart = @ptrToInt(buf.ptr);

        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(u8, user.gecos.len);

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

        const inner = Inner{
            .uid = user.uid,
            .gid = user.gid,
            .additional_gids_offset = std.math.maxInt(u29),
            .shell_here = shellIndexFn(user.shell) == null,
            .shell_len_or_idx = 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,
        };
        const innerBytes = mem.asBytes(&inner);
        buf.len += innerBytes.len +
            user.home.len +
            user.gecos.len;

        mem.copy(u8, buf[0..innerBytes.len], innerBytes);
        buf = buf[innerBytes.len..];
        mem.copy(u8, buf[0..user.home.len], user.home);
        buf = buf[user.home.len..];

        if (!inner.name_is_a_suffix) {
            buf.len += user.name.len;
            mem.copy(u8, buf[0..user.name.len], user.name);
            buf = buf[user.name.len..];
        }
        mem.copy(u8, buf, user.gecos);
        buf = buf[user.gecos.len..];
        if (inner.shell_here) {
            buf.len += user.shell.len;
            mem.copy(u8, buf[0..user.shell.len], user.shell);
            buf = buf[user.shell.len..];
        }

        const bufLen = @ptrToInt(buf.ptr) - bufStart;
        const padding = pad.roundUpPadding(u64, PackedUserAlignmentBits, bufLen);
        buf.len += padding;
        mem.set(u8, buf[0..padding], 0);
    }

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

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

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

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

    pub fn shell(self: *const PackedUser, shellIndex: shellIndexProto) []const u8 {
        if (self.inner.shell_here) {
            const shell_pos = self.inner.maybeShellPos();
            const shell_len = self.inner.shellLen();
            return self.userdata[shell_pos .. shell_pos + shell_len];
        }
        return shellIndex(self.inner.shell_len_or_idx);
    }
};

pub const UserReader = struct {
    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.homeLen()];
        var name: []const u8 = undefined;
        var pos: usize = undefined;

        if (u.name_is_a_suffix) {
            const name_start = u.homeLen() - u.nameLen();
            name = section[name_start..u.homeLen()];
            pos = u.homeLen();
        } else {
            const name_start = u.homeLen();
            name = section[name_start .. name_start + u.nameLen()];
            pos = name_start + u.nameLen();
        }
        const gecos = section[pos .. pos + u.gecosLen()];
        pos += u.gecosLen();

        var shell: []const u8 = undefined;
        if (u.shell_here) {
            shell = section[pos .. pos + u.shellLen()];
        } 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" {
    // External padding (PackedUserAlignmentBits) 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.expect(@sizeOf(PackedUser) * 8 - @bitSizeOf(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();

    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" ** 255,
        .home = "h" ** 64,
        .shell = "s" ** 64,
    } };
    for (users) |user| {
        try buf.ensureUnusedCapacity(PackedUser.maxSize());
        try PackedUser.packTo(&buf.items, user, testShellIndex);
    }

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