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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
open Rresult
let ( <.> ) f g x = f (g x)
module type FLOW = sig
type t
type error
type +'a io
val read :
t -> bytes -> int -> int -> ([ `End | `Len of int ], error) result io
val fully_write : t -> string -> int -> int -> (unit, error) result io
val close : t -> unit io
val bind : 'a io -> ('a -> 'b io) -> 'b io
val map : ('a -> 'b) -> 'a io -> 'b io
val return : 'a -> 'a io
end
module Make (Flow : FLOW) = struct
type error =
| Alert of Tls.Packet.alert_type
| Failure of Tls.Engine.failure
| Flow_error of Flow.error
| Closed
let ( >>= ) = Flow.bind
let ( >>| ) x f = Flow.map f x
let return = Flow.return
type t = {
socket : Flow.t;
mutable state : [ `Active of Tls.Engine.state
| `Read_closed of Tls.Engine.state
| `Write_closed of Tls.Engine.state
| `Closed
| `Error of error ];
mutable linger : Cstruct.t list;
}
let fully_write socket ({ Cstruct.len; _ } as cs) =
Flow.fully_write socket (Cstruct.to_string cs) 0 len
let read socket =
let buf = Bytes.create 0x1000 in
Flow.read socket buf 0 (Bytes.length buf) >>| function
| Ok `End -> Ok `Eof
| Ok (`Len len) -> Ok (`Data (Cstruct.of_bytes buf ~off:0 ~len))
| Error _ as err -> err
let half_close state mode = match state, mode with
| `Active tls, `read -> `Read_closed tls
| `Active tls, `write -> `Write_closed tls
| `Active _, `read_write -> `Closed
| `Read_closed tls, `read -> `Read_closed tls
| `Read_closed _, (`write | `read_write) -> `Closed
| `Write_closed tls, `write -> `Write_closed tls
| `Write_closed _, (`read | `read_write) -> `Closed
| (`Closed | `Error _) as e, (`read | `write | `read_write) -> e
let inject_state tls = function
| `Active _ -> `Active tls
| `Read_closed _ -> `Read_closed tls
| `Write_closed _ -> `Write_closed tls
| (`Closed | `Error _) as e -> e
let write_flow flow buf =
fully_write flow.socket buf >>= function
| Ok _ as o -> return o
| Error e ->
flow.state <- `Error (Flow_error e);
return (Error (Flow_error e))
let read_react flow =
let handle tls buf =
match Tls.Engine.handle_tls tls buf with
| Ok (state, eof, `Response resp, `Data data) ->
let state = inject_state state flow.state in
let state = Option.(value ~default:state (map (fun `Eof -> half_close state `read) eof)) in
flow.state <- state;
( match resp with
| None -> return (Ok ())
| Some buf -> write_flow flow buf) >>= fun _ ->
return @@ `Ok data
| Error (fail, `Response resp) ->
let r = `Error (Failure fail) in
flow.state <- r ;
fully_write flow.socket resp >>= fun _ -> return r
in
match flow.state with
| `Error _ as e -> return e
| `Read_closed _ | `Closed -> return `Eof
| `Active _ | `Write_closed _ ->
read flow.socket >>= function
| Error e -> flow.state <- `Error (Flow_error e); return (`Error (Flow_error e))
| Ok `Eof ->
flow.state <- half_close flow.state `read;
return `Eof
| Ok `Data buf -> match flow.state with
| `Active tls | `Write_closed tls -> handle tls buf
| `Read_closed _ | `Closed -> return `Eof
| `Error _ as e -> return e
let rec read flow =
match flow.linger with
| [] ->
( read_react flow >>= function
| `Ok None -> read flow
| `Ok (Some buf) -> return (Result.Ok (`Data buf))
| `Eof -> return (Result.Ok `Eof)
| `Error e -> return (Result.Error e))
| bufs ->
flow.linger <- [] ;
return @@ Ok (`Data (Cstruct.concat @@ List.rev bufs))
let writev flow bufs =
match flow.state with
| `Closed | `Write_closed _ -> return (Result.Error Closed)
| `Error err -> return (Result.Error err)
| `Active tls | `Read_closed tls ->
match Tls.Engine.send_application_data tls bufs with
| Some (tls, answer) ->
flow.state <- `Active tls ;
write_flow flow answer
| None ->
assert false
let write flow cs = writev flow [ cs ]
let close flow =
( match flow.state with
| `Active tls | `Read_closed tls ->
let tls, buf = Tls.Engine.send_close_notify tls in
flow.state <- inject_state tls flow.state;
flow.state <- `Closed;
fully_write flow.socket buf >>= fun _ -> return ()
| `Write_closed _ ->
flow.state <- `Closed;
return ()
| _ -> return () ) >>= fun () ->
Flow.close flow.socket
let rec drain_handshake flow =
match flow.state with
| `Active tls when not (Tls.Engine.handshake_in_progress tls) ->
return (Ok flow)
| _ -> (
read_react flow >>= function
| `Ok (Some mbuf) ->
flow.linger <- mbuf :: flow.linger ;
drain_handshake flow
| `Ok None -> drain_handshake flow
| `Error err -> return (Result.Error err)
| `Eof -> return (Result.Error Closed))
let init_client cfg socket =
let tls, init = Tls.Engine.client cfg in
let flow = { socket; state = `Active tls; linger = [] } in
fully_write socket init >>= fun _ -> drain_handshake flow
let init_server cfg socket =
let flow =
{ socket; state = `Active (Tls.Engine.server cfg); linger = [] } in
drain_handshake flow
end