Source file base64_rfc2045.ml

(*
 * Copyright (c) 2018 Romain Calascibetta <romain.calascibetta@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *)

let default_alphabet =
  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

let io_buffer_size = 65536
let invalid_arg fmt = Format.ksprintf (fun s -> invalid_arg s) fmt

let invalid_bounds off len =
  invalid_arg "Invalid bounds (off: %d, len: %d)" off len

let malformed source off pos len =
  `Malformed (Bytes.sub_string source (off + pos) len)

let unsafe_byte source off pos = Bytes.unsafe_get source (off + pos)
let unsafe_blit = Bytes.unsafe_blit
let unsafe_chr = Char.unsafe_chr
let unsafe_set_chr source off chr = Bytes.unsafe_set source off chr

type state = {quantum: int; size: int; buffer: Bytes.t}

let continue state (quantum, size) = `Continue {state with quantum; size}
let flush state = `Flush {state with quantum= 0; size= 0}

let table =
  "\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\255\062\255\255\255\063\052\053\054\055\056\057\058\059\060\061\255\255\255\255\255\255\255\000\001\002\003\004\005\006\007\008\009\010\011\012\013\014\015\016\017\018\019\020\021\022\023\024\025\255\255\255\255\255\255\026\027\028\029\030\031\032\033\034\035\036\037\038\039\040\041\042\043\044\045\046\047\048\049\050\051\255\255\255\255\255"

let r_repr ({quantum; size; _} as state) chr =
  (* assert (0 <= off && 0 <= len && off + len <= String.length source); *)
  (* assert (len >= 1); *)
  let code = Char.code table.[Char.code chr] in
  match size with
  | 0 -> continue state (code, 1)
  | 1 -> continue state ((quantum lsl 6) lor code, 2)
  | 2 -> continue state ((quantum lsl 6) lor code, 3)
  | 3 ->
      unsafe_set_chr state.buffer 0 (unsafe_chr ((quantum lsr 10) land 255)) ;
      unsafe_set_chr state.buffer 1 (unsafe_chr ((quantum lsr 2) land 255)) ;
      unsafe_set_chr state.buffer 2
        (unsafe_chr ((quantum lsl 6) lor code land 255)) ;
      flush state
  | _ -> malformed (Bytes.make 1 chr) 0 0 1

let r_crlf source off len =
  (* assert (0 <= off && 0 <= len && off + len <= String.length source); *)
  (* assert (len = 2); *)
  match Bytes.sub_string source off len with
  | "\r\n" -> `Line_break
  | _ -> malformed source off 0 len

type src = [`Channel of in_channel | `String of string | `Manual]

type decode =
  [`Await | `End | `Wrong_padding | `Malformed of string | `Flush of string]

type input =
  [`Line_break | `Wsp | `Padding | `Malformed of string | `Flush of state]

type decoder =
  { src: src
  ; mutable i: Bytes.t
  ; mutable i_off: int
  ; mutable i_pos: int
  ; mutable i_len: int
  ; mutable s: state
  ; h: Bytes.t
  ; mutable h_len: int
  ; mutable h_need: int
  ; mutable padding: int
  ; mutable unsafe: bool
  ; mutable byte_count: int
  ; mutable limit_count: int
  ; mutable pp: decoder -> input -> decode
  ; mutable k: decoder -> decode }

let i_rem decoder = decoder.i_len - decoder.i_pos + 1

let end_of_input decoder =
  decoder.i <- Bytes.empty ;
  decoder.i_off <- 0 ;
  decoder.i_pos <- 0 ;
  decoder.i_len <- min_int

let src decoder source off len =
  if off < 0 || len < 0 || off + len > Bytes.length source then
    invalid_bounds off len
  else if len = 0 then end_of_input decoder
  else (
    decoder.i <- source ;
    decoder.i_off <- off ;
    decoder.i_pos <- 0 ;
    decoder.i_len <- len - 1 )

let refill k decoder =
  match decoder.src with
  | `Manual ->
      decoder.k <- k ;
      `Await
  | `String _ -> end_of_input decoder ; k decoder
  | `Channel ic ->
      let len = input ic decoder.i 0 (Bytes.length decoder.i) in
      src decoder decoder.i 0 len ;
      k decoder

let dangerous decoder v = decoder.unsafe <- v
let reset decoder = decoder.limit_count <- 0

let ret k v byte_count decoder =
  decoder.k <- k ;
  decoder.byte_count <- decoder.byte_count + byte_count ;
  decoder.limit_count <- decoder.limit_count + byte_count ;
  if decoder.limit_count > 78 then dangerous decoder true ;
  decoder.pp decoder v

[@@@warning "-32"]

let is_b64 = function
  | 'A' .. 'Z' | 'a' .. 'z' | '0' .. '9' | '+' | '/' -> true
  | _ -> false

let t_need decoder need =
  decoder.h_len <- 0 ;
  decoder.h_need <- need

let rec t_fill k decoder =
  let blit decoder len =
    unsafe_blit decoder.i
      (decoder.i_off + decoder.i_pos)
      decoder.h decoder.h_len len ;
    decoder.i_pos <- decoder.i_pos + len ;
    decoder.h_len <- decoder.h_len + len
  in
  let rem = i_rem decoder in
  if rem < 0 (* end of input *) then k decoder
  else
    let need = decoder.h_need - decoder.h_len in
    if rem < need then (
      blit decoder rem ;
      refill (t_fill k) decoder )
    else ( blit decoder need ; k decoder )

type flush_and_malformed = [`Flush of state | `Malformed of string]

let padding {size; _} padding =
  match (size, padding) with
  | 0, 0 -> true
  | 1, _ -> false
  | 2, 2 -> true
  | 3, 1 -> true
  | _ -> false

let rec t_crlf decoder =
  if decoder.h_len < decoder.h_need then
    ret decode_base64
      (malformed decoder.h 0 0 decoder.h_len)
      decoder.h_len decoder
  else
    ret decode_base64 (r_crlf decoder.h 0 decoder.h_len) decoder.h_len decoder

and t_flush {quantum; size; buffer} =
  match size with
  | 0 | 1 -> `Flush {quantum; size; buffer= Bytes.empty}
  | 2 ->
      let quantum = quantum lsr 4 in
      `Flush
        { quantum
        ; size
        ; buffer= Bytes.make 1 (unsafe_chr (quantum land 255)) }
  | 3 ->
      let quantum = quantum lsr 2 in
      unsafe_set_chr buffer 0 (unsafe_chr ((quantum lsr 8) land 255)) ;
      unsafe_set_chr buffer 1 (unsafe_chr (quantum land 255)) ;
      `Flush {quantum; size; buffer= Bytes.sub buffer 0 2}
  | _ -> malformed buffer 0 0 3

and t_decode_base64 chr decoder =
  if decoder.padding = 0 then
    let rec go pos = function
      | `Continue state ->
          if decoder.i_len - (decoder.i_pos + pos) + 1 > 0
          then (
            match unsafe_byte decoder.i decoder.i_off (decoder.i_pos + pos) with
            | ('A' .. 'Z' | 'a' .. 'z' | '0' .. '9' | '+' | '/') as chr -> go (succ pos) (r_repr state chr)
            | '=' ->
                decoder.i_pos <- decoder.i_pos + pos ;
                decoder.s <- state ;
                ret decode_base64 `Padding pos decoder
            | ' ' | '\t' ->
                decoder.i_pos <- decoder.i_pos + pos ;
                decoder.s <- state ;
                ret decode_base64 `Wsp pos decoder
            | '\r' ->
                decoder.i_pos <- decoder.i_pos + pos ;
                decoder.s <- state ;
                t_need decoder 2 ;
                t_fill t_crlf decoder
            | chr -> malformed (Bytes.make 1 chr) 0 0 1
          ) else (
            decoder.i_pos <- decoder.i_pos + pos ;
            decoder.byte_count <- decoder.byte_count + pos ;
            decoder.s <- state ;
            refill decode_base64 decoder )
      | #flush_and_malformed as v ->
          decoder.i_pos <- decoder.i_pos + pos ;
          ret decode_base64 v pos decoder
    in
    go 1 (r_repr decoder.s chr)
  else malformed (Bytes.make 1 chr) 0 0 1

and decode_base64 decoder =
  let rem = i_rem decoder in
  if rem <= 0 then
    if rem < 0 then
      ret
        (fun decoder ->
          if padding decoder.s decoder.padding then `End else `Wrong_padding )
        (t_flush decoder.s) 0 decoder
    else refill decode_base64 decoder
  else
    match unsafe_byte decoder.i decoder.i_off decoder.i_pos with
    | ('A' .. 'Z' | 'a' .. 'z' | '0' .. '9' | '+' | '/') as chr ->
        t_decode_base64 chr decoder
    | '=' ->
        decoder.padding <- decoder.padding + 1 ;
        decoder.i_pos <- decoder.i_pos + 1 ;
        ret decode_base64 `Padding 1 decoder
    | ' ' | '\t' ->
        decoder.i_pos <- decoder.i_pos + 1 ;
        ret decode_base64 `Wsp 1 decoder
    | '\r' -> t_need decoder 2 ; t_fill t_crlf decoder
    | chr -> malformed (Bytes.make 1 chr) 0 0 1

let pp_base64 decoder = function
  | `Line_break -> reset decoder ; decoder.k decoder
  | `Wsp | `Padding -> decoder.k decoder
  | `Flush state ->
      decoder.s <- state ;
      `Flush (Bytes.to_string state.buffer)
  | `Malformed _ as v -> v

let decoder src =
  let pp = pp_base64 in
  let k = decode_base64 in
  let i, i_off, i_pos, i_len =
    match src with
    | `Manual -> (Bytes.empty, 0, 1, 0)
    | `Channel _ -> (Bytes.create io_buffer_size, 0, 1, 0)
    | `String s -> (Bytes.unsafe_of_string s, 0, 0, String.length s - 1)
  in
  { src
  ; i_off
  ; i_pos
  ; i_len
  ; i
  ; s= {quantum= 0; size= 0; buffer= Bytes.create 3}
  ; h= Bytes.create 2
  ; h_len= 0
  ; h_need= 0
  ; padding= 0
  ; unsafe= false
  ; byte_count= 0
  ; limit_count= 0
  ; pp
  ; k }

let decode decoder = decoder.k decoder
let decoder_byte_count decoder = decoder.byte_count
let decoder_src decoder = decoder.src
let decoder_dangerous decoder = decoder.unsafe

(* / *)

let invalid_encode () = invalid_arg "Expected `Await encode"

type dst = [`Channel of out_channel | `Buffer of Buffer.t | `Manual]
type encode = [`Await | `End | `Char of char]

type encoder =
  { dst: dst
  ; mutable o: Bytes.t
  ; mutable o_off: int
  ; mutable o_pos: int
  ; mutable o_len: int
  ; mutable c_col: int
  ; i: Bytes.t
  ; mutable s: int
  ; t: Bytes.t
  ; mutable t_pos: int
  ; mutable t_len: int
  ; mutable k: encoder -> encode -> [`Ok | `Partial] }

let o_rem encoder = encoder.o_len - encoder.o_pos + 1

let dst encoder source off len =
  if off < 0 || len < 0 || off + len > Bytes.length source then
    invalid_bounds off len ;
  encoder.o <- source ;
  encoder.o_off <- off ;
  encoder.o_pos <- 0 ;
  encoder.o_len <- len - 1

let dst_rem = o_rem

let partial k encoder = function
  | `Await -> k encoder
  | `Char _ | `End -> invalid_encode ()

let flush k encoder =
  match encoder.dst with
  | `Manual ->
      encoder.k <- partial k ;
      `Partial
  | `Channel oc ->
      output oc encoder.o encoder.o_off encoder.o_pos ;
      encoder.o_pos <- 0 ;
      k encoder
  | `Buffer b ->
      let o = Bytes.unsafe_to_string encoder.o in
      Buffer.add_substring b o encoder.o_off encoder.o_pos ;
      encoder.o_pos <- 0 ;
      k encoder

let t_range encoder len =
  encoder.t_pos <- 0 ;
  encoder.t_len <- len

let rec t_flush k encoder =
  let blit encoder len =
    unsafe_blit encoder.t encoder.t_pos encoder.o encoder.o_pos len ;
    encoder.o_pos <- encoder.o_pos + len ;
    encoder.t_pos <- encoder.t_pos + len
  in
  let rem = o_rem encoder in
  let len = encoder.t_len - encoder.t_pos + 1 in
  if rem < len then (
    blit encoder rem ;
    flush (t_flush k) encoder )
  else ( blit encoder len ; k encoder )

let rec encode_line_break k encoder =
  let rem = o_rem encoder in
  let s, j, k =
    if rem < 2 then (
      t_range encoder 2 ;
      (encoder.t, 0, t_flush k) )
    else
      let j = encoder.o_pos in
      encoder.o_pos <- encoder.o_pos + 2 ;
      (encoder.o, encoder.o_off + j, k)
  in
  unsafe_set_chr s j '\r' ;
  unsafe_set_chr s (j + 1) '\n' ;
  encoder.c_col <- 0 ;
  k encoder

and encode_char chr k (encoder : encoder) =
  if encoder.s >= 2 then (
    let a, b, c =
      (unsafe_byte encoder.i 0 0, unsafe_byte encoder.i 0 1, chr)
    in
    encoder.s <- 0 ;
    let quantum = (Char.code a lsl 16) + (Char.code b lsl 8) + Char.code c in
    let a = quantum lsr 18 in
    let b = (quantum lsr 12) land 63 in
    let c = (quantum lsr 6) land 63 in
    let d = quantum land 63 in
    let rem = o_rem encoder in
    let s, j, k =
      if rem < 4 then (
        t_range encoder 4 ;
        (encoder.t, 0, t_flush (k 4)) )
      else
        let j = encoder.o_pos in
        encoder.o_pos <- encoder.o_pos + 4 ;
        (encoder.o, encoder.o_off + j, k 4)
    in
    unsafe_set_chr s j default_alphabet.[a] ;
    unsafe_set_chr s (j + 1) default_alphabet.[b] ;
    unsafe_set_chr s (j + 2) default_alphabet.[c] ;
    unsafe_set_chr s (j + 3) default_alphabet.[d] ;
    flush k encoder )
  else (
    unsafe_set_chr encoder.i encoder.s chr ;
    encoder.s <- encoder.s + 1 ;
    k 0 encoder )

and encode_trailing k encoder =
  match encoder.s with
  | 2 ->
      let b, c = (unsafe_byte encoder.i 0 0, unsafe_byte encoder.i 0 1) in
      encoder.s <- 0 ;
      let quantum = (Char.code b lsl 10) + (Char.code c lsl 2) in
      let b = (quantum lsr 12) land 63 in
      let c = (quantum lsr 6) land 63 in
      let d = quantum land 63 in
      let rem = o_rem encoder in
      let s, j, k =
        if rem < 4 then (
          t_range encoder 4 ;
          (encoder.t, 0, t_flush (k 4)) )
        else
          let j = encoder.o_pos in
          encoder.o_pos <- encoder.o_pos + 4 ;
          (encoder.o, encoder.o_off + j, k 4)
      in
      unsafe_set_chr s j default_alphabet.[b] ;
      unsafe_set_chr s (j + 1) default_alphabet.[c] ;
      unsafe_set_chr s (j + 2) default_alphabet.[d] ;
      unsafe_set_chr s (j + 3) '=' ;
      flush k encoder
  | 1 ->
      let c = unsafe_byte encoder.i 0 0 in
      encoder.s <- 0 ;
      let quantum = Char.code c lsl 4 in
      let c = (quantum lsr 6) land 63 in
      let d = quantum land 63 in
      let rem = o_rem encoder in
      let s, j, k =
        if rem < 4 then (
          t_range encoder 4 ;
          (encoder.t, 0, t_flush (k 4)) )
        else
          let j = encoder.o_pos in
          encoder.o_pos <- encoder.o_pos + 4 ;
          (encoder.o, encoder.o_off + j, k 4)
      in
      unsafe_set_chr s j default_alphabet.[c] ;
      unsafe_set_chr s (j + 1) default_alphabet.[d] ;
      unsafe_set_chr s (j + 2) '=' ;
      unsafe_set_chr s (j + 3) '=' ;
      flush k encoder
  | 0 -> k 0 encoder
  | _ -> assert false

and encode_base64 encoder v =
  let k col_count encoder =
    encoder.c_col <- encoder.c_col + col_count ;
    encoder.k <- encode_base64 ;
    `Ok
  in
  match v with
  | `Await -> k 0 encoder
  | `End ->
      if encoder.c_col = 76 then
        encode_line_break (fun encoder -> encode_base64 encoder v) encoder
      else encode_trailing k encoder
  | `Char chr ->
      let rem = o_rem encoder in
      if rem < 1 then flush (fun encoder -> encode_base64 encoder v) encoder
      else if encoder.c_col = 76 then
        encode_line_break (fun encoder -> encode_base64 encoder v) encoder
      else encode_char chr k encoder

let encoder dst =
  let o, o_off, o_pos, o_len =
    match dst with
    | `Manual -> (Bytes.empty, 1, 0, 0)
    | `Buffer _ | `Channel _ ->
        (Bytes.create io_buffer_size, 0, 0, io_buffer_size - 1)
  in
  { dst
  ; o_off
  ; o_pos
  ; o_len
  ; o
  ; t= Bytes.create 4
  ; t_pos= 1
  ; t_len= 0
  ; c_col= 0
  ; i= Bytes.create 3
  ; s= 0
  ; k= encode_base64 }

let encode encoder = encoder.k encoder
let encoder_dst encoder = encoder.dst