[PATCH v2] pickles: Add mbr.pk for MBR partition table

[Mohammad-Reza Nabipoor]

2021-01-28  Mohammad-Reza Nabipoor  <m.nabipoor@yahoo.com>

        * pickles/mbr.pk: New pickle.
        * pickles/Makefile.am (dist_pickles_DATA): Add `mbr.pk`.
        * testsuite/poke.pickles/mbr-test.pk: New test.
        * testsuite/Makefile.am (EXTRA_DIST): Update.
---
 ChangeLog                          |   7 ++
 pickles/Makefile.am                |   2 +-
 pickles/mbr.pk                     | 104 ++++++++++++++++++++++++++
 testsuite/Makefile.am              |   1 +
 testsuite/poke.pickles/mbr-test.pk | 116 +++++++++++++++++++++++++++++
 5 files changed, 229 insertions(+), 1 deletion(-)
 create mode 100644 pickles/mbr.pk
 create mode 100644 testsuite/poke.pickles/mbr-test.pk

diff --git a/ChangeLog b/ChangeLog
index 87f22258..336fcc6a 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,10 @@
+2021-01-28  Mohammad-Reza Nabipoor  <m.nabipoor@yahoo.com>
+
+       * pickles/mbr.pk: New pickle.
+       * pickles/Makefile.am (dist_pickles_DATA): Add `mbr.pk`.
+       * testsuite/poke.pickles/mbr-test.pk: New test.
+       * testsuite/Makefile.am (EXTRA_DIST): Update.
+
 2021-01-28  Mohammad-Reza Nabipoor  <m.nabipoor@yahoo.com>
 
        * libpoke/pkl-gen.pks (struct_writer): Fix bit-offset of mapped
diff --git a/pickles/Makefile.am b/pickles/Makefile.am
index d5452e52..7dada3dd 100644
--- a/pickles/Makefile.am
+++ b/pickles/Makefile.am
@@ -2,4 +2,4 @@ picklesdir = $(pkgdatadir)/pickles
 dist_pickles_DATA = elf.pk ctf.pk leb128.pk bpf.pk btf.pk bmp.pk \
                     color.pk rgb24.pk id3v1.pk \
                     dwarf.pk dwarf-common.pk dwarf-frame.pk dwarf-pubnames.pk \
-                    dwarf-types.pk time.pk argp.pk pktest.pk
+                    dwarf-types.pk time.pk argp.pk pktest.pk mbr.pk
diff --git a/pickles/mbr.pk b/pickles/mbr.pk
new file mode 100644
index 00000000..4bf5172b
--- /dev/null
+++ b/pickles/mbr.pk
@@ -0,0 +1,104 @@
+/* mbr.pk - MBR (Master Boot Record) partition table.  */
+
+/* Copyright (C) 2021 The poke authors */
+
+/* This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+/* set_endian(ENDIAN_LITTLE); */
+
+type MBR_CHS =
+  struct
+  {
+    type CySec = /* FIXME use unnamed type for `cs` */
+      struct byte
+      {
+        uint<2> cylinder_hi;
+        uint<6> sector;
+      };
+    byte head;
+    CySec cs;
+    byte cylinder_lo;
+
+    method get_cylinder = uint<10>:
+      {
+        return cs.cylinder_hi ::: cylinder_lo;
+      }
+    method set_cylinder = (uint<10> c) void:
+      {
+        cs.cylinder_hi = c .>> 8;
+        cylinder_lo = c;
+      }
+    method _print = void:
+      {
+        printf ("#<head=%u8d,sector=%u6d,cylinder=%u10d>",
+                head, cs.sector, get_cylinder);
+      }
+  };
+
+/* MBR Partition Table Entry (PTE)
+ *
+ *  Offset  Size (bytes)  Description
+ *
+ *  0x00    1             Drive attributes (bit 7 set = active or bootable)
+ *  0x01    3             CHS Address of partition start
+ *  0x04    1             Partition type
+ *  0x05    3             CHS address of last partition sector
+ *  0x08    4             LBA of partition start
+ *  0x0C    4             Number of sectors in partition
+ *
+ * ref: https://wiki.osdev.org/MBR_(x86)
+ */
+type MBR_PTE =
+  struct
+  {
+    struct byte
+    {
+      uint<1> active;  /* bootable */
+      uint<7>;
+
+      method _print = void:
+        {
+          printf "#<active=%u1d>", active;
+        }
+    } attr;
+    MBR_CHS start_chs;
+    byte part_type;  /* partition type*/
+    MBR_CHS end_chs;
+    uint32 lba;
+    uint32 sector_count;
+  };
+
+/*
+ *  Offset  Size (bytes)  Description
+ *
+ *  0x000   440           MBR Bootstrap (flat binary executable code)
+ *  0x1B8   4             Optional "Unique Disk ID / Signature"
+ *  0x1BC   2             Optional, reserved 0x0000
+ *  0x1BE   16            First partition table entry
+ *  0x1CE   16            Second partition table entry
+ *  0x1DE   16            Third partition table entry
+ *  0x1EE   16            Fourth partition table entry
+ *  0x1FE   2             (0x55, 0xAA) "Valid bootsector" signature bytes
+ *
+ * ref: https://wiki.osdev.org/MBR_(x86)
+ */
+type MBR = struct
+  {
+    byte[440#B] bootstrap;
+    uint32 signature;
+    byte[2] reserved;
+    MBR_PTE[4] pte;
+    uint16 magic = 0xaa55UH;
+  };
diff --git a/testsuite/Makefile.am b/testsuite/Makefile.am
index 15c490ce..3139b33e 100644
--- a/testsuite/Makefile.am
+++ b/testsuite/Makefile.am
@@ -481,6 +481,7 @@ EXTRA_DIST = \
   poke.pickles/pickles.exp \
   poke.pickles/argp-test.pk \
   poke.pickles/color-test.pk \
+  poke.pickles/mbr-test.pk \
   poke.pickles/id3v1-test.pk \
   poke.pickles/rgb24-test.pk \
   poke.pkl/pkl.exp \
diff --git a/testsuite/poke.pickles/mbr-test.pk 
b/testsuite/poke.pickles/mbr-test.pk
new file mode 100644
index 00000000..962fcc77
--- /dev/null
+++ b/testsuite/poke.pickles/mbr-test.pk
@@ -0,0 +1,116 @@
+/* mbr-test.pk - Tests for the mbr pickle.  */
+
+/* Copyright (C) 2021 The poke authors */
+
+/* This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+load pktest;
+load mbr;
+
+set_endian(ENDIAN_LITTLE);
+
+var data = open ("*data*");
+
+var tests = [
+  PkTest {
+    name = "magic number",
+    func = lambda (string name) void:
+      {
+        byte[512] @ data : 0#B = byte[512] ();
+
+        try
+          {
+            var m = MBR @ data : 0#B;
+
+            assert(0, "unreachable reached!");
+          }
+        catch if E_constraint
+          {
+            assert (1, "the magic number cannot be zero!");
+          }
+
+        /* Fix the magic number */
+        byte[2] @ data : 510#B = [0x55UB, 0xaaUB];
+
+        var m = MBR @ data : 0#B;
+      },
+  },
+  PkTest {
+    name = "CHS",
+    func = lambda (string name) void:
+      {
+        byte[510] @ data : 0#B = byte[510] ();
+        byte[2] @ data : 510#B = [0x55UB, 0xaaUB];
+
+        var m = MBR @ data : 0#B;
+        var chs = m.pte[0].start_chs;
+
+        assert(chs.head == 0);
+        assert(chs.cs.cylinder_hi == 0);
+        assert(chs.cs.sector == 0);
+        assert(chs.cylinder_lo == 0);
+        assert(chs.get_cylinder == 0);
+        assert((byte[chs'size] @ data : chs'offset) == [0UB, 0UB, 0UB]);
+
+        chs.set_cylinder(0x3ff);
+        assert(chs.head == 0);
+        assert(chs.cs.cylinder_hi == 3);
+        assert(chs.cs.sector == 0);
+        assert(chs.cylinder_lo == 0xff);
+        assert(chs.get_cylinder == 0x3ff);
+        assert((byte[chs'size] @ data : chs'offset) == [0UB, 0xc0UB, 0xffUB]);
+
+        chs.cs.cylinder_hi = 1;
+        assert(chs.head == 0);
+        assert(chs.cs.cylinder_hi == 1);
+        assert(chs.cs.sector == 0);
+        assert(chs.cylinder_lo == 0xff);
+        assert(chs.get_cylinder == 0x1ff);
+        assert((byte[chs'size] @ data : chs'offset) == [0UB, 0x40UB, 0xffUB]);
+
+        chs.head = 0xa5;
+        assert(chs.head == 0xa5);
+        assert(chs.cs.cylinder_hi == 1);
+        assert(chs.cs.sector == 0);
+        assert(chs.cylinder_lo == 0xff);
+        assert(chs.get_cylinder == 0x1ff);
+        assert((byte[chs'size] @ data : chs'offset) ==
+               [0xa5UB, 0x40UB, 0xffUB]);
+
+        chs.cs.sector = 0xff;  /* Only the 6 LSBs are relevant */
+        assert(chs.head == 0xa5);
+        assert(chs.cs.cylinder_hi == 1);
+        assert(chs.cs.sector == 0x3f);
+        assert(chs.cylinder_lo == 0xff);
+        assert(chs.get_cylinder == 0x1ff);
+        assert((byte[chs'size] @ data : chs'offset) ==
+               [0xa5UB, 0x7fUB, 0xffUB]);
+
+        chs.cylinder_lo = 0x81;
+        assert(chs.head == 0xa5);
+        assert(chs.cs.cylinder_hi == 1);
+        assert(chs.cs.sector == 0x3f);
+        assert(chs.cylinder_lo == 0x81);
+        assert(chs.get_cylinder == 0x181);
+        assert((byte[chs'size] @ data : chs'offset) ==
+               [0xa5UB, 0x7fUB, 0x81UB]);
+      },
+  },
+];
+
+var ec = pktest_run (tests) ? 0 : 1;
+
+close(data);
+exit (ec);
-- 
2.30.0