Source file SecPackageRegistry.ml

let package_id_for_library library_id =
  if String.equal (MlFront_Core.LibraryId.vendor library_id) "Our" then None
  else if
    String.equal
      (MlFront_Core.LibraryId.full_name library_id)
      "MlFront_Attestation"
  then None
  else if
    String.equal (MlFront_Core.LibraryId.full_name library_id) "MlFront_Std"
  then None
  else Some (MlFront_Core.PackageId.of_library_id library_id)

let package_id_for_module m =
  Assumptions
  .distributions_fetched_for_all_modules_except_Our_vendor_and_mlfront_modules
    ();
  let library_id = MlFront_Core.StandardModuleId.library_id m in
  package_id_for_library library_id

let package_id_for_distcore (dist_core : MlFront_Thunk.ThunkDist.DistCore.t) =
  match dist_core with
  | { id = None; _ } -> None
  | { id = Some (_range, lib, _ver); _ } ->
      Assumptions
      .distributions_fetched_for_all_modules_except_Our_vendor_and_mlfront_modules
        ();
      package_id_for_library lib

let package_id_for_distribution (distribution : MlFront_Thunk.ThunkDist.t) =
  match distribution.id with
  | _range, lib, _ver ->
      Assumptions
      .distributions_fetched_for_all_modules_except_Our_vendor_and_mlfront_modules
        ();
      package_id_for_library lib

let package_id_for_distribution'
    (distribution : MlFront_Thunk.ThunkAst.distribution) =
  match distribution.distribution_id with
  | _range, lib, _ver ->
      Assumptions
      .distributions_fetched_for_all_modules_except_Our_vendor_and_mlfront_modules
        ();
      package_id_for_library lib

module DistCoreEntry = struct
  type t = MlFront_Thunk.ThunkDist.DistCore.t

  let package_id : t -> _ = package_id_for_distcore

  let version : t -> MlFront_Thunk.ThunkSemver64.t option = function
    | { id = None; _ } -> None
    | { id = Some (_range, _lib, ver); _ } -> Some ver

  let continuations_to_sign :
      t -> MlFront_Thunk.ThunkDist.DistCore.continuation_to_sign list = function
    | {
        continuations = { continuations_attestation = _; continuations_to_sign };
        _;
      } ->
        continuations_to_sign

  let compare a b =
    match (package_id a, package_id b) with
    | None, None -> 0
    | None, Some _ -> -1
    | Some _, None -> 1
    | Some p1, Some p2 ->
    match MlFront_Core.PackageId.compare p1 p2 with
    | 0 -> begin
        match (version a, version b) with
        | None, None -> 0
        | None, Some _ -> -1
        | Some _, None -> 1
        | Some v1, Some v2 -> MlFront_Thunk.ThunkSemver64.compare v1 v2
      end
    | n -> n
end

module DistCoreMap = Map.Make (DistCoreEntry)
module PackageIdMap = Map.Make (MlFront_Core.PackageId)

type 'a packages = 'a DistCoreMap.t PackageIdMap.t

type 'a t = {
  packages : 'a packages;
  show : MlFront_Thunk.ThunkDist.DistCore.t -> 'a -> string;
}

let create ~show () = { packages = PackageIdMap.empty; show }

