From 44107bac51269964b13ab21a75ffa84900670854 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Fri, 18 Feb 2022 07:42:43 +0200
Subject: [PATCH] adding simple header validation

---
 src/header.zig | 66 ++++++++++++++++++++++++++++++++++++++++++++++----
 src/shell.zig  |  8 +++---
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/src/header.zig b/src/header.zig
index 8ab6785..bfff6c7 100644
--- a/src/header.zig
+++ b/src/header.zig
@@ -1,4 +1,5 @@
 const std = @import("std");
+const shells = @import("shell.zig");
 
 const HeaderSize = @sizeOf(Header);
 const Magic = [4]u8{ 0xf0, 0x9f, 0xa4, 0xb7 };
@@ -10,6 +11,7 @@ const HeaderError = error{
     InvalidVersion,
     InvalidBom,
     TooManyShells,
+    InvalidOffset,
 };
 
 const Header = packed struct {
@@ -29,7 +31,37 @@ const Header = packed struct {
     offset_additional_gids: u32,
 
     pub fn init(blob: [HeaderSize]u8) HeaderError!Header {
-        return @bitCast(Header, blob);
+        const self = @bitCast(Header, blob);
+        for (Magic) |item, index| {
+            if (self.magic[index] != item) return error.InvalidMagic;
+        }
+        if (self.version != 0) {
+            return error.InvalidVersion;
+        }
+        if (self.bom != Bom) {
+            return error.InvalidBom;
+        }
+        if (self.num_shells > shells.MaxShells) {
+            return error.TooManyShells;
+        }
+
+        const offsets = [_]u32{
+            self.offset_cmph_uid2user,
+            self.offset_cmph_groupname2group,
+            self.offset_cmph_username2user,
+            self.offset_idx,
+            self.offset_groups,
+            self.offset_users,
+            self.offset_groupmembers,
+            self.offset_additional_gids,
+        };
+        for (offsets) |offset| {
+            if (offset & @intCast(u32, 63) != 0) {
+                return error.InvalidOffset;
+            }
+        }
+
+        return self;
     }
 
     pub fn asArray(self: *const Header) [HeaderSize]u8 {
@@ -44,7 +76,7 @@ test "header size is byte-aligned" {
 }
 
 test "header pack and unpack" {
-    const header = Header{
+    const goodHeader = Header{
         .magic = Magic,
         .version = Version,
         .bom = Bom,
@@ -61,7 +93,31 @@ test "header pack and unpack" {
         .offset_additional_gids = 0,
     };
 
-    const blob = header.asArray();
-    const header2 = try Header.init(blob);
-    try testing.expectEqual(header, header2);
+    const blob = goodHeader.asArray();
+    const gotHeader = try Header.init(blob);
+    try testing.expectEqual(goodHeader, gotHeader);
+
+    {
+        var header = goodHeader;
+        header.magic[0] = 0;
+        try testing.expectError(error.InvalidMagic, Header.init(header.asArray()));
+    }
+
+    {
+        var header = goodHeader;
+        header.bom = 0x3412;
+        try testing.expectError(error.InvalidBom, Header.init(header.asArray()));
+    }
+
+    {
+        var header = goodHeader;
+        header.num_shells = shells.MaxShells + 1;
+        try testing.expectError(error.TooManyShells, Header.init(header.asArray()));
+    }
+
+    {
+        var header = goodHeader;
+        header.offset_cmph_uid2user = 128;
+        try testing.expectError(error.InvalidOffset, Header.init(header.asArray()));
+    }
 }
diff --git a/src/shell.zig b/src/shell.zig
index 366735b..cd0a279 100644
--- a/src/shell.zig
+++ b/src/shell.zig
@@ -7,8 +7,8 @@ const BoundedArray = std.BoundedArray;
 const StringContext = std.hash_map.StringContext;
 
 // MaxShells is the maximum number of "popular" shells.
-const MaxShells = 63;
-const MaxShellLen = 64;
+pub const MaxShells = 63;
+pub const MaxShellLen = 64;
 
 // 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
@@ -22,7 +22,7 @@ const ShellIndex = struct {
 };
 
 // ShellReader interprets "Shell Index" and "Shell Blob" sections.
-const ShellReader = struct {
+pub const ShellReader = struct {
     sectionIndex: []const ShellIndex,
     sectionBlob: []const u8,
 
@@ -44,7 +44,7 @@ const ShellReader = struct {
 
 // ShellWriter is a shell popularity contest: collect shells and return the
 // popular ones, sorted by score. score := len(shell) * number_of_shells.
-const ShellWriter = struct {
+pub const ShellWriter = struct {
     counts: std.StringHashMap(u32),
     allocator: Allocator,
     const KV = struct { shell: []const u8, score: u32 };