1
Fork 0

refactor ShellReader, so it can be included

main
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);
errdefer blob.deinit();
for (corpus.users) |user, i| {
const userOffset = try math.cast(u32, blob.items.len);
std.debug.assert(userOffset & 7 == 0);
idx2offset[i] = userOffset;
const user_offset = try math.cast(u32, blob.items.len);
std.debug.assert(user_offset & 7 == 0);
idx2offset[i] = user_offset;
try userImport.PackedUser.packTo(
&blob,
user,

View File

@ -10,7 +10,7 @@ const StringContext = std.hash_map.StringContext;
// maxShells is the maximum number of "popular" shells.
pub const max_shells = 63;
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
// (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.
// The actual shell length is len+1: we don't allow empty shells, and the real
// length of the shell is 1-64 bytes.
const ShellIndex = packed struct {
pub const ShellIndex = packed struct {
offset: u10,
len: u6,
};
// ShellReader interprets "Shell Index" and "Shell Blob" sections.
pub const ShellReader = struct {
sectionIndex: []const ShellIndex,
sectionBlob: []const u8,
section_index: []const ShellIndex,
section_blob: []const u8,
pub fn init(index: []const u8, blob: []const u8) ShellReader {
return ShellReader{
.sectionIndex = std.mem.bytesAsSlice(ShellIndex, index),
.sectionBlob = blob,
.section_index = std.mem.bytesAsSlice(ShellIndex, index),
.section_blob = blob,
};
}
// get returns a shell at the given index.
pub fn get(self: *const ShellReader, idx: u6) []const u8 {
const shellIndex = self.sectionIndex[idx];
const start = shellIndex.offset << 2;
const end = start + shellIndex.len + 1;
return self.sectionBlob[start..end];
const shell_index = self.section_index[idx];
const start = shell_index.offset << 2;
const end = start + shell_index.len + 1;
return self.section_blob[start..end];
}
};
@ -72,31 +72,31 @@ pub const ShellWriter = struct {
.indices = StringHashMap(u6).init(allocator),
};
errdefer self.indices.deinit();
var fullOffset: u12 = 0;
var full_offset: u12 = 0;
var idx: u6 = 0;
while (idx < shells.len) : (idx += 1) {
const len = try std.math.cast(u6, shells.get(idx).len);
try self.blob.appendSlice(shells.get(idx));
const ourShell = self.blob.constSlice()[fullOffset .. fullOffset + len];
try self.indices.put(ourShell, idx);
const our_shell = self.blob.constSlice()[full_offset .. full_offset + len];
try self.indices.put(our_shell, idx);
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,
});
fullOffset += len;
const padding = pad.roundUpPadding(u12, ShellAlignment, fullOffset);
fullOffset += padding;
full_offset += len;
const padding = pad.roundUpPadding(u12, shell_alignment_bits, full_offset);
full_offset += padding;
try self.blob.appendNTimes(0, padding);
}
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());
}
pub fn sectionBlob(self: *const ShellSections) []const u8 {
pub fn section_blob(self: *const ShellSections) []const u8 {
return self.blob.constSlice();
}
@ -198,17 +198,17 @@ test "basic shellpopcon" {
try testing.expectEqual(sections.getIndex(bash).?, 2);
try testing.expectEqual(sections.getIndex(nobody), null);
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),
);
const shellReader = ShellReader.init(
sections.sectionIndex(),
sections.sectionBlob(),
sections.section_index(),
sections.section_blob(),
);
try testing.expectEqualStrings(shellReader.get(0), long);
try testing.expectEqualStrings(shellReader.get(1), zsh);
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 validate = @import("validate.zig");
const compress = @import("compress.zig");
const shellImport = @import("shell.zig");
const InvalidRecord = validate.InvalidRecord;
const assert = std.debug.assert;
@ -180,7 +181,7 @@ pub const PackedUser = struct {
pub const Iterator = struct {
section: ?[]const u8,
shellIndex: Idx2ShellProto,
shell_reader: shellImport.ShellReader,
pub fn next(it: *Iterator) error{Overflow}!?Self {
if (it.section) |section| {
@ -192,8 +193,8 @@ pub const PackedUser = struct {
}
};
pub fn iterator(section: []const u8, idxFn: Idx2ShellProto) Iterator {
return Iterator{ .section = section, .shellIndex = idxFn };
pub fn iterator(section: []const u8, shell_reader: shellImport.ShellReader) Iterator {
return Iterator{ .section = section, .shell_reader = shell_reader };
}
// 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];
}
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) {
const shell_pos = self.inner.maybeShellStart();
const shell_len = self.inner.shellLen();
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;
}
fn testShell(index: u6) []const u8 {
return switch (index) {
0 => "/bin/bash",
1 => "/bin/zsh",
else => unreachable,
};
}
const test_shell_reader = shellImport.ShellReader{
.section_blob = "/bin/bash000/bin/zsh",
.section_index = &[_]shellImport.ShellIndex{
shellImport.ShellIndex{ .offset = 0, .len = 9 - 1 },
shellImport.ShellIndex{ .offset = 12 >> 2, .len = 8 - 1 },
},
};
test "construct PackedUser section" {
var buf = ArrayList(u8).init(testing.allocator);
@ -344,7 +345,7 @@ test "construct PackedUser section" {
}
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) {
try testing.expectEqual(users[i].uid, user.uid());
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].gecos, user.gecos());
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);
}