let get_package_ids ({ packages; show = _ } : 'a t) :
    MlFront_Core.PackageId.t list =
  PackageIdMap.fold (fun pkg _acc acc -> pkg :: acc) packages []

let accept_if_not_present ?on_duplicate
    (dist_core : MlFront_Thunk.ThunkDist.DistCore.t) item
    ({ packages; show } : 'a t) : 'a t =
  Assumptions.accepted_distributions_are_present_in_trace_store ();
  let package_id = DistCoreEntry.package_id dist_core in
  let version = DistCoreEntry.version dist_core in
  match (package_id, version) with
  | None, _ | _, None ->
      (* No id or no version: cannot add to registry *)
      { packages; show }
  | Some package_id, Some package_semver -> begin
      let packages =
        PackageIdMap.update package_id
          (fun dist_core_map_opt ->
            match dist_core_map_opt with
            | None ->
                (* No package. Add it. *)
                Some (DistCoreMap.singleton dist_core item)
            | Some dist_core_map ->
                (* Add to the dist cores *)
                let dist_core_map =
                  DistCoreMap.update dist_core
                    (function
                      | None -> Some item
                      | Some item' -> begin
                          match on_duplicate with
                          | None -> Some item'
                          | Some f ->
                              f item'
                                (Printf.sprintf
                                   "Duplicate distribution id `%s` found in \
                                    `%s` and `%s`."
                                   (Printf.sprintf "%s@%s"
                                      (MlFront_Core.PackageId.full_name
                                         package_id)
                                      (MlFront_Thunk.ThunkSemver64.to_string
                                         package_semver))
                                   (show dist_core item) (show dist_core item'))
                        end)
                    dist_core_map
                in
                Some (DistCoreMap.add dist_core item dist_core_map))
          packages
      in
      { packages; show }
    end

let compare_int64_desc (maj1, min1) (maj2, min2) =
  match Int64.compare maj2 maj1 with 0 -> Int64.compare min2 min1 | n -> n

let sort_int64_desc = List.sort_uniq compare_int64_desc

(* runtime sorting test *)
let () =
  assert (
    sort_int64_desc [ (1L, 2L); (1L, 1L); (2L, 0L); (1L, 2L) ]
    = [ (2L, 0L); (1L, 2L); (1L, 1L) ])

let get_majmin_from_lookup :
    MlFront_Thunk.ThunkDist.DistCore.t * 'a -> int64 * int64 =
 fun (dist_core, _item) ->
  match DistCoreEntry.version dist_core with
  | None -> (0L, 0L)
  | Some { major; minor; patch = _; prerelease = _; build = _ } -> (major, minor)

let lookup_desc ({ packages; show = _ } : 'a t) (pkg : MlFront_Core.PackageId.t)
    : (MlFront_Thunk.ThunkDist.DistCore.t * 'a) list =
  match PackageIdMap.find_opt pkg packages with
  | None -> []
  | Some distcore_map ->
      DistCoreMap.fold
        (fun dist_core item acc -> (dist_core, item) :: acc)
        distcore_map []
      |> List.sort (fun a b ->
             compare_int64_desc (get_majmin_from_lookup a)
               (get_majmin_from_lookup b))

let get_next_versions_desc ({ packages; show = _ } : 'a t)
    (pkg : MlFront_Core.PackageId.t) : (int64 * int64) list =
  match PackageIdMap.find_opt pkg packages with
  | None -> []
  | Some distcore_map -> begin
      DistCoreMap.fold
        (fun dist_core _item acc ->
          let items =
            List.map
              (fun ({ majmin = _range, major, minor; producer = _ } :
                     MlFront_Thunk.ThunkDist.DistCore.continuation_to_sign) ->
                (major, minor))
              (DistCoreEntry.continuations_to_sign dist_core)
          in
          items :: acc)
        distcore_map []
      |> List.flatten |> sort_int64_desc
    end

let get_released_versions dist_core_map :
    (MlFront_Core.PackageId.t * MlFront_Thunk.ThunkSemver64.t) list =
  DistCoreMap.fold
    (fun dist_core _item acc ->
      match DistCoreEntry.(package_id dist_core, version dist_core) with
      | None, _ | _, None -> acc
      | Some pkg, Some ver -> (pkg, ver) :: acc)
    dist_core_map []
  |> List.sort (fun (pkg1, ver1) (pkg2, ver2) ->
         match MlFront_Core.PackageId.compare pkg1 pkg2 with
         | 0 -> MlFront_Thunk.ThunkSemver64.compare ver1 ver2
         | n -> n)

let create_dist_core () : MlFront_Thunk.ThunkDist.DistCore.t =
  {
    id = None;
    producer =
      {
        producer_application = None;
        producer_openbsd_signify = None;
        producer_github_slsa_v1_l2 = None;
        producer_github_slsa_v1_l3 = None;
      };
    license = { spdx = None; plaintext = None; markdown = None };
    continuations =
      { continuations_attestation = None; continuations_to_sign = [] };
  }

(** Load existing core distribution that matches [majmin], or continue with a
    continuation that matches [majmin], or make a new core distribution.

    The existing registry are unmodified.

    The return values are:
    - [`Exists (dist_core, dist_file)]: an existing distribution file matches
      the requested [majmin]. The [dist_core] is complete and can be used as is.
      The [dist_file] is the path to the existing distribution file.
    - [`Continued dist_core]: a continuation in an existing distribution file
      matches the requested [majmin]. The [dist_core] is {b partial} and its
      [None] fields must be filled in.
    - [`NoMatchingContinuation]: there are existing distribution files, but none
      of their continuations match the requested [majmin]. This is an error.
    - [`Superseded latest]: the requested [majmin] is less than the latest
      released version [latest]. This is an error.
    - [`Error msg]: some other error occurred, with [msg] describing the error.
*)
let characterize ~majmin pkg ({ packages; show } : 'a t) :
    [ `Exists of MlFront_Thunk.ThunkDist.DistCore.t * 'a
    | `Continued of MlFront_Thunk.ThunkDist.DistCore.t
    | `NoMatchingContinuation
    | `Superseded of MlFront_Thunk.ThunkSemver64.t
    | `Error of string ] =
  let find_major, find_minor = (fst majmin, snd majmin) in
  let find_ver =
    MlFront_Thunk.ThunkSemver64.from_parts find_major find_minor 0L [] []
  in
  (* Find the distribution core map for the package [pkg] *)
  let distcore_map =
    match PackageIdMap.find_opt pkg packages with
    | None -> DistCoreMap.empty
    | Some distcore_map -> distcore_map
  in
  (* Find latest version *)
  let latest_ver =
    let released_versions = get_released_versions distcore_map in
    List.fold_left
      (fun acc (_lib, released_ver) ->
        match acc with
        | None -> Some released_ver
        | Some latest_ver ->
            if MlFront_Thunk.ThunkSemver64.compare latest_ver released_ver < 0
            then Some released_ver
            else acc)
      None released_versions
  in
  match (find_ver, latest_ver) with
  | Some find_ver, Some latest_ver
    when MlFront_Thunk.ThunkSemver64.compare find_ver latest_ver < 0 ->
      (* If the requested version is less than latest released version, stop *)
      `Superseded latest_ver
  | _ -> begin
      (* Find the MAJOR.MINOR.0 version *)
      let same_major_minor =
        DistCoreMap.fold
          (fun dist_core item
               (acc : (MlFront_Thunk.ThunkDist.DistCore.t * 'a) option) ->
            (* OPTIMIZATION ALERT: Filter out anything not matching (major,minor) *)
            match DistCoreEntry.version dist_core with
            | Some ({ major; minor; patch; _ } as ver)
              when Int64.equal major find_major && Int64.equal minor find_minor
              ->
                (* Warn if patch >= 1 *)
                if Int64.compare patch 1L >= 0 then begin
                  Printf.eprintf
                    "[warning] Distribution file `%s` has patch version %Ld in \
                     its `id` version %s. Distributions contain next version \
                     \"continuations\" that are immutable; they cannot be \
                     patched.\n\
                     %!"
                    (show dist_core item) patch
                    (MlFront_Thunk.ThunkSemver64.to_string ver);
                  acc
                end
                else Some (dist_core, item)
            | _ -> acc)
          distcore_map None
      in
      match same_major_minor with
      | Some (dist_core, item) -> `Exists (dist_core, item)
      | None -> (
          (* Find the continuation MAJOR.MINOR with the highest `id` version *)
          let next_major_minor =
            DistCoreMap.fold
              (fun dist_core _item
                   (acc :
                     (MlFront_Thunk.ThunkDist.DistCore.t
                     * MlFront_Thunk.ThunkDist.DistCore.continuation_to_sign)
                     option) ->
                match DistCoreEntry.version dist_core with
                | Some ver -> (
                    (* Any continuation have the searched-for major.minor? *)
                    let konts = DistCoreEntry.continuations_to_sign dist_core in
                    let kont =
                      List.find_opt
                        (fun ({ majmin = _range, major, minor; _ } :
                               MlFront_Thunk.ThunkDist.DistCore
                               .continuation_to_sign) ->
                          Int64.equal major find_major
                          && Int64.equal minor find_minor)
                        konts
                    in
                    match (acc, kont) with
                    | _, None ->
                        (* Continuation not found *)
                        acc
                    | None, Some kont ->
                        (* First match *)
                        Some (dist_core, kont)
                    | Some (acc_core, _acc_kont), Some kont -> (
                        (* Break ties. If the incoming `id` is higher than the existing `id`
                         (ie. incoming is newer), then use incoming. *)
                        let acc_ver = DistCoreEntry.version acc_core in
                        match acc_ver with
                        | None -> Some (dist_core, kont)
                        | Some acc_ver ->
                            if
                              MlFront_Thunk.ThunkSemver64.compare ver acc_ver
                              > 0
                            then Some (dist_core, kont)
                            else acc))
                | _ -> acc)
              distcore_map None
          in
          match next_major_minor with
          | Some
              ( ({ id = Some (idrange, idlibrary, _idversion); _ } as dist_core),
                {
                  majmin = _;
                  producer =
                    {
                      producer_openbsd_signify =
                        Some kont_producer_openbsd_signify;
                      _;
                    };
                } ) ->
              (* Rotate the continuations *)
              let kont_idrange =
                (* nit: idrange is pre-continuation so it is incorrect *)
                idrange
              in
              let kont_library = idlibrary in
              let kont_version =
                MlFront_Thunk.ThunkSemver64.from_parts find_major find_minor 0L
                  [] []
              in
              let kont_producer : MlFront_Thunk.ThunkDist.DistCore.producer =
                (* Keep original producer _except_ the signify key is rotated in *)
                {
                  producer_application = dist_core.producer.producer_application;
                  producer_openbsd_signify = Some kont_producer_openbsd_signify;
                  producer_github_slsa_v1_l2 =
                    dist_core.producer.producer_github_slsa_v1_l2;
                  producer_github_slsa_v1_l3 =
                    dist_core.producer.producer_github_slsa_v1_l3;
                }
              in
              let kont_continuation_to_sign :
                  MlFront_Thunk.ThunkDist.DistCore.continuation_to_sign list =
                (* keep all continuations that are strictly after this continuation *)
                List.filter
                  (fun ({ majmin = _range, major, minor; _ } :
                         MlFront_Thunk.ThunkDist.DistCore.continuation_to_sign)
                     ->
                    let major_c = Int64.compare major find_major in
                    if major_c > 0 then true
                    else if major_c = 0 then Int64.compare minor find_minor > 0
                    else false)
                  dist_core.continuations.continuations_to_sign
              in
              begin
                match kont_version with
                | None -> `Error "Invalid continuation version"
                | Some kont_version ->
                    `Continued
                      {
                        id = Some (kont_idrange, kont_library, kont_version);
                        producer = kont_producer;
                        license = dist_core.license;
                        continuations =
                          {
                            continuations_attestation = None;
                            continuations_to_sign = kont_continuation_to_sign;
                          };
                      }
              end
          | _ ->
              (* No matching continuation. That is an error! *)
              `NoMatchingContinuation)
    end

module DistributedModules = struct
  type majmin = { major : int64; minor : int64 } [@@deriving ord]

  module MajMinSet = Set.Make (struct
    type t = majmin

    let compare = compare_majmin
  end)

  module ModuleMap = Map.Make (MlFront_Core.StandardModuleId)

  type t = MajMinSet.t ModuleMap.t

  type is_distributed_response =
    | Distributed
    | NotDistributed of {
        available_majmin_versions : (int64 * int64) list;
        available_modules : MlFront_Core.StandardModuleId.t list;
      }

  let empty : t = ModuleMap.empty

  let add (module_id : MlFront_Core.StandardModuleId.t)
      (semver : MlFront_Thunk.ThunkSemver64.t) (dm : t) : t =
    let existing =
      match ModuleMap.find_opt module_id dm with
      | None -> MajMinSet.empty
      | Some s -> s
    in
    let majmin = { major = semver.major; minor = semver.minor } in
    let updated = MajMinSet.add majmin existing in
    ModuleMap.add module_id updated dm

  module ResultObserver =
    MlFront_Thunk.ThunkParsers.Results.MakeObserverWithDiagnoseErrors
      (MlFront_Thunk.Diagnose.Diagnose.ConsolePlainStyle)

  let from_distribution
      ({
         id = _;
         producer = _;
         license = _;
         continuations = _;
         build =
           {
             build_attestation = _;
             build_to_sign =
                {
                  build_bundle_modver = _;
                  build_modules;
                  build_producer_accepts = _;
                  build_bundle_canonical = _;
                  build_distmeta = _;
                  build_package = _;
                };
            };
        } :
        MlFront_Thunk.ThunkDist.t) : t =
    List.fold_left
      (fun acc (_range, module_id, module_semver) ->
        add module_id module_semver acc)
      empty build_modules

  let is_distributed (module_id : MlFront_Core.StandardModuleId.t)
      (semver : MlFront_Thunk.ThunkSemver64.t) (dm : t) =
    let get_available_modules () =
      ModuleMap.bindings dm |> List.map fst
      |> List.sort_uniq MlFront_Core.StandardModuleId.compare
    in
    match ModuleMap.find_opt module_id dm with
    | None ->
        NotDistributed
          {
            available_majmin_versions = [];
            available_modules = get_available_modules ();
          }
    | Some s ->
        let majmin = { major = semver.major; minor = semver.minor } in
        if MajMinSet.mem majmin s then Distributed
        else
          let available_majmin_versions =
            MajMinSet.elements s
            |> List.map (fun { major; minor } -> (major, minor))
          in
          NotDistributed
            {
              available_majmin_versions;
              available_modules = get_available_modules ();
            }
end