refactor ShellReader, so it can be included

This commit is contained in:
Motiejus Jakštys 2022-03-10 19:19:50 +02:00 committed by Motiejus Jakštys
parent 1751f6c996
commit 3b2f135d5c
3 changed files with 41 additions and 40 deletions

View File

@ -248,9 +248,9 @@ pub fn usersSection(
var blob = try ArrayList(u8).initCapacity(allocator, 24 * corpus.users.len); var blob = try ArrayList(u8).initCapacity(allocator, 24 * corpus.users.len);
errdefer blob.deinit(); errdefer blob.deinit();
for (corpus.users) |user, i| { for (corpus.users) |user, i| {
const userOffset = try math.cast(u32, blob.items.len); const user_offset = try math.cast(u32, blob.items.len);
std.debug.assert(userOffset & 7 == 0); std.debug.assert(user_offset & 7 == 0);
idx2offset[i] = userOffset; idx2offset[i] = user_offset;
try userImport.PackedUser.packTo( try userImport.PackedUser.packTo(
&blob, &blob,
user, user,

View File

@ -10,7 +10,7 @@ const StringContext = std.hash_map.StringContext;
// maxShells is the maximum number of "popular" shells. // maxShells is the maximum number of "popular" shells.
pub const max_shells = 63; pub const max_shells = 63;
pub const max_shell_len = 64; pub const max_shell_len = 64;
const ShellAlignment = 2; // bits pub const shell_alignment_bits = 2; // bits
// ShellIndex is an index to the shell strings. As shell can be up to 64 bytes // ShellIndex is an index to the shell strings. As shell can be up to 64 bytes
// (1<<6), maximum number of shells is 63 (1<<6-1), the maximum location offset // (1<<6), maximum number of shells is 63 (1<<6-1), the maximum location offset
@ -18,29 +18,29 @@ const ShellAlignment = 2; // bits
// to 4 bytes. // to 4 bytes.
// The actual shell length is len+1: we don't allow empty shells, and the real // The actual shell length is len+1: we don't allow empty shells, and the real
// length of the shell is 1-64 bytes. // length of the shell is 1-64 bytes.
const ShellIndex = packed struct { pub const ShellIndex = packed struct {
offset: u10, offset: u10,
len: u6, len: u6,
}; };
// ShellReader interprets "Shell Index" and "Shell Blob" sections. // ShellReader interprets "Shell Index" and "Shell Blob" sections.
pub const ShellReader = struct { pub const ShellReader = struct {
sectionIndex: []const ShellIndex, section_index: []const ShellIndex,
sectionBlob: []const u8, section_blob: []const u8,
pub fn init(index: []const u8, blob: []const u8) ShellReader { pub fn init(index: []const u8, blob: []const u8) ShellReader {
return ShellReader{ return ShellReader{
.sectionIndex = std.mem.bytesAsSlice(ShellIndex, index), .section_index = std.mem.bytesAsSlice(ShellIndex, index),
.sectionBlob = blob, .section_blob = blob,
}; };
} }
// get returns a shell at the given index. // get returns a shell at the given index.
pub fn get(self: *const ShellReader, idx: u6) []const u8 { pub fn get(self: *const ShellReader, idx: u6) []const u8 {
const shellIndex = self.sectionIndex[idx]; const shell_index = self.section_index[idx];
const start = shellIndex.offset << 2; const start = shell_index.offset << 2;
const end = start + shellIndex.len + 1; const end = start + shell_index.len + 1;
return self.sectionBlob[start..end]; return self.section_blob[start..end];
} }
}; };
@ -72,31 +72,31 @@ pub const ShellWriter = struct {
.indices = StringHashMap(u6).init(allocator), .indices = StringHashMap(u6).init(allocator),
}; };
errdefer self.indices.deinit(); errdefer self.indices.deinit();
var fullOffset: u12 = 0; var full_offset: u12 = 0;
var idx: u6 = 0; var idx: u6 = 0;
while (idx < shells.len) : (idx += 1) { while (idx < shells.len) : (idx += 1) {
const len = try std.math.cast(u6, shells.get(idx).len); const len = try std.math.cast(u6, shells.get(idx).len);
try self.blob.appendSlice(shells.get(idx)); try self.blob.appendSlice(shells.get(idx));
const ourShell = self.blob.constSlice()[fullOffset .. fullOffset + len]; const our_shell = self.blob.constSlice()[full_offset .. full_offset + len];
try self.indices.put(ourShell, idx); try self.indices.put(our_shell, idx);
self.index.set(idx, ShellIndex{ self.index.set(idx, ShellIndex{
.offset = try std.math.cast(u10, fullOffset >> 2), .offset = try std.math.cast(u10, full_offset >> 2),
.len = len - 1, .len = len - 1,
}); });
fullOffset += len; full_offset += len;
const padding = pad.roundUpPadding(u12, ShellAlignment, fullOffset); const padding = pad.roundUpPadding(u12, shell_alignment_bits, full_offset);
fullOffset += padding; full_offset += padding;
try self.blob.appendNTimes(0, padding); try self.blob.appendNTimes(0, padding);
} }
return self; return self;
} }
pub fn sectionIndex(self: *const ShellSections) []const u8 { pub fn section_index(self: *const ShellSections) []const u8 {
return std.mem.sliceAsBytes(self.index.constSlice()); return std.mem.sliceAsBytes(self.index.constSlice());
} }
pub fn sectionBlob(self: *const ShellSections) []const u8 { pub fn section_blob(self: *const ShellSections) []const u8 {
return self.blob.constSlice(); return self.blob.constSlice();
} }
@ -198,17 +198,17 @@ test "basic shellpopcon" {
try testing.expectEqual(sections.getIndex(bash).?, 2); try testing.expectEqual(sections.getIndex(bash).?, 2);
try testing.expectEqual(sections.getIndex(nobody), null); try testing.expectEqual(sections.getIndex(nobody), null);
try testing.expectEqual( try testing.expectEqual(
sections.sectionBlob().len, sections.section_blob().len,
pad.roundUp(u12, 2, bash.len) + pad.roundUp(u12, 2, zsh.len) + pad.roundUp(u12, 2, long.len), pad.roundUp(u12, 2, bash.len) + pad.roundUp(u12, 2, zsh.len) + pad.roundUp(u12, 2, long.len),
); );
const shellReader = ShellReader.init( const shellReader = ShellReader.init(
sections.sectionIndex(), sections.section_index(),
sections.sectionBlob(), sections.section_blob(),
); );
try testing.expectEqualStrings(shellReader.get(0), long); try testing.expectEqualStrings(shellReader.get(0), long);
try testing.expectEqualStrings(shellReader.get(1), zsh); try testing.expectEqualStrings(shellReader.get(1), zsh);
try testing.expectEqualStrings(shellReader.get(2), bash); try testing.expectEqualStrings(shellReader.get(2), bash);
try testing.expectEqual(shellReader.sectionIndex.len, 3); try testing.expectEqual(shellReader.section_index.len, 3);
} }

View File

@ -3,6 +3,7 @@ const std = @import("std");
const pad = @import("padding.zig"); const pad = @import("padding.zig");
const validate = @import("validate.zig"); const validate = @import("validate.zig");
const compress = @import("compress.zig"); const compress = @import("compress.zig");
const shellImport = @import("shell.zig");
const InvalidRecord = validate.InvalidRecord; const InvalidRecord = validate.InvalidRecord;
const assert = std.debug.assert; const assert = std.debug.assert;
@ -180,7 +181,7 @@ pub const PackedUser = struct {
pub const Iterator = struct { pub const Iterator = struct {
section: ?[]const u8, section: ?[]const u8,
shellIndex: Idx2ShellProto, shell_reader: shellImport.ShellReader,
pub fn next(it: *Iterator) error{Overflow}!?Self { pub fn next(it: *Iterator) error{Overflow}!?Self {
if (it.section) |section| { if (it.section) |section| {
@ -192,8 +193,8 @@ pub const PackedUser = struct {
} }
}; };
pub fn iterator(section: []const u8, idxFn: Idx2ShellProto) Iterator { pub fn iterator(section: []const u8, shell_reader: shellImport.ShellReader) Iterator {
return Iterator{ .section = section, .shellIndex = idxFn }; return Iterator{ .section = section, .shell_reader = shell_reader };
} }
// packTo packs the User record and copies it to the given byte slice. // packTo packs the User record and copies it to the given byte slice.
@ -269,13 +270,13 @@ pub const PackedUser = struct {
return self.bytes[gecos_pos .. gecos_pos + gecos_len]; return self.bytes[gecos_pos .. gecos_pos + gecos_len];
} }
pub fn shell(self: Self, idxFn: Idx2ShellProto) []const u8 { pub fn shell(self: Self, shell_reader: shellImport.ShellReader) []const u8 {
if (self.inner.shell_here) { if (self.inner.shell_here) {
const shell_pos = self.inner.maybeShellStart(); const shell_pos = self.inner.maybeShellStart();
const shell_len = self.inner.shellLen(); const shell_len = self.inner.shellLen();
return self.bytes[shell_pos .. shell_pos + shell_len]; return self.bytes[shell_pos .. shell_pos + shell_len];
} }
return idxFn(self.inner.shell_len_or_idx); return shell_reader.get(self.inner.shell_len_or_idx);
} }
}; };
@ -295,13 +296,13 @@ fn testShellIndex(allocator: Allocator) StringHashMap(u6) {
return result; return result;
} }
fn testShell(index: u6) []const u8 { const test_shell_reader = shellImport.ShellReader{
return switch (index) { .section_blob = "/bin/bash000/bin/zsh",
0 => "/bin/bash", .section_index = &[_]shellImport.ShellIndex{
1 => "/bin/zsh", shellImport.ShellIndex{ .offset = 0, .len = 9 - 1 },
else => unreachable, shellImport.ShellIndex{ .offset = 12 >> 2, .len = 8 - 1 },
}; },
} };
test "construct PackedUser section" { test "construct PackedUser section" {
var buf = ArrayList(u8).init(testing.allocator); var buf = ArrayList(u8).init(testing.allocator);
@ -344,7 +345,7 @@ test "construct PackedUser section" {
} }
var i: u29 = 0; var i: u29 = 0;
var it1 = PackedUser.iterator(buf.items, testShell); var it1 = PackedUser.iterator(buf.items, test_shell_reader);
while (try it1.next()) |user| : (i += 1) { while (try it1.next()) |user| : (i += 1) {
try testing.expectEqual(users[i].uid, user.uid()); try testing.expectEqual(users[i].uid, user.uid());
try testing.expectEqual(users[i].gid, user.gid()); try testing.expectEqual(users[i].gid, user.gid());
@ -355,7 +356,7 @@ test "construct PackedUser section" {
try testing.expectEqualStrings(users[i].name, user.name()); try testing.expectEqualStrings(users[i].name, user.name());
try testing.expectEqualStrings(users[i].gecos, user.gecos()); try testing.expectEqualStrings(users[i].gecos, user.gecos());
try testing.expectEqualStrings(users[i].home, user.home()); try testing.expectEqualStrings(users[i].home, user.home());
try testing.expectEqualStrings(users[i].shell, user.shell(testShell)); try testing.expectEqualStrings(users[i].shell, user.shell(test_shell_reader));
} }
try testing.expectEqual(users.len, i); try testing.expectEqual(users.len, i);
} }