plug_parsers_related

defmodule Plug.Parsers.RELATED do
  @moduledoc """
  Parses multipart/related request body.

  This is a partial implementation of RFC2387. This implementation is
  heavily influence by `Plug.Parsers.MULTIPART` which is shipped with Plug.

  The following choices have been made regarding support for the RFC2387
  specification:
  * start-info parameter - The start-info parameter is ignored if specified.
  * If no start parameter is given the first MIME part is marked as the root
    part.
  * The Content-ID part header is REQUIRED for each part. If Content-ID is not
    present the part will be ignored.
  * Content-Disposition is used to create `Plug.Upload` structs in the output

  A HTTP request with the following Content-Type header and body:
  ```
  Content-Type: multipart/related; type="application/soap+xml"; boundary=----w58EW1cEpjzydSCq; start="part1"
  ```

  ```
  ------w58EW1cEpjzydSCq\r
  Content-Disposition: form-data; name=\"info\"; filename=\"foo.txt\"\r
  Content-ID: part1\r
  Content-Type: text/plain\r
  \r
  bar
  \r
  ------w58EW1cEpjzydSCq\r
  Content-ID: part2\r
  Content-Type: text/plain\r
  \r
  foobar\r
  ------w58EW1cEpjzydSCq\r
  Content-Type: text/plain\r
  \r
  ignored\r
  ------w58EW1cEpjzydSCq--\r
  ```
  Will be available in `Plug.Conn` as follows:
  ```
  %{
    "part1" => %{
      "body" => %Plug.Upload{
        content_type: "text/plain",
        filename: "foo.txt",
        path: "/var/folders/hf/pd3rksm54wj6hfznxz9tv36c0000gn/T//plug-1540/multipart-1540203239-599604070647209-1"
      },
      "start" => true
  },
    "part2" => %{
      "body" => "foobar"
    }
  }
  ```
  ## Options

  All options supported by `Plug.Conn.read_body/2` are also supported here.
  They are repeated here for convenience:

    * `:length` - sets the maximum number of bytes to read from the request,
      defaults to 8_000_000 bytes
    * `:read_length` - sets the amount of bytes to read at one time from the
      underlying socket to fill the chunk, defaults to 1_000_000 bytes
    * `:read_timeout` - sets the timeout for each socket read, defaults to
      15_000ms

  So by default, `Plug.Parsers` will read 1_000_000 bytes at a time from the
  socket with an overall limit of 8_000_000 bytes.

  Besides the options supported by `Plug.Conn.read_body/2`, the multipart parser
  also checks for `:headers` option that contains the same `:length`, `:read_length`
  and `:read_timeout` options which are used explicitly for parsing multipart
  headers.
  """

  require Record
  Record.defrecordp :state, start: nil, count: 0, parts: []

  @behaviour Plug.Parsers

  def init(opts) do
    opts
  end

  def parse(conn, "multipart", "related", %{"boundary" => _} = headers, opts) do
    try do
      start = Map.get(headers, "start", nil)
      acc = state(start: start)
      parse_multipart(conn, opts, acc)
    rescue
      # Do not ignore upload errors
      e in Plug.UploadError ->
        reraise e, System.stacktrace()

      # All others are wrapped
      e ->
        reraise Plug.Parsers.ParseError.exception(exception: e), System.stacktrace()
    end
  end

  #####
  # Shouldn't get here
  #####
  def parse(conn, _type, _subtype, _headers, _opts) do
    {:next, conn}
  end

  ## Multipart

  #####
  # Only gets here once
  #####
  defp parse_multipart(conn, opts, acc) do
    # Remove the length from options as it would attempt
    # to eagerly read the body on the limit value.
    {limit, opts} = Keyword.pop(opts, :length, 8_000_000)

    # The read length is now our effective length per call.
    {read_length, opts} = Keyword.pop(opts, :read_length, 1_000_000)
    opts = [length: read_length, read_length: read_length] ++ opts

    # The header options are handled individually.
    {headers_opts, opts} = Keyword.pop(opts, :headers, [])

    read_result = Plug.Conn.read_part_headers(conn, headers_opts)
    {:ok, limit, acc, conn} = parse_multipart(read_result, limit, opts, headers_opts, acc)

    if limit > 0 do
      # Was able to read the multipart because it was less then the limit
      parts = state(acc, :parts)
      |> Enum.map(fn(r) -> {elem(r, 0), Map.delete(elem(r,1), "index")} end)
      |> Enum.reduce(%{}, &Plug.Conn.Query.decode_pair/2)
      {:ok, parts, conn}
    else
      {:error, :too_large, conn}
    end
  end

  #####
  # Process multipart, gets here for each part
  #####
  defp parse_multipart({:ok, headers, conn}, limit, opts, headers_opts, acc) when limit >= 0 do
    {conn, limit, acc} = parse_multipart_headers(headers, conn, limit, opts, acc)
    read_result = Plug.Conn.read_part_headers(conn, headers_opts)
    parse_multipart(read_result, limit, opts, headers_opts, acc)
  end

  #####
  # Shouldn't get here
  #####
  defp parse_multipart({:ok, _headers, conn}, limit, _opts, _headers_opts, acc) do
    {:ok, limit, acc, conn}
  end

  #####
  # All the part headers have been processed
  #####
  defp parse_multipart({:done, conn}, limit, _opts, _headers_opts, acc) do
    {:ok, limit, acc, conn}
  end

  defp parse_multipart_headers(headers, conn, limit, opts, acc) do
    case multipart_type(headers, List.keyfind(headers, "content-id", 0), List.keyfind(headers, "content-disposition", 0), acc) do
      {:binary, cid, part} ->
        {:ok, limit, body, conn} =
          parse_multipart_body(Plug.Conn.read_part_body(conn, opts), limit, opts, "")

        Plug.Conn.Utils.validate_utf8!(body, Plug.Parsers.BadEncodingError, "multipart body")
        {conn, limit,
         state(acc, count: state(acc, :count) + 1, parts: [{cid, Map.put(part, "body", body)} | state(acc, :parts)])}

      {:file, name, path, %{"body" => %Plug.Upload{}} = uploaded} ->
        {:ok, file} = File.open(path, [:write, :binary, :delayed_write, :raw])

        {:ok, limit, conn} =
          parse_multipart_file(Plug.Conn.read_part_body(conn, opts), limit, opts, file)

        :ok = File.close(file)
        {conn, limit,
         state(acc, count: state(acc, :count) + 1, parts: [{name, uploaded} | state(acc, :parts)])}

      :skip ->
        {conn, limit, acc}
    end
  end

  defp parse_multipart_body({:more, tail, conn}, limit, opts, body)
       when limit >= byte_size(tail) do
    read_result = Plug.Conn.read_part_body(conn, opts)
    parse_multipart_body(read_result, limit - byte_size(tail), opts, body <> tail)
  end

  defp parse_multipart_body({:more, tail, conn}, limit, _opts, body) do
    {:ok, limit - byte_size(tail), body, conn}
  end

  defp parse_multipart_body({:ok, tail, conn}, limit, _opts, body)
       when limit >= byte_size(tail) do
    {:ok, limit - byte_size(tail), body <> tail, conn}
  end

  defp parse_multipart_body({:ok, tail, conn}, limit, _opts, body) do
    {:ok, limit - byte_size(tail), body, conn}
  end

  defp parse_multipart_file({:more, tail, conn}, limit, opts, file)
       when limit >= byte_size(tail) do
    IO.binwrite(file, tail)
    read_result = Plug.Conn.read_part_body(conn, opts)
    parse_multipart_file(read_result, limit - byte_size(tail), opts, file)
  end

  defp parse_multipart_file({:more, tail, conn}, limit, _opts, _file) do
    {:ok, limit - byte_size(tail), conn}
  end

  defp parse_multipart_file({:ok, tail, conn}, limit, _opts, file)
       when limit >= byte_size(tail) do
    IO.binwrite(file, tail)
    {:ok, limit - byte_size(tail), conn}
  end

  defp parse_multipart_file({:ok, tail, conn}, limit, _opts, _file) do
    {:ok, limit - byte_size(tail), conn}
  end

  ## Helpers

  #
  # Header could have the following keys:
  #  Content-Type: application/octet-stream
  #  Content-Transfer-Encoding: base64
  #  Content-Id: <PO_Image@example.com>
  #  Content-Description: WSS XML Encryption message; type="image/jpeg"
  #  Content-Disposition: attachment; filename="fname.ext"
  #
  # multipart/related doesn't have to have a Content-Disposition but it must have a Content-Id
  defp multipart_type(_, nil, _, _) do
    :skip
  end

  # Header has a Content-Id 
  defp multipart_type(_, {"content-id", cid}, nil, acc) do
    {:binary, cid, %{"index" => state(acc, :count)} |> set_start(acc, cid)}
  end

  # Header has a Content-Id and Content-Disposition
  defp multipart_type(headers, {"content-id", _}, {"content-disposition", disposition}, acc) do
    with [_, params] <- :binary.split(disposition, ";"),
         %{"filename" => name} = params <- Plug.Conn.Utils.params(params) do
      handle_disposition(params, name, headers, acc)
    else
      _ -> :skip
    end
  end

  defp part_info(acc, cid, body) do
    %{"index" => state(acc, :count), "body" => body}
    |> set_start(acc, cid)
  end

  defp set_start(part, {:state, nil, _, _} = acc, _) do
    if (state(acc, :count) == 0) do
      Map.put(part, "start", true)
    else
      part
    end
  end

  defp set_start(part, acc, cid) do
    if (state(acc, :start) == cid) do
      Map.put(part, "start", true)
    else
      part
    end
  end

  defp handle_disposition(params, name, headers, acc) do
    case Map.fetch(params, "filename") do
      {:ok, ""} ->
        :skip

      {:ok, filename} ->
        path = Plug.Upload.random_file!("multipart")
        content_type = get_header(headers, "content-type")
        cid = get_header(headers, "content-id")
        upload = %Plug.Upload{filename: filename, path: path, content_type: content_type}
        {:file, cid, path, part_info(acc, cid, upload)}

      :error ->
        {:binary, name}
    end
  end

  defp get_header(headers, key) do
    case List.keyfind(headers, key, 0) do
      {^key, value} -> value
      nil -> nil
    end
  end
end