Source file request.ml

module Option = struct
  let equal eq a b =
    match a, b with
    | Some a, Some b -> eq a b
    | None, None -> true
    | _, _ -> false
end

module String = struct
  include String

  (* (c) @dbuenzli with love *)

  let add_sub ~no_empty s ~start ~stop acc =
    if start = stop
    then ( if no_empty then acc else "" :: acc )
    else String.sub s start (stop - start) :: acc

  let fcuts ~no_empty ~sep s =
    let sep_len = String.length sep in
    if sep_len = 0 then invalid_arg "empty separator" ;
    let s_len = String.length s in
    let max_sep_idx = sep_len - 1 in
    let max_s_idx = s_len - sep_len in
    let rec check_sep start i k acc =
      if k > max_sep_idx
      then
        let new_start = i + sep_len in
        scan new_start new_start (add_sub ~no_empty s ~start ~stop:i acc)
      else
        if s.[i + k] = sep.[k]
        then check_sep start i (k + 1) acc
        else scan start (i + 1) acc
    and scan start i acc =
      if i > max_s_idx
      then
        if start = 0
        then ( if no_empty && s_len = 0 then [] else [ s ] )
        else List.rev (add_sub ~no_empty s ~start ~stop:s_len acc)
      else
        if s.[i] = sep.[0]
        then check_sep start i 1 acc
        else scan start (i + 1) acc in
    scan 0 0 []

  let cuts ?(empty= true) ~sep s =
    fcuts ~no_empty:(not empty) ~sep s
end

