Source file vif_handler_unix.ml

open Vif_core

type ('c, 'value) t = (Httpcats.Server.flow, 'c, 'value) Handler.t

let pwd = Fpath.v (Unix.getcwd ())
let tree = Conan_light.tree
let ( let@ ) finally fn = Fun.protect ~finally fn

let file ?offset path =
  let open Flux in
  let buf = Bytes.create 0x7ff in
  let fn bqueue =
    let fd = Unix.openfile path Unix.[ O_RDONLY ] 0o644 in
    let resource = Miou.Ownership.create ~finally:Unix.close fd in
    Miou.Ownership.own resource;
    let@ () =
     fun () ->
      Miou.Ownership.release resource;
      Bqueue.close bqueue
    in
    let _ =
      match offset with
      | Some offset -> Unix.lseek fd offset Unix.SEEK_SET
      | None -> 0
    in
    let rec go () =
      let len = Unix.read fd buf 0 (Bytes.length buf) in
      if len > 0 then (
        let str = Bytes.sub_string buf 0 len in
        Bqueue.put bqueue str; go ())
    in
    go ()
  in
  Source.with_task ~size:0x7ff fn

let sha256sum path =
  let path = Fpath.to_string path in
  if Sys.file_exists path = false || Sys.is_directory path then
    invalid_arg "sha256sum";
  let fd = Unix.openfile path Unix.[ O_RDONLY ] 0o644 in
  let finally () = Unix.close fd in
  Fun.protect ~finally @@ fun () ->
  let stat = Unix.fstat fd in
  let ba =
    Unix.map_file fd Bigarray.char Bigarray.c_layout false
      [| stat.Unix.st_size |]
  in
  let ba = Bigarray.array1_of_genarray ba in
  let hash = Digestif.SHA256.digest_bigstring ba in
  Digestif.SHA256.to_hex hash

let mime_type path =
  match Conan_unix.run_with_tree Conan_light.tree (Fpath.to_string path) with
  | Ok m -> Conan.Metadata.mime m
  | Error _ -> None
  | exception _ -> None

(* Returns [None] if we have a cache hit; otherwise returns [Some etag'] where
   [etag'] is [etag] if present otherwise the sha256 computed etag. *)
let cached_or_etag ?etag req target =
  let hdrs = Request.headers req in
  let etag =
    match etag with
    | Some etag -> etag
    | None -> Fmt.str {|"%s"|} (sha256sum target)
  in
  match Headers.get hdrs "if-none-match" with
  | Some etag' when String.equal etag etag' -> None
  | Some _ | None -> Some etag

let valid ~top target =
  Fpath.is_prefix top target
  && Sys.file_exists (Fpath.to_string target)
  && Sys.is_directory (Fpath.to_string target) = false

let pp_msg ppf (`Msg msg) = Fmt.string ppf msg

let trim lst =
  let lst = List.drop_while (( = ) "") lst in
  let lst = List.drop_while (( = ) "") (List.rev lst) in
  List.rev lst

module K = struct
  type t = Fpath.t

  let equal = Fpath.equal
  let hash = Hashtbl.hash
end

module V = struct
  type t = { mtime: float; mime: string }

  let weight _ = 1
end

module Cache = Lru.M.Make (K) (V)

let cached_on_server_size stat abs_path cache =
  match Cache.find abs_path cache with
  | Some { mtime; mime } when mtime >= stat.Unix.st_mtime -> Some mime
  | Some _ ->
      Cache.remove abs_path cache;
      None
  | None -> None

let static ?(top = pwd) =
  ();
  let cache = Cache.create ~random:true 0x100 in
  fun req target _server _ ->
    let target = String.split_on_char '/' target in
    let target = trim target in
    let target = String.concat "/" target in
    let abs_path =
      let ( let* ) = Result.bind in
      let* x = Fpath.of_string target in
      Ok Fpath.(normalize (top // x))
    in
    match (Request.meth req, abs_path) with
    | `GET, Ok abs_path when valid ~top abs_path -> begin
        let ( let* ) = Response.bind in
        let process =
          match cached_or_etag req abs_path with
          | None ->
              let* () = Response.with_string req "" in
              Response.respond `Not_modified
          | Some etag ->
              let stat = Unix.stat (Fpath.to_string abs_path) in
              let mime =
                match cached_on_server_size stat abs_path cache with
                | Some _ as value -> value
                | None ->
                    let mime = mime_type abs_path in
                    let fn mime =
                      let value = { V.mtime= stat.Unix.st_mtime; mime } in
                      Cache.add abs_path value cache
                    in
                    Option.iter fn mime; mime
              in
              let src = file (Fpath.to_string abs_path) in
              let* _ = Response.content_length stat.Unix.st_size in
              let* () =
                match mime with
                | Some mime -> Response.add ~field:"content-type" mime
                | None -> Response.return ()
              in
              let field = "etag" in
              let* () = Response.add ~field etag in
              let* () = Response.with_source req src in
              Response.respond `OK
        in
        Some process
      end
    | _ -> None