module Reader : sig
type err =
[ `Unknown_type of char
| `Length
| `Needed_bytes of int
| `Invalid_frame
]
type 'a t
val run : buf:Bytes.t -> pos:int -> len:int -> 'a t -> ('a * int, err * int) result
val ( >>= ) : 'a t -> ('a -> 'b t) -> 'b t
val return : 'a -> 'a t
val fail : err -> 'a t
val int32 : int32 t
val bytes : int -> string t
val string : string t
val string_list : unit -> string list t
val int8 : int t
val int16 : int t
val repeat : int -> 'a t -> 'a list t
val consume : int -> 'a t -> 'a list t
end = struct
module State = struct
type t = {
buf : Bytes.t;
pos : int;
stop : int;
}
end
type err =
[ `Unknown_type of char
| `Length
| `Needed_bytes of int
| `Invalid_frame
]
type 'a t = State.t -> ('a, err) result * State.t
let ( >>= ) t f st =
match t st with
| Ok v, st -> f v st
| (Error _, _) as r -> r
let return v st = (Ok v, st)
let fail err st = (Error err, st)
let run ~buf ~pos ~len t =
let st = State.{ buf; pos; stop = pos + len } in
assert (st.State.stop Error (err, st.State.pos)
| Ok v, st -> Ok (v, st.State.pos)
let int32 st =
if st.State.pos + 4 = 0);
if st.State.pos + n
let s = Bytes.sub_string st.State.buf st.State.pos (idx - st.State.pos) in
(* +1 to consume the end null byte *)
(Ok s, { st with State.pos = st.State.pos + (idx - st.State.pos) + 1 })
| _ -> (Error `Length, st)
let rec string_list () =
string
>>= function
| "" -> return []
| auth_mechanism -> string_list () >>= fun rest -> return (auth_mechanism :: rest)
let int8 st =
if st.State.pos + 1 (
match repeat (n - 1) t st with
| Ok vs, st -> (Ok (v :: vs), st)
| (Error _ as err), st -> (err, st))
| (Error _ as err), st -> (err, st)
else (Ok [], st)
let rec consume len t st =
if len > 0 then
let pos = st.State.pos in
match t st with
| Ok v, st -> (
let consumed = st.State.pos - pos in
let len = len - consumed in
match consume len t st with
| Ok vs, st -> (Ok (v :: vs), st)
| (Error _ as err), st -> (err, st))
| (Error _ as err), st -> (err, st)
else (Ok [], st)
end
module Writer : sig
type t
val run : Buffer.t -> t -> unit
val msg_code : char -> t
val int32 : int32 -> t
val int16 : int -> t
val bytes : string -> t
val string : string -> t
val chr : char -> t
val iter : ('a -> t) -> 'a list -> t
val ( >>= ) : t -> (unit -> t) -> t
end = struct
type t = Buffer.t * Buffer.t -> unit
let msg_code ch (b, _) = Buffer.add_char b ch
let int32 n (_, buf) =
let b = Bytes.create 4 in
EndianBytes.BigEndian.set_int32 b 0 n;
Buffer.add_bytes buf b
let int16 n (_, buf) =
let b = Bytes.create 2 in
EndianBytes.BigEndian.set_int16 b 0 n;
Buffer.add_bytes buf b
let bytes s (_, buf) = Buffer.add_string buf s
let string s (_, buf) =
Buffer.add_string buf s;
Buffer.add_char buf '\000'
let chr ch (_, buf) = Buffer.add_char buf ch
let iter f xs bs = List.iter (fun x -> f x bs) xs
let run prepend_buf t =
let buf = Buffer.create 1024 in
t (prepend_buf, buf);
let len = 4 + Buffer.length buf in
int32 (Int32.of_int len) (prepend_buf, prepend_buf);
Buffer.add_buffer prepend_buf buf
let ( >>= ) t f buf =
let r = t buf in
f r buf
end
module Frame = struct
module Backend = struct
type row = {
name : string;
table_id : int32;
column_attr : int;
data_type_id : int32;
data_type_size : int;
data_type_mod : int32;
format_code : int;
}
[@@deriving show, eq]
type t =
| AuthenticationOk
| AuthenticationKerberosV5
| AuthenticationCleartextPassword
| AuthenticationMD5Password of { salt : string }
| AuthenticationSCMCredential
| AuthenticationGSS
| AuthenticationSSPI
| AuthenticationGSSContinue of { data : string }
| AuthenticationSASL of { auth_mechanisms : string list }
| AuthenticationSASLContinue of { data : string }
| AuthenticationSASLFinal of { data : string }
| BackendKeyData of {
pid : int32;
secret_key : int32;
}
| BindComplete
| CloseComplete
| CommandComplete of { tag : string }
| CopyData of { data : string }
| CopyDone
| CopyInResponse of {
format : int;
column_formats : int list;
}
| CopyOutResponse of {
format : int;
column_formats : int list;
}
| CopyBothResponse of {
format : int;
column_formats : int list;
}
| DataRow of { data : string option list }
| EmptyQueryResponse
| ErrorResponse of { msgs : (char * string) list }
| FunctionCallResponse of { result : string option }
| NegotiateProtocolVersion of {
minor_version : int32;
unrecognized_options : string list;
}
| NoData
| NoticeResponse of { msgs : (char * string) list }
| NotificationResponse of {
pid : int32;
channel : string;
payload : string;
}
| ParameterDescription of { object_ids : int32 list }
| ParameterStatus of {
name : string;
value : string;
}
| ParseComplete
| PortalSuspended
| ReadyForQuery of { status : char }
| RowDescription of { rows : row list }
[@@deriving show, eq]
end
module Frontend = struct
type t =
| Bind of {
portal : string;
stmt : string;
format_codes : bool list;
values : string option list;
result_format_codes : bool list;
}
| CancelRequest of {
pid : int32;
secret_key : int32;
}
| Close of {
typ : char;
name : string;
}
| CopyData of { data : string }
| CopyDone
| CopyFail of { message : string }
| Describe of {
typ : char;
name : string;
}
| Execute of {
portal : string;
max_rows : int32;
}
| Flush
| FunctionCall (* TODO *)
| GSSResponse of { data : string }
| Parse of {
stmt : string;
query : string;
data_types : int32 list;
}
| PasswordMessage of { password : string }
| Query of { query : string }
| SASLInitialResponse of {
auth_mechanism : string;
data : string;
}
| SASLResponse of { data : string }
| SSLRequest
| StartupMessage of { msgs : (string * string) list }
| Sync
| Terminate
[@@deriving show, eq]
end
end
module Decode = struct
type err =
[ `Unknown_type of char
| `Invalid_frame
]
[@@deriving show, eq]
type t = {
buf : Buffer.t;
mutable needed_bytes : int;
}
let create () = { buf = Buffer.create 4096; needed_bytes = -1 }
let dispatch_backend_msg len =
let open Reader in
let open Frame.Backend in
function
| 'R' -> (
int32
>>= fun n ->
match Int32.to_int n with
| 0 -> return AuthenticationOk
| 2 -> return AuthenticationKerberosV5
| 3 -> return AuthenticationCleartextPassword
| 5 -> bytes 4 >>= fun salt -> return (AuthenticationMD5Password { salt })
| 6 -> return AuthenticationSCMCredential
| 7 -> return AuthenticationGSS
| 9 -> return AuthenticationSSPI
| 8 -> bytes (len - 4) >>= fun data -> return (AuthenticationGSSContinue { data })
| 10 ->
string_list ()
>>= fun auth_mechanisms -> return (AuthenticationSASL { auth_mechanisms })
| 11 -> bytes (len - 4) >>= fun data -> return (AuthenticationSASLContinue { data })
| 12 -> bytes (len - 4) >>= fun data -> return (AuthenticationSASLFinal { data })
| _ -> fail `Invalid_frame)
| 'K' ->
int32 >>= fun pid -> int32 >>= fun secret_key -> return (BackendKeyData { pid; secret_key })
| '2' -> return BindComplete
| '3' -> return CloseComplete
| 'C' -> string >>= fun tag -> return (CommandComplete { tag })
| 'G' ->
int8
>>= fun format ->
int16
>>= fun num_columns ->
repeat num_columns int16
>>= fun column_formats -> return (CopyInResponse { format; column_formats })
| 'H' ->
int8
>>= fun format ->
int16
>>= fun num_columns ->
repeat num_columns int16
>>= fun column_formats -> return (CopyOutResponse { format; column_formats })
| 'W' ->
int8
>>= fun format ->
int16
>>= fun num_columns ->
repeat num_columns int16
>>= fun column_formats -> return (CopyBothResponse { format; column_formats })
| 'd' -> bytes len >>= fun data -> return (CopyData { data })
| 'c' -> return CopyDone
| 'D' ->
int16
>>= fun columns ->
assert (columns > 0);
repeat
columns
(int32
>>= fun n ->
match Int32.to_int n with
| -1 -> return None
| n -> bytes n >>= fun s -> return (Some s))
>>= fun data -> return (DataRow { data })
| 'I' -> return EmptyQueryResponse
| 'E' ->
consume (len - 4) (bytes 1 >>= fun code -> string >>= fun msg -> return (code.[0], msg))
>>= fun msgs ->
bytes 1
>>= fun s ->
assert (s.[0] = '\000');
return (ErrorResponse { msgs })
| 'V' ->
int32
>>= fun n ->
(match Int32.to_int n with
| -1 -> return None
| n -> bytes n >>= fun s -> return (Some s))
>>= fun result -> return (FunctionCallResponse { result })
| 'v' ->
int32
>>= fun minor_version ->
int32
>>= fun n ->
repeat (Int32.to_int n) string
>>= fun unrecognized_options ->
return (NegotiateProtocolVersion { minor_version; unrecognized_options })
| 'n' -> return NoData
| 'N' ->
consume (len - 4) (bytes 1 >>= fun code -> string >>= fun msg -> return (code.[0], msg))
>>= fun msgs ->
bytes 1
>>= fun s ->
assert (s.[0] = '\000');
return (NoticeResponse { msgs })
| 'A' ->
int32
>>= fun pid ->
string
>>= fun channel ->
string >>= fun payload -> return (NotificationResponse { pid; channel; payload })
| 't' ->
int16
>>= fun n ->
repeat n int32 >>= fun object_ids -> return (ParameterDescription { object_ids })
| 'S' -> string >>= fun name -> string >>= fun value -> return (ParameterStatus { name; value })
| '1' -> return ParseComplete
| 's' -> return PortalSuspended
| 'Z' -> bytes 1 >>= fun status -> return (ReadyForQuery { status = status.[0] })
| 'T' ->
int16
>>= fun num_fields ->
repeat
num_fields
(string
>>= fun name ->
int32
>>= fun table_id ->
int16
>>= fun column_attr ->
int32
>>= fun data_type_id ->
int16
>>= fun data_type_size ->
int32
>>= fun data_type_mod ->
int16
>>= fun format_code ->
return
{
name;
table_id;
column_attr;
data_type_id;
data_type_size;
data_type_mod;
format_code;
})
>>= fun rows -> return (RowDescription { rows })
| t -> fail (`Unknown_type t)
let rec backend_msg' pos buf len =
assert (pos >= 0);
assert (len >= 0);
assert (len >= fun msg_typ ->
int32 >>= fun len -> dispatch_backend_msg (Int32.to_int len - 4) msg_typ.[0])
in
match res with
| Ok (frame, pos') -> (
match backend_msg' pos' buf (len - (pos' - pos)) with
| Ok (frames, pos) -> Ok (frame :: frames, pos)
| Error (frames, `Length, _) -> Ok (frame :: frames, pos')
| Error (frames, err, pos) -> Error (frame :: frames, err, pos))
| Error (err, _) -> Error ([], err, pos)
let run_backend_msg t pos buf len =
match backend_msg' pos buf len with
| Error ([], (`Unknown_type _ as err), _) | Error ([], (`Invalid_frame as err), _) -> Error err
| Error (frames, `Needed_bytes n, pos) ->
assert (n > 0);
Buffer.clear t.buf;
Buffer.add_subbytes t.buf buf pos (len - pos);
t.needed_bytes
Buffer.clear t.buf;
Buffer.add_subbytes t.buf buf pos (len - pos);
Ok frames
let backend_msg t ~pos ~len buf =
if Buffer.length t.buf = 0 && t.needed_bytes < 0 then run_backend_msg t pos buf len
else if t.needed_bytes < 0 then (
Buffer.add_subbytes t.buf buf pos len;
let b = Buffer.to_bytes t.buf in
run_backend_msg t 0 b (Bytes.length b))
else (
Buffer.add_subbytes t.buf buf pos len;
t.needed_bytes