ocaml-mbr / lib / mbr.ml

(*
 * Copyright (C) 2013 Citrix Inc
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *)

open Result.Syntax

(* Error helpers *)
let err_lba_max x = Error (Fmt.str "sector count exceeds LBA max: %Ld" x)

let err_bad_signature s1 s2 =
  Error (Fmt.str "Invalid signature: %02x %02x <> 0x55 0xaa" s1 s2)

let err_too_small n = Error (Fmt.str "MBR too small: %d < 512" n)

module Reader = Bytesrw.Bytes.Reader
module Writer = Bytesrw.Bytes.Writer
module Slice = Bytesrw.Bytes.Slice

(* Read exactly n bytes from reader *)
let read_exactly reader n =
  if n = 0 then Ok Bytes.empty
  else
    let buf = Bytes.create n in
    let rec loop pos remaining =
      if remaining = 0 then Ok buf
      else
        match Reader.read reader with
        | slice when Slice.is_eod slice -> Error "unexpected end of data"
        | slice ->
            let str = Slice.to_string slice in
            let available = String.length str in
            let to_copy = min available remaining in
            Bytes.blit_string str 0 buf pos to_copy;
            if to_copy < available then begin
              let leftover = String.sub str to_copy (available - to_copy) in
              Reader.push_back reader (Slice.of_string leftover)
            end;
            loop (pos + to_copy) (remaining - to_copy)
    in
    loop 0 n

module Geometry = struct
  type t = { cylinders : int; heads : int; sectors : int }

  let kib = 1024L
  let mib = Int64.mul kib 1024L

  let of_lba_size x =
    let sectors = 63 in
    let* heads =
      if x < Int64.(mul 504L mib) then Ok 16
      else if x < Int64.(mul 1008L mib) then Ok 64
      else if x < Int64.(mul 4032L mib) then Ok 128
      else if x < Int64.(add (mul 8032L mib) (mul 512L kib)) then Ok 255
      else err_lba_max x
    in
    let cylinders =
      Int64.(to_int (div (div x (of_int sectors)) (of_int heads)))
    in
    Ok { cylinders; heads; sectors }

  let to_chs g x =
    let open Int64 in
    let cylinders = to_int (div x (mul (of_int g.sectors) (of_int g.heads))) in
    let heads = to_int (rem (div x (of_int g.sectors)) (of_int g.heads)) in
    let sectors = to_int (succ (rem x (of_int g.sectors))) in
    { cylinders; heads; sectors }
end

module Partition = struct
  type t = {
    active : bool;
    first_absolute_sector_chs : Geometry.t;
    ty : int;
    last_absolute_sector_chs : Geometry.t;
    first_absolute_sector_lba : int32;
    sectors : int32;
  }

  let sector_start t =
    Int64.(logand (of_int32 t.first_absolute_sector_lba) 0xFFFF_FFFFL)

  let size_sectors t = Int64.(logand (of_int32 t.sectors) 0xFFFF_FFFFL)

  let make ?(active = false) ~partition_type:(ty : int)
      first_absolute_sector_lba sectors =
    let* () =
      if ty > 0 && ty < 256 then Ok ()
      else Error "Mbr.Partition.make: ty must be between 1 and 255"
    in
    let first_absolute_sector_chs =
      { Geometry.cylinders = 0; heads = 0; sectors = 0 }
    in
    let last_absolute_sector_chs = first_absolute_sector_chs in
    Ok
      {
        active;
        first_absolute_sector_chs;
        ty;
        last_absolute_sector_chs;
        first_absolute_sector_lba;
        sectors;
      }

  let make' ?active ~partition_type:(ty : int) sector_start size_sectors =
    if
      Int64.(
        logand sector_start 0xFFFF_FFFFL = sector_start
        && logand size_sectors 0xFFFF_FFFFL = size_sectors)
    then
      let sector_start = Int64.to_int32 sector_start in
      let size_sectors = Int64.to_int32 size_sectors in
      make ?active ~partition_type:ty sector_start size_sectors
    else Error "partition parameters do not fit in int32"

  let decode_chs s =
    let heads = Char.code s.[0] in
    let y = Char.code s.[1] in
    let z = Char.code s.[2] in
    { Geometry.cylinders = (y lsl 2) lor z; heads; sectors = y land 0x3F }

  let f_status = Wire.Field.v "status" Wire.uint8

  let f_first_chs =
    Wire.Field.v "first_chs" (Wire.byte_array ~size:(Wire.int 3))

  let f_ty = Wire.Field.v "type" Wire.uint8
  let f_last_chs = Wire.Field.v "last_chs" (Wire.byte_array ~size:(Wire.int 3))
  let f_lba_start = Wire.Field.v "lba_start" Wire.uint32
  let f_sectors = Wire.Field.v "sectors" Wire.uint32

  let codec =
    Wire.Codec.v "MbrPartition"
      (fun status first_chs ty last_chs lba sectors ->
        {
          active = status = 0x80;
          first_absolute_sector_chs = decode_chs first_chs;
          ty;
          last_absolute_sector_chs = decode_chs last_chs;
          first_absolute_sector_lba =
            Int32.of_int (Wire.Private.UInt32.to_int lba);
          sectors = Int32.of_int (Wire.Private.UInt32.to_int sectors);
        })
      Wire.Codec.
        [
          (f_status $ fun t -> if t.active then 0x80 else 0);
          (f_first_chs $ fun _ -> "\000\000\000");
          (f_ty $ fun t -> t.ty);
          (f_last_chs $ fun _ -> "\000\000\000");
          ( f_lba_start $ fun t ->
            Wire.Private.UInt32.of_int
              (Int32.to_int t.first_absolute_sector_lba) );
          ( f_sectors $ fun t ->
            Wire.Private.UInt32.of_int (Int32.to_int t.sectors) );
        ]
