Source file decoder.ml

type decoder = { buffer : Bytes.t; mutable pos : int; mutable max : int }

let pp ppf { buffer; pos; max } =
  Fmt.pf ppf "%S" (Bytes.sub_string buffer pos (max - pos))

let io_buffer_size = 65536
let decoder () = { buffer = Bytes.create io_buffer_size; pos = 0; max = 0 }

let decoder_from_string x =
  let max = String.length x in
  let buffer = Bytes.of_string x in
  { buffer; pos = 0; max }

let decoder_from_preallocated_bytes buffer = { buffer; pos = 0; max = 0 }

type error =
  [ `End_of_input
  | `Expected_char of char
  | `Unexpected_char of char
  | `Expected_string of string
  | (* | Invalid_command of string *)
    `Expected_eol
  | `Expected_eol_or_space
  | `Not_enough_space
  | `Assert_predicate of char -> bool ]
(* | Invalid_code of int *)

let pp_error ppf = function
  | `End_of_input -> Fmt.string ppf "End of input"
  | `Expected_char chr -> Fmt.pf ppf "Expected char: %02x" (Char.code chr)
  | `Unexpected_char chr -> Fmt.pf ppf "Unexpected char: %02x" (Char.code chr)
  | `Expected_string s -> Fmt.pf ppf "Expected_string: %s" s
  | `Expected_eol -> Fmt.string ppf "Expected end-of-line"
  | `Expected_eol_or_space -> Fmt.string ppf "Expected end-of-line or space"
  | `Not_enough_space -> Fmt.string ppf "Not enough space"
  | `Assert_predicate _ -> Fmt.string ppf "Assert predicate"

type 'err info = { error : 'err; buffer : Bytes.t; committed : int }

type ('v, 'err) state =
  | Done of 'v
  | Read of {
      buffer : Bytes.t;
      off : int;
      len : int;
      continue : [ `End | `Len of int ] -> ('v, 'err) state;
    }
  | Error of 'err info

exception Leave of error info

let return (type v) (v : v) _ : (v, 'err) state = Done v

let safe :
    (decoder -> ('v, ([> error ] as 'err)) state) -> decoder -> ('v, 'err) state
    =
 fun k decoder ->
  try k decoder
  with Leave { error = #error as error; buffer; committed } ->
    Error { error = (error :> 'err); buffer; committed }

let end_of_input decoder = decoder.max

let peek_char decoder =
  if decoder.pos < end_of_input decoder
  then Some (Bytes.unsafe_get decoder.buffer decoder.pos)
  else None
(* XXX(dinosaure): in [angstrom] world, [peek_char] should try to read input
    again. However, SMTP is a line-directed protocol where we can ensure to
    have the full line at the top (with a queue) instead to have a
    systematic check (which slow-down the process). *)

let leave_with (decoder : decoder) error =
  raise (Leave { error; buffer = decoder.buffer; committed = decoder.pos })

let fail (decoder : decoder) error =
  Error { error; buffer = decoder.buffer; committed = decoder.pos }

let string str decoder =
  let idx = ref 0 in
  let len = String.length str in
  while
    decoder.pos + !idx < end_of_input decoder
    && !idx < len
    && Char.equal
         (Bytes.unsafe_get decoder.buffer (decoder.pos + !idx))
         (String.unsafe_get str !idx)
  do
    incr idx
  done ;
  if !idx = len
  then decoder.pos <- decoder.pos + len
  else leave_with decoder (`Expected_string str)

let junk_char decoder =
  if decoder.pos < end_of_input decoder
  then decoder.pos <- decoder.pos + 1
  else leave_with decoder `End_of_input

let while1 predicate decoder =
  let idx = ref decoder.pos in
  while
    !idx < end_of_input decoder
    && predicate (Bytes.unsafe_get decoder.buffer !idx)
  do
    incr idx
  done ;
  if !idx - decoder.pos = 0
  then leave_with decoder (`Assert_predicate predicate) ;
  let sub = (decoder.buffer, decoder.pos, !idx - decoder.pos) in
  (* XXX(dinosaure): avoid sub-string operation. *)
  decoder.pos <- !idx ;
  sub

let at_least_one_line ?(relax = true) decoder =
  let pos = ref decoder.pos in
  let chr = ref '\000' in
  let has_cr = ref false in
  while
    !pos < decoder.max
    &&
    (chr := Bytes.unsafe_get decoder.buffer !pos ;
     not (!chr = '\n' && (!has_cr || relax)))
  do
    has_cr := !chr = '\r' ;
    incr pos
  done ;
  !pos < decoder.max && !chr = '\n' && (!has_cr || relax)

let prompt :
    ?relax:bool ->
    (decoder -> ('v, ([> error ] as 'err)) state) ->
    decoder ->
    ('v, 'err) state =
 fun ?relax k decoder ->
  if decoder.pos > 0
  then (
    (* XXX(dinosaure): compress *)
    let rest = decoder.max - decoder.pos in
    Bytes.unsafe_blit decoder.buffer decoder.pos decoder.buffer 0 rest ;
    decoder.max <- rest ;
    decoder.pos <- 0) ;
  let rec go off =
    if off = Bytes.length decoder.buffer
    then
      Error
        {
          error = `Not_enough_space;
          buffer = decoder.buffer;
          committed = decoder.pos;
        }
    else if not (at_least_one_line ?relax { decoder with max = off })
            (* XXX(dinosaure): we make a new decoder here and we did __not__ set [decoder.max] owned by end-user,
               and this is exactly what we want. *)
    then
      let continue = function
        | `Len len -> go (off + len)
        | `End ->
            Error
              {
                error = `End_of_input;
                buffer = decoder.buffer;
                committed = decoder.pos;
              } in
      Read
        {
          buffer = decoder.buffer;
          off;
          len = Bytes.length decoder.buffer - off;
          continue;
        }
    else (
      decoder.max <- off ;
      safe k decoder) in
  go decoder.max

let peek_while_eol ?(relax = false) decoder =
  let idx = ref decoder.pos in
  let chr = ref '\000' in
  let has_cr = ref false in

  while
    !idx < end_of_input decoder
    &&
    (chr := Bytes.unsafe_get decoder.buffer !idx ;
     not (!chr == '\n' && (!has_cr || relax)))
  do
    has_cr := !chr == '\r' ;
    incr idx
  done ;

  if !idx < end_of_input decoder && !chr == '\n' && (!has_cr || relax)
  then (
    assert (!idx + 1 - decoder.pos > 0) ;
    (decoder.buffer, decoder.pos, !idx + 1 - decoder.pos))
  else leave_with decoder `Expected_eol

let peek_while_eol_or_space ?(relax = false) decoder =
  let idx = ref decoder.pos in
  let chr = ref '\000' in
  let has_cr = ref false in

  while
    !idx < end_of_input decoder
    &&
    (chr := Bytes.unsafe_get decoder.buffer !idx ;
     (not (!chr = '\n' && (!has_cr || relax))) && !chr <> ' ')
  do
    has_cr := !chr = '\r' ;
    incr idx
  done ;

  if !idx < end_of_input decoder
     && ((!chr = '\n' && (!has_cr || relax)) || !chr = ' ')
  then (decoder.buffer, decoder.pos, !idx + 1 - decoder.pos)
  else leave_with decoder `Expected_eol_or_space