hcs / lib / multipart.ml

(** Multipart form data parsing.

    {[
      (* Non-streaming *)
      match Multipart.parse req with
      | Ok parts -> Multipart.find_part "name" parts
      | Error e -> Response.bad_request (Multipart.error_to_string e)

      (* Streaming large files *)
      match Multipart.create_parser req with
      | Ok parser ->
          Multipart.iter_parts (fun part ->
            Stream.Async.iter write_chunk part.body
          ) parser
      | Error e -> ...
    ]} *)

module Hmf = Http_multipart_formdata

type part = {
  name : string;
  filename : string option;
  content_type : string;
  data : string;
}

type error =
  | Missing_content_type
  | Not_multipart
  | Missing_boundary
  | Invalid_boundary of string
  | Parse_error of string

let error_to_string = function
  | Missing_content_type -> "Missing Content-Type header"
  | Not_multipart -> "Content-Type is not multipart/form-data"
  | Missing_boundary -> "Missing boundary parameter in Content-Type"
  | Invalid_boundary msg -> "Invalid boundary: " ^ msg
  | Parse_error msg -> "Parse error: " ^ msg

let is_multipart (req : Request.t) =
  match Request.content_type req with
  | None -> false
  | Some ct ->
      let ct_lower = String.lowercase_ascii ct in
      let prefix = "multipart/form-data" in
      String.length ct_lower >= String.length prefix
      && String.sub ct_lower 0 (String.length prefix) = prefix

let boundary (req : Request.t) : (Hmf.boundary, error) result =
  match Request.content_type req with
  | None -> Error Missing_content_type
  | Some ct -> (
      if not (is_multipart req) then Error Not_multipart
      else
        match Hmf.boundary ct with
        | Ok b -> Ok b
        | Error msg ->
            if String.length msg > 0 && msg.[0] = '\'' then
              Error Missing_boundary
            else Error (Invalid_boundary msg))

let parse (req : Request.t) : (part list, error) result =
  match boundary req with
  | Error e -> Error e
  | Ok b -> (
      let body = Request.body req in
      match Hmf.parts b body with
      | Error msg -> Error (Parse_error msg)
      | Ok parts ->
          Ok
            (List.map
               (fun (_name, (header, body)) ->
                 {
                   name = Hmf.name header;
                   filename = Hmf.filename header;
                   content_type = Hmf.content_type header;
                   data = body;
                 })
               parts))

let find_part name parts = List.find_opt (fun p -> p.name = name) parts

let find_file name parts =
  List.find_opt (fun p -> p.name = name && Option.is_some p.filename) parts

let to_assoc parts : (string * part) list =
  List.map (fun p -> (p.name, p)) parts

type stream_part = {
  name : string;
  filename : string option;
  content_type : string;
  body : Cstruct.t Stream.Async.t;
}

type parser = {
  reader : Hmf.reader;
  body_stream : unit -> Cstruct.t option;
  mutable current_state : Hmf.read;
  mutable finished : bool;
}

let create_parser (req : Request.t) : (parser, error) result =
  match boundary req with
  | Error e -> Error e
  | Ok b ->
      let body_stream =
        Request.body_reader req |> fun r -> r.H1_server.read_stream
      in
      let reader = Hmf.reader b `Incremental in
      let current_state = Hmf.read reader in
      Ok { reader; body_stream; current_state; finished = false }

let rec advance_state parser =
  match parser.current_state with
  | `Awaiting_input continue -> (
      match parser.body_stream () with
      | None ->
          parser.current_state <- continue `Eof;
          advance_state parser
      | Some chunk ->
          parser.current_state <- continue (`Cstruct chunk);
          advance_state parser)
  | other -> other

let make_body_stream parser : Cstruct.t Stream.Async.t =
  let finished = ref false in
  fun () ->
    if !finished then None
    else
      let rec read_body () =
        match advance_state parser with
        | `Body chunk ->
            parser.current_state <- Hmf.read parser.reader;
            Some chunk
        | `Body_end ->
            finished := true;
            parser.current_state <- Hmf.read parser.reader;
            None
        | `End ->
            finished := true;
            parser.finished <- true;
            None
        | `Error _ ->
            finished := true;
            parser.finished <- true;
            None
        | `Header _ ->
            finished := true;
            None
        | `Awaiting_input _ -> read_body ()
      in
      read_body ()

let next_part parser : stream_part option =
  if parser.finished then None
  else
    let rec find_header () =
      match advance_state parser with
      | `Header header ->
          parser.current_state <- Hmf.read parser.reader;
          let body = make_body_stream parser in
          Some
            {
              name = Hmf.name header;
              filename = Hmf.filename header;
              content_type = Hmf.content_type header;
              body;
            }
      | `End ->
          parser.finished <- true;
          None
      | `Error _ ->
          parser.finished <- true;
          None
      | `Body _ | `Body_end ->
          parser.current_state <- Hmf.read parser.reader;
          find_header ()
      | `Awaiting_input _ -> find_header ()
    in
    find_header ()

let iter_parts f parser : (unit, error) result =
  let rec loop () =
    match next_part parser with
    | None ->
        if parser.finished then Ok () else Error (Parse_error "Unexpected end")
    | Some part ->
        f part;
        Stream.Async.drain part.body;
        loop ()
  in
  loop ()

let fold_parts f init parser =
  let rec loop acc =
    match next_part parser with
    | None ->
        if parser.finished then Ok acc else Error (Parse_error "Unexpected end")
    | Some part ->
        let acc' = f part acc in
        Stream.Async.drain part.body;
        loop acc'
  in
  loop init

let collect_body (part : stream_part) : string =
  Stream.Async.cstructs_to_string part.body

let parse_assoc req =
  match parse req with Error e -> Error e | Ok parts -> Ok (to_assoc parts)