imageserver.ml

open Astring

let request_uri =
  try Unix.unsafe_getenv "REQUEST_URI"
  with Not_found -> ""

let request_method =
  match Unix.unsafe_getenv "REQUEST_METHOD" with
  | exception Not_found -> `Get
  | "GET" -> `Get
  | "POST" -> `Post
  | "HEAD" -> `Head
  | m -> `Other m

type request = {
  name: string;
  start: int;
  end_: int;
}

let sanitise s =
  String.filter (function '/' -> false | _ -> true) s

let parse_filename filename =
  let (let*) = Option.bind in
  let to_pos_int x = Option.bind (String.to_int x) (fun x -> if x >= 0 then Some x else None) in
  let* (x, extension) = String.cut ~rev:true ~sep:"." filename in
  let* (x, end_) = String.cut ~rev:true ~sep:"-" x in
  let* end_ = to_pos_int end_ in
  let* (base, start) = String.cut ~rev:true ~sep:"-" x in
  let* start = to_pos_int start in
  if start < end_ then
    Some {
      name = sanitise (base ^ "." ^ extension);
      start;
      end_;
    }
  else None

let respond_with_error status fmt =
  begin match status with
  | `Not_found ->
    print_string "Status: 404 Not found\n"
  | `Method_not_allowed ->
    print_string "Status: 405 Method Not Allowed\n"
  end;
  print_string "Content-type: text/plain\n\n";
  Printf.ksprintf print_endline fmt

let respond_with_unix_error status (e, func, param) =
  respond_with_error status "%s %s %s" (Unix.error_message e) func param

let catch_unix f =
  try Ok (f ()) with Unix.Unix_error (e, func, param) -> Error (e, func, param)

let write_stdout s =
  let bytes_written = Unix.write_substring Unix.stdout s 0 (String.length s) in
  assert (bytes_written = String.length s)

let rec respond_loop remaining fd buffer did_send_headers =
  if remaining >= 0 then (
    match catch_unix (fun () -> Unix.read fd buffer 0 (min remaining (Bytes.length buffer))) with
    | Error e ->
      if not did_send_headers then respond_with_unix_error `Not_found e
    | Ok bytes_read when bytes_read <= 0 ->
      if not did_send_headers then respond_with_error `Not_found "read returned %d" bytes_read
    | Ok bytes_read ->
      if not did_send_headers then (
        write_stdout "Content-type: application/octet-stream\n\n"
      );
      match catch_unix (fun () -> Unix.write Unix.stdout buffer 0 bytes_read) with
      | Error e -> if not did_send_headers then respond_with_unix_error `Not_found e
      | Ok bytes_written ->
        assert (bytes_written = bytes_read);
        respond_loop (remaining - bytes_read) fd buffer true
  )

let () =
  let source_dir =
    try Unix.unsafe_getenv "IMAGESERVER_DIR"
    with Not_found -> ""
  in
  if request_method <> `Get && request_method <> `Head then
    respond_with_error `Method_not_allowed "Bad method"
  else
  let filename = Filename.basename request_uri in
  match parse_filename filename with
  | None ->
    respond_with_error `Not_found "Bad range"
  | Some r ->
    let fd = catch_unix (fun () -> Unix.openfile (Filename.concat source_dir r.name) [Unix.O_RDONLY] 0) in
    match fd with
    | Error e ->
      respond_with_unix_error `Not_found e
    | Ok fd ->
      Fun.protect
        ~finally:(fun () ->
            let _ : (unit, _) result = catch_unix (fun () -> Unix.close fd) in
            ()
          )
        (fun () ->
           let offset = catch_unix (fun () -> Unix.lseek fd r.start Unix.SEEK_SET) in
           match offset with
           | Error e -> respond_with_unix_error `Not_found e
           | Ok _offset -> (* Note: Actual seek offset provides no useful information *)
             let length = r.end_ - r.start in
             let chunk_size = min 65536 length in
             let buffer = Bytes.create chunk_size in
             respond_loop length fd buffer false)