Source file `TFramedTransport.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
open Thrift

module T = Transport

let c_0xff_32 = Int32.of_string "0xff"

(* Copied from OCamlnet rtypes.ml *)
let encode_frame_size x =
	let s = Bytes.create 4 in
	let n3 = Int32.to_int (Int32.shift_right_logical x 24) land 0xff in
	let n2 = Int32.to_int (Int32.shift_right_logical x 16) land 0xff in
	let n1 = Int32.to_int (Int32.shift_right_logical x 8) land 0xff in
	let n0 = Int32.to_int (Int32.logand x c_0xff_32) in
		Bytes.unsafe_set s 0 (Char.unsafe_chr n3);
		Bytes.unsafe_set s 1 (Char.unsafe_chr n2);
		Bytes.unsafe_set s 2 (Char.unsafe_chr n1);
		Bytes.unsafe_set s 3 (Char.unsafe_chr n0);
		s
		
let decode_frame_size s = 
	let n3 = Int32.of_int (Char.code @@ Bytes.get s 0) in
	let n2 = Int32.of_int (Char.code @@ Bytes.get s 1) in
	let n1 = Int32.of_int (Char.code @@ Bytes.get s 2) in
	let n0 = Int32.of_int (Char.code @@ Bytes.get s 3) in
		Int32.logor
		(Int32.shift_left n3 24)
		(Int32.logor
			(Int32.shift_left n2 16)
			(Int32.logor
				(Int32.shift_left n1 8)
				n0))

class t ?(max_length=Sys.max_string_length) (transport: T.t) =
object (self)
	inherit T.t

	method isOpen = transport#isOpen
	method opn = transport#opn
	method close = transport#close
 
	val mutable read_buf = None
	val mutable read_buf_offset = 0
	val mutable write_buf = "" (* TODO: use Buffer.t instead *)

	method private read_frame =
		let len_buf = Bytes.create 4 in
		assert (transport#readAll len_buf 0 4 = 4); 
		
		let size = Int32.to_int (decode_frame_size len_buf) in
		
		(if size < 0
		then failwith (Printf.sprintf "Read a negative frame size (%i)!" size));
		
		(if size > max_length
		then failwith (Printf.sprintf "Frame size (%i) larger than max length (%i)!" size max_length));

		let buf = Bytes.create size in
			assert (transport#readAll buf 0 size = size);
			read_buf <- Some buf;
			read_buf_offset <- 0

	method private read_from_frame frame buf off len =
		let to_copy = min len ((Bytes.length frame) - read_buf_offset) in
			Bytes.blit frame read_buf_offset buf off to_copy;
			read_buf_offset <- read_buf_offset + to_copy;
			to_copy

	method read buf off len = 
		match read_buf with
		| Some frame -> 
			let i = self#read_from_frame frame buf off len in
			if i > 0
			then i
			else begin
				self#read_frame; 
				self#read_from_frame frame buf off len
			end
		| None ->
				self#read_frame;
				self#read buf off len 
	 
	method write buf off len = 
		write_buf <- write_buf ^ (Bytes.sub_string buf off len)

	method write_string buf off len = 
		write_buf <- write_buf ^ (String.sub buf off len)

	method flush = 
		let encoded_size = encode_frame_size (Int32.of_int (String.length write_buf)) in
			transport#write encoded_size 0 (Bytes.length encoded_size);
			transport#write_string write_buf 0 (String.length write_buf);
			transport#flush; 
			write_buf <- ""
end