end

type t = {
  bootstrap_code : string;
  original_physical_drive : int;
  seconds : int;
  minutes : int;
  hours : int;
  disk_signature : int32;
  partitions : Partition.t list;
}

let pp ppf t =
  Fmt.pf ppf "@[<v>MBR(sig=%ld, %d partitions)@]" t.disk_signature
    (List.length t.partitions)

let partitions t = t.partitions

(* Security helper: check for int32 addition overflow *)
let int32_add_overflow a b =
  let sum = Int32.add a b in
  (* Overflow if: a > 0 && b > 0 && sum < 0, or a < 0 && b < 0 && sum > 0 *)
  (* For unsigned interpretation, check if sum < a when both are positive *)
  Int32.unsigned_compare sum a < 0

let validate_partitions partitions =
  let* () =
    if List.length partitions <= 4 then Ok () else Error "Too many partitions"
  in
  let num_active =
    List.fold_left
      (fun acc p -> if p.Partition.active then succ acc else acc)
      0 partitions
  in
  let* () =
    if num_active <= 1 then Ok ()
    else Error "More than one active/boot partitions is not advisable"
  in
  let* () =
    List.fold_left
      (fun r p ->
        let* () = r in
        if
          int32_add_overflow p.Partition.first_absolute_sector_lba
            p.Partition.sectors
        then
          Error
            (Fmt.str "Partition start %lu + size %lu overflows int32"
               p.Partition.first_absolute_sector_lba p.Partition.sectors)
        else Ok ())
      (Ok ()) partitions
  in
  let sorted =
    List.sort
      (fun p1 p2 ->
        Int32.unsigned_compare p1.Partition.first_absolute_sector_lba
          p2.Partition.first_absolute_sector_lba)
      partitions
  in
  let* (_ : int32) =
    List.fold_left
      (fun r p ->
        let* offset = r in
        if
          Int32.unsigned_compare offset p.Partition.first_absolute_sector_lba
          <= 0
        then
          Ok
            (Int32.add p.Partition.first_absolute_sector_lba p.Partition.sectors)
        else Error "Partitions overlap")
      (Ok 1l) sorted
  in
  Ok sorted

let v ?(disk_signature = 0l) partitions =
  let* partitions = validate_partitions partitions in
  Ok
    {
      bootstrap_code = String.init (218 + 216) (Fun.const '\000');
      original_physical_drive = 0;
      seconds = 0;
      minutes = 0;
      hours = 0;
      disk_signature;
      partitions;
    }

(* MBR layout: Wire codec for the full 512-byte sector.
   The codec decodes/encodes a raw record with all 4 partition slots;
   conversion to/from the public [t] type filters empty entries. *)

let sizeof = 512
let default_partition_start = 2048l
let bootstrap_code1_len = 218
let bootstrap_code2_len = 216

(* Wire fields for the MBR header *)
let f_bootstrap1 =
  Wire.Field.v "bootstrap1"
    (Wire.byte_array ~size:(Wire.int bootstrap_code1_len))

let funused1 = Wire.Field.v "unused1" (Wire.byte_array ~size:(Wire.int 2))
let f_drive = Wire.Field.v "drive" Wire.uint8
let f_seconds = Wire.Field.v "seconds" Wire.uint8
let f_minutes = Wire.Field.v "minutes" Wire.uint8
let f_hours = Wire.Field.v "hours" Wire.uint8

let f_bootstrap2 =
  Wire.Field.v "bootstrap2"
    (Wire.byte_array ~size:(Wire.int bootstrap_code2_len))

let f_disk_sig = Wire.Field.v "disk_signature" Wire.uint32
let funused2 = Wire.Field.v "unused2" (Wire.byte_array ~size:(Wire.int 2))
let f_part1 = Wire.Field.v "partition1" (Wire.codec Partition.codec)
let f_part2 = Wire.Field.v "partition2" (Wire.codec Partition.codec)
let f_part3 = Wire.Field.v "partition3" (Wire.codec Partition.codec)
let f_part4 = Wire.Field.v "partition4" (Wire.codec Partition.codec)
let f_sig1 = Wire.Field.v "signature1" Wire.uint8
let f_sig2 = Wire.Field.v "signature2" Wire.uint8