type t =
  [ `Hello of Domain.t
  | `Mail of Reverse_path.t * (string * string option) list
  | `Recipient of Forward_path.t * (string * string option) list
  | `Expand of string
  | `Data
  | `Data_end
  | `Help of string option
  | `Noop of string option
  | `Verify of string
  | `Reset
  | `Quit
  | `Verb of string * string list
  | `Payload of string ]

let equal_parameters a b =
  let a = List.sort (fun (ka, _) (kb, _) -> String.compare ka kb) a in
  let b = List.sort (fun (ka, _) (kb, _) -> String.compare ka kb) b in
  let equal_values a b = match a, b with
    | Some a, Some b -> String.equal a b
    | None, None -> true
    | _, _ -> false in
  try
    List.for_all2 (fun (ka, va) (kb, vb) -> String.equal ka kb && equal_values va vb)
      a b
  with _ -> false

let equal a b = match a, b with
  | `Hello a, `Hello b -> Domain.equal a b
  | `Mail (a, pa), `Mail (b, pb) ->
    Reverse_path.equal a b && equal_parameters pa pb
  | `Recipient (a, pa), `Recipient (b, pb) ->
    Forward_path.equal a b && equal_parameters pa pb
  | `Expand a, `Expand b -> String.equal a b
  | `Data, `Data -> true
  | `Data_end, `Data_end -> true
  | `Help a, `Help b -> Option.equal String.equal a b
  | `Noop a, `Noop b -> Option.equal String.equal a b
  | `Verify a, `Verify b -> String.equal a b
  | `Reset, `Reset -> true
  | `Quit, `Quit -> true
  | _, _ -> false

let pp ppf = function
  | `Hello domain ->
    Fmt.pf ppf "(Hello @[<hov>%a@])" Domain.pp domain
  | `Mail (reverse_path, parameters) ->
    Fmt.pf ppf "(Mail { @[<hov>path= %a;@ parameters= %a;@] })"
      Reverse_path.pp reverse_path
      Fmt.(Dump.list (pair string (option string))) parameters
  | `Recipient (forward_path, parameters) ->
    Fmt.pf ppf "(Recipient { @[<hov>path= %a;@ parameters= %a;@] })"
      Forward_path.pp forward_path
      Fmt.(Dump.list (pair string (option string))) parameters
  | `Expand data -> Fmt.pf ppf "(Expand %s)" data
  | `Data -> Fmt.string ppf "Data"
  | `Data_end -> Fmt.string ppf "<dot>"
  | `Help data -> Fmt.pf ppf "(Help %a)" Fmt.(option string) data
  | `Noop data -> Fmt.pf ppf "(Noop %a)" Fmt.(option string) data
  | `Verify data -> Fmt.pf ppf "(Verify %s)" data
  | `Reset -> Fmt.string ppf "Reset"
  | `Quit -> Fmt.string ppf "Quit"
  | `Verb (v, a) -> Fmt.pf ppf "(Verb %s %a)" v Fmt.(Dump.list string) a
  | `Payload v -> Fmt.pf ppf "%s" v

module Decoder = struct
  open Decoder

  type nonrec error = [ `Invalid_command of string | `Invalid_domain of string | error ]

  let pp_error ppf = function
    | `Invalid_command command -> Fmt.pf ppf "Invalid command: %S" command
    | `Invalid_domain domain -> Fmt.pf ppf "Invalid domain: %S" domain
    | #Decoder.error as err -> pp_error ppf err

  (* According to RFC 5321. *)

  let trie = Hashtbl.create 16
  let () = Hashtbl.add trie "EHLO" `Hello
  let () = Hashtbl.add trie "HELO" `Hello
  let () = Hashtbl.add trie "MAIL" `Mail
  let () = Hashtbl.add trie "RCPT" `Recipient
  let () = Hashtbl.add trie "DATA" `Data
  let () = Hashtbl.add trie "." `Data_end
  let () = Hashtbl.add trie "RSET" `Reset
  let () = Hashtbl.add trie "VRFY" `Verify
  let () = Hashtbl.add trie "EXPN" `Expand
  let () = Hashtbl.add trie "HELP" `Help
  let () = Hashtbl.add trie "NOOP" `Noop
  let () = Hashtbl.add trie "QUIT" `Quit

  type command =
    [ `Hello | `Mail | `Recipient | `Data | `Data_end | `Reset | `Verify
    | `Expand | `Help | `Noop | `Quit ]

  let add_extension v = Hashtbl.add trie (String.uppercase_ascii v) `Extension
  type extension = [ `Extension ]

  let without_eol (raw, off, len) =
    if len > 1 && Bytes.unsafe_get raw (off + len - 2) = '\r'
    then Bytes.sub_string raw off (len - 2)
    else Bytes.sub_string raw off (len - 1)

  let command ?relax k decoder =
    let raw, off, len = peek_while_eol_or_space ?relax decoder in
    let data = match Bytes.unsafe_get raw (off + len - 1) with
      | ' ' -> Bytes.sub_string raw off (len - 1)
      | '\n' ->
        if len > 1
        && Bytes.unsafe_get raw (off + len - 2) = '\r'
        then Bytes.sub_string raw off (len - 2)
        else Bytes.sub_string raw off (len - 1)
      | _ -> leave_with decoder `Expected_eol_or_space in
    match Hashtbl.find trie (String.uppercase_ascii data) with
    | #extension ->
      if Bytes.unsafe_get raw (off + len - 1) = '\n'
      then ( decoder.pos <- decoder.pos + len ; k (`Only_verb data) decoder )
      else ( decoder.pos <- decoder.pos + len ; k (`Verb data) decoder )
    | #command as command ->
      decoder.pos <- decoder.pos + len ; k command decoder
    | exception Not_found ->
      let raw_eol, off, len = peek_while_eol ?relax decoder in
      let v = without_eol (raw_eol, off, len) in
      decoder.pos <- decoder.pos + len ;
      k (`Payload v) decoder

  let hello ?relax (decoder : decoder) =
    let raw_eol, off, len = peek_while_eol ?relax decoder in
    let domain = without_eol (raw_eol, off, len) in
    let domain =
      try Ok (Domain.of_string_exn domain)
      with _exn -> match Ipaddr.of_string domain with
        | Ok (Ipaddr.V4 ipv4) -> Ok (Domain.IPv4 ipv4)
        | Ok (Ipaddr.V6 ipv6) -> Ok (Domain.IPv6 ipv6)
        | Error _ -> Error (`Invalid_domain domain) in
    match domain with
    | Ok domain ->
      decoder.pos <- decoder.pos + len ;
      return (`Hello domain) decoder
    | Error err -> fail decoder err

  let mail ?relax decoder =
    let raw_eol, off, len = peek_while_eol ?relax decoder in
    let v = without_eol (raw_eol, off, len) in
    let reverse_path = Reverse_path.Decoder.of_string v in
    decoder.pos <- decoder.pos + len ; return (`Mail reverse_path) decoder

  let recipient ?relax decoder =
    let raw_eol, off, len = peek_while_eol ?relax decoder in
    let v = without_eol (raw_eol, off, len) in
    let forward_path =
      Forward_path.Decoder.of_string v in
    decoder.pos <- decoder.pos + len ; return (`Recipient forward_path) decoder

  let help ?relax decoder =
    match peek_char decoder with
    | Some _ ->
      let raw_eol, off, len = peek_while_eol ?relax decoder in
      let v = without_eol (raw_eol, off, len) in
      let v = `Help (Some v) in
      decoder.pos <- decoder.pos + len ; return v decoder
    | None -> return (`Help None) decoder

  let noop ?relax decoder =
    match peek_char decoder with
    | Some _ ->
      let raw_eol, off, len = peek_while_eol ?relax decoder in
      let v = without_eol (raw_eol, off, len) in
      let v = `Noop (Some v) in
      decoder.pos <- decoder.pos + len ; return v decoder
    | None -> return (`Noop None) decoder

  let request ?relax decoder =
    let k v decoder = match v with
      | `Hello -> hello ?relax decoder
      | `Mail ->
        string "FROM:" decoder ;
        mail ?relax decoder
      | `Recipient ->
        string "TO:" decoder ;
        recipient ?relax decoder
      | `Data -> (* assert (decoder.pos = end_of_input decoder) ; *) return `Data decoder
      | `Data_end -> (* assert (decoder.pos = end_of_input decoder) ; *) return `Data_end decoder
      | `Reset -> (* assert (decoder.pos = end_of_input decoder) ; *) return `Reset decoder
      | `Verify ->
        let raw_eol, off, len = peek_while_eol ?relax decoder in
        let v = without_eol (raw_eol, off, len) in
        let v = `Verify v in
        decoder.pos <- decoder.pos + len ; return v decoder
      | `Expand ->
        let raw_eol, off, len = peek_while_eol ?relax decoder in
        let v = without_eol (raw_eol, off, len) in
        let v = `Expand v in
        decoder.pos <- decoder.pos + len ; return v decoder
      | `Help -> help decoder
      | `Noop -> noop decoder
      | `Quit -> (* assert (decoder.pos = end_of_input decoder) ; *) return `Quit decoder
      | `Only_verb verb -> return (`Verb (verb, [])) decoder
      | `Verb verb ->
        let raw_eol, off, len = peek_while_eol ?relax decoder in
        let raw = without_eol (raw_eol, off, len) in
        let v = `Verb (verb, String.cuts ~sep:" " ~empty:true raw) in
        decoder.pos <- decoder.pos + len ; return v decoder
      | `Payload payload -> return (`Payload payload) decoder
    in command ?relax k decoder

  let request ?relax decoder =
    if at_least_one_line ?relax decoder
    then safe (request ?relax) decoder
    else prompt ?relax (request ?relax) decoder

  let of_string x =
    let decoder = decoder_from x in
    let go x : (t, [> error ]) result = match x with
      | Read _ -> Error `End_of_input
      | Error { error; _ } ->  Error error
      | Done v -> Ok v in
    go (request decoder)

  let of_string_raw x r =
    let decoder = decoder_from x in
    let go x : (t, [> error ]) result = match x with
      | Read _ -> Error `End_of_input
      | Error { error; _ } ->  Error error
      | Done v -> r := decoder.pos ; Ok v in
    go (request decoder)
end

module Encoder = struct
  open Encoder

  type nonrec error = error

  let pp_error = pp_error

  let crlf encoder = write "\r\n" encoder

  let hello domain encoder =
    write "EHLO" encoder ; (* TODO: can write HELO. *)
    write " " encoder ;
    write (Domain.to_string domain) encoder ;
    crlf encoder

  let write_parameters parameters encoder =
    let rec go = function
      | [] -> ()
      | (k, Some v) :: r ->
        write " " encoder ;
        write k encoder ;
        write "=" encoder ;
        write v encoder ;
        go r
      | (k, None) :: r ->
        write " " encoder ;
        write k encoder ;
        go r in
    go parameters

  let mail reverse_path parameters encoder =
    write "MAIL FROM:" encoder ;
    write (Reverse_path.Encoder.to_string reverse_path) encoder ;
    match parameters with
    | [] -> crlf encoder
    | parameters ->
      write_parameters parameters encoder ; crlf encoder

  let recipient forward_path parameters encoder =
    write "RCPT TO:" encoder ;
    write (Forward_path.Encoder.to_string forward_path) encoder ;
    match parameters with
    | [] -> crlf encoder
    | parameters ->
      write_parameters parameters encoder ; crlf encoder

  let verify argument encoder =
    write "VRFY " encoder ;
    write argument encoder ;
    crlf encoder

  let expand argument encoder =
    write "EXPN " encoder ;
    write argument encoder ;
    crlf encoder

  let help argument encoder =
    write "HELP" encoder ;
    match argument with
    | Some argument ->
      write " " encoder ;
      write argument encoder ;
      crlf encoder
    | None -> crlf encoder

  let noop argument encoder =
    write "NOOP" encoder ;
    match argument with
    | Some argument ->
      write " " encoder ;
      write argument encoder ;
      crlf encoder
    | None -> crlf encoder

  let request command encoder =
    match command with
    | `Hello domain -> hello domain encoder
    | `Mail (reverse_path, parameters) ->
      mail reverse_path parameters encoder
    | `Recipient (forward_path, parameters) ->
      recipient forward_path parameters encoder
    | `Data -> write "DATA" encoder ; crlf encoder
    | `Data_end -> write "." encoder ; crlf encoder
    | `Reset -> write "RSET" encoder ; crlf encoder
    | `Verify argument ->
      verify argument encoder
    | `Expand argument ->
      expand argument encoder
    | `Help argument ->
      help argument encoder
    | `Noop argument ->
      noop argument encoder
    | `Quit -> write "QUIT" encoder ; crlf encoder
    | `Verb (verb, []) ->
      write verb encoder ; crlf encoder
    | `Verb (verb, args) ->
      write verb encoder
    ; List.iter (fun arg -> write " " encoder ; write arg encoder) args
    ; crlf encoder
    | `Payload v ->
      write v encoder ; crlf encoder

  let request command encoder =
    let k encoder = request command encoder ; Done in
    match safe k encoder with
    | Write _ -> assert false (* XXX(dinosaure): request never calls [flush]. *)
    | v -> flush (fun _ -> v) encoder

  let to_string x =
    let encoder = encoder () in
    let res = Buffer.create 16 in
    let rec go x : (string, error) result = match x with
      | Write { buffer; off; len; continue } ->
        Buffer.add_substring res buffer off len ;
        go (continue len)
      | Error error -> Error error
      | Done -> Ok (Buffer.contents res) in
    go (request x encoder)
end