Source

msgpack-ocaml / msgpack.ml

Full commit
(*pp camlp4o $(ocamlfind query -i-format bitstring) bitstring.cma bitstring_persistent.cma pa_bitstring.cmo *)

module StdBuffer = Buffer
module Bitstring = struct
  include Bitstring

  let extract_char_signed data off len flen = 
    let res = extract_char_unsigned data off len flen in
    let sign = res lsr (flen - 1) in
    if sign = 0 then res else res - (1 lsl flen)

  let extract_int_be_signed data off len flen =
    let res = extract_int_be_unsigned data off len flen in
    let sign = res lsr (flen - 1) in
    if sign = 0 then res else res - (1 lsl flen)

  let extract_int32_be_signed = extract_int32_be_unsigned
  let extract_int64_be_signed data off len flen =
    let res = extract_int64_be_unsigned data off len flen in
    if flen = 64 then res
    else
      let sign = Int64.shift_right res (flen - 1) in
      if sign = 0L then res else Int64.sub res (Int64.shift_left 1L flen)
  
end

open Bitstring

type t = 
  | Nil
  | Bool of bool
  | Int of int  (* for ease of development *)
  | Int64 of Int64.t
  | UInt64 of Int64.t
  | Float of float
  | Raw of string
  | Array of t array
  | Map of (t, t) Hashtbl.t

type rettype =
    Success of t * bitstring

let recover_sign_8 x = if x >= 0x80 then x - 0x100 else x
let recover_sign_16 x = if x >= 0x8000 then x - 0x10000 else x
let recover_sign_32 x = 
  if Int32.compare x 0x80000000l >= 0 
  then Int64.sub (Int64.of_int32 x) 0x100000000L
  else Int64.of_int32 x

let rec unpack buf = 
  bitmatch buf with
    | { false: 1; x: 7; rest: -1: bitstring } (* Positive FixNum; false = 0 *)
      -> Success (Int64 (Int64.of_int x), rest)
    | { 0b111: 3; x: 5; rest: -1: bitstring } (* Negative FixNum *)
      -> Success (Int64 (Int64.of_int (-x)), rest)
    | { 0xc0: 8; rest: -1: bitstring } (* Nil *)
      -> Success (Nil, rest)
    | { 0xc2: 8; rest: -1: bitstring } (* false *)
      -> Success (Bool false, rest)
    | { 0xc3: 8; rest: -1: bitstring } (* true *)
      -> Success (Bool true, rest)
    | { 0xca: 8; payload: 32: bigendian; rest: -1: bitstring } (* float32 *)
      -> Success (Float (Int32.float_of_bits payload), rest)
    | { 0xcb: 8; payload: 64: bigendian; rest: -1: bitstring } (* float64 *)
      -> Success (Float (Int64.float_of_bits payload), rest)
    | { 0xcc: 8; x: 8; rest: -1: bitstring } (* unsigned 8-bit *)
      -> Success (UInt64 (Int64.of_int x), rest)
    | { 0xcd: 8; x: 16: bigendian; rest: -1: bitstring } (* unsigned 16-bit *)
      -> Success (UInt64 (Int64.of_int x), rest)
    | { 0xce: 8; x: 32: bigendian; rest: -1: bitstring } (* unsigned 32-bit *)
      -> Success (UInt64 (Int64.of_int32 x), rest)
    | { 0xcf: 8; x: 64: bigendian; rest: -1: bitstring } (* unsigned 64-bit *)
      -> Success (UInt64 x, rest)
    | { 0xd0: 8; x: 8: signed; rest: -1: bitstring } (* signed 8-bit *)
      -> Success (Int64 (Int64.of_int x), rest)
    | { 0xd1: 8; x: 16: bigendian, signed; rest: -1: bitstring } (* signed 16-bit *)
      -> Success (Int64 (Int64.of_int x), rest)
    | { 0xd2: 8; x: 32: bigendian, signed; rest: -1: bitstring } (* signed 32-bit *)
      -> Success (Int64 (Int64.of_int32 x), rest)
    | { 0xd3: 8; x: 64: bigendian, signed; rest: -1: bitstring } (* signed 64-bit *)
      -> Success (Int64 x, rest)
    | { 0b101: 3; n: 5; payload: n*8: string; rest: -1: bitstring } (* Fix raw *)
      -> Success (Raw payload, rest)
    | { 0xda: 8; n: 16: bigendian; payload: n*8: string; rest: -1: bitstring } (* 16-bit raw *)
      -> Success (Raw payload, rest)
    | { 0xdb: 8; n: 32: bigendian;
	payload: (Int32.to_int n * 8): string; rest: -1: bitstring } (* 32-bit raw *)
      -> Success (Raw payload, rest) (* warning: potential DOS risk *)
    | { 0b1001: 4; n: 4; payload: -1: bitstring } (* fix array *)
      -> unpack_array n payload
    | { 0xdc: 8; n: 16: bigendian; payload: -1: bitstring } (* array 16-bit *)
      -> unpack_array n payload
    | { 0xdd: 8; n: 32: bigendian; payload: -1: bitstring } (* array 32-bit *)
      -> unpack_array (Int32.to_int n) payload
    | { 0b1000: 4; n: 4; payload: -1: bitstring } (* Fix map *)
      -> unpack_map n payload
    | { 0xde: 8; n: 16: bigendian; payload: -1: bitstring } (* 16-bit map *)
      -> unpack_map n payload
    | { 0xdf: 8; n: 32: bigendian; payload: -1: bitstring } (* 32-bit map *)
      -> unpack_map (Int32.to_int n) payload
    | { c: 8 }
      -> failwith (Printf.sprintf "Msgpack: unknown protocol (code: 0x%2X)" c)
and unpack_array n payload = (* fixme: cannot perform buffered I/O *)
  let ar = Array.make n Nil in
  let rec iter i payload = 
    if i = n 
    then Success (Array ar, payload)
    else
      match unpack payload with
	  Success (v, payload) ->
	    ar.(i) <- v;
	    iter (i + 1) payload
  in
  iter 0 payload
and unpack_map n payload = 
  let m = Hashtbl.create n in
  let rec iter i payload = 
    if i = n then Success (Map m, payload)
    else 
      match unpack payload with
	  Success (k, payload) ->
	    begin
	      match unpack payload with
		  Success (v, payload) ->
		    Hashtbl.add m k v;
		    iter (i + 1) payload
	    end
  in
  iter 0 payload

open StdBuffer

let (@+=) buf (s, _, _) = add_string buf s

let rec pack e buf = 
  match e with
    | Nil -> add_char buf '\xc0'
    | Bool false -> add_char buf '\xc2'
    | Bool true -> add_char buf '\xc3'
    | Int x when x >= 0 && x < 128 -> 
	add_char buf (char_of_int x)
    | Int x when x < 0 && x >= -32 ->
	add_char buf (char_of_int (256 + x))
    | Int x when x < 0 && x >= -128 ->
	add_char buf '\xd0';
	add_char buf (char_of_int (256 + x))
    | Int x when x < 0x8000 && x >= -0x8000 ->
	buf @+= BITSTRING { 0xd1: 8; x: 16: bigendian }
    | Int x ->
	let lx = Int64.of_int x in
	if Int64.compare lx 0x80000000L < 0 && Int64.compare lx (-0x80000000L) >= 0 then
	  buf @+= BITSTRING { 0xd2: 8; Int32.of_int x: 32: bigendian }
	else
	  buf @+= BITSTRING { 0xd3: 8; lx: 64: bigendian }
    | _ -> failwith ""
(*
    | Bool of bool
  | Int of int  (* for ease of development *)
  | Int64 of Int64.t
  | UInt64 of Int64.t
  | Float of float
  | Raw of string
  | Array of t array
  | Map of (t, t) Hashtbl.t

*)