(* A partition entry that encodes as 16 zero bytes *)
let empty_partition =
  {
    Partition.active = false;
    first_absolute_sector_chs =
      { Geometry.cylinders = 0; heads = 0; sectors = 0 };
    ty = 0;
    last_absolute_sector_chs =
      { Geometry.cylinders = 0; heads = 0; sectors = 0 };
    first_absolute_sector_lba = 0l;
    sectors = 0l;
  }

(* Raw decoded MBR: all 4 partition slots present, signature bytes included *)
type raw = {
  bootstrap1 : string;
  unused1 : string;
  drive : int;
  seconds : int;
  minutes : int;
  hours : int;
  bootstrap2 : string;
  disk_signature_raw : int;
  unused2 : string;
  part1 : Partition.t;
  part2 : Partition.t;
  part3 : Partition.t;
  part4 : Partition.t;
  sig1 : int;
  sig2 : int;
}

let raw_codec =
  Wire.Codec.v "Mbr"
    (fun bootstrap1 unused1 drive seconds minutes hours bootstrap2
         disk_signature_raw unused2 part1 part2 part3 part4 sig1 sig2 ->
      {
        bootstrap1;
        unused1;
        drive;
        seconds;
        minutes;
        hours;
        bootstrap2;
        disk_signature_raw;
        unused2;
        part1;
        part2;
        part3;
        part4;
        sig1;
        sig2;
      })
    Wire.Codec.
      [
        (f_bootstrap1 $ fun r -> r.bootstrap1);
        (funused1 $ fun r -> r.unused1);
        (f_drive $ fun r -> r.drive);
        (f_seconds $ fun r -> r.seconds);
        (f_minutes $ fun r -> r.minutes);
        (f_hours $ fun r -> r.hours);
        (f_bootstrap2 $ fun r -> r.bootstrap2);
        (f_disk_sig $ fun r -> r.disk_signature_raw);
        (funused2 $ fun r -> r.unused2);
        (f_part1 $ fun r -> r.part1);
        (f_part2 $ fun r -> r.part2);
        (f_part3 $ fun r -> r.part3);
        (f_part4 $ fun r -> r.part4);
        (f_sig1 $ fun r -> r.sig1);
        (f_sig2 $ fun r -> r.sig2);
      ]

(* Filter a decoded partition: ty=0 means empty slot *)
let filter_partition (p : Partition.t) =
  if p.ty <> 0 then Ok (Some p)
  else if (not p.active) && p.first_absolute_sector_lba = 0l && p.sectors = 0l
  then Ok None
  else Error "Non-zero empty partition type"

(* Internal unmarshal from bytes buffer *)
let unmarshal_bytes (buf : bytes) : (t, string) result =
  let* () =
    if Bytes.length buf < sizeof then err_too_small (Bytes.length buf)
    else Ok ()
  in
  let* raw =
    Wire.Codec.decode raw_codec buf 0
    |> Result.map_error (Fmt.str "%a" Wire.pp_parse_error)
  in
  let* () =
    if raw.sig1 = 0x55 && raw.sig2 = 0xaa then Ok ()
    else err_bad_signature raw.sig1 raw.sig2
  in
  let* p1 = filter_partition raw.part1 in
  let* p2 = filter_partition raw.part2 in
  let* p3 = filter_partition raw.part3 in
  let* p4 = filter_partition raw.part4 in
  let partitions = List.filter_map Fun.id [ p1; p2; p3; p4 ] in
  Ok
    {
      bootstrap_code = raw.bootstrap1 ^ raw.bootstrap2;
      original_physical_drive = raw.drive;
      seconds = raw.seconds;
      minutes = raw.minutes;
      hours = raw.hours;
      disk_signature =
        Int32.of_int (Wire.Private.UInt32.to_int raw.disk_signature_raw);
      partitions;
    }

(* Internal marshal to bytes buffer *)
let marshal_bytes (buf : bytes) t =
  let parts = t.partitions in
  let nth i =
    try List.nth parts i
    with Failure _ | Invalid_argument _ -> empty_partition
  in
  let raw =
    {
      bootstrap1 = String.sub t.bootstrap_code 0 bootstrap_code1_len;
      unused1 = "\000\000";
      drive = t.original_physical_drive;
      seconds = t.seconds;
      minutes = t.minutes;
      hours = t.hours;
      bootstrap2 =
        String.sub t.bootstrap_code bootstrap_code1_len bootstrap_code2_len;
      disk_signature_raw =
        Wire.Private.UInt32.of_int (Int32.to_int t.disk_signature);
      unused2 = "\000\000";
      part1 = nth 0;
      part2 = nth 1;
      part3 = nth 2;
      part4 = nth 3;
      sig1 = 0x55;
      sig2 = 0xaa;
    }
  in
  Wire.Codec.encode raw_codec raw buf 0

(* Streaming API *)

let of_string s = unmarshal_bytes (Bytes.of_string s)

let to_string t =
  let buf = Bytes.make sizeof '\x00' in
  marshal_bytes buf t;
  Bytes.to_string buf

let read reader =
  match read_exactly reader sizeof with
  | Error e -> Error e
  | Ok buf -> unmarshal_bytes buf

let write writer t =
  let s = to_string t in
  Writer.write writer (Slice.of_string s)