86 lines
2.7 KiB
OCaml
86 lines
2.7 KiB
OCaml
|
(* (c) 2017, 2018 Hannes Mehnert, all rights reserved *)
|
||
|
|
|
||
|
|
open Lwt.Infix
|
||
|
|
|
||
|
|
let src = Logs.Src.create "dns_mirage" ~doc:"effectful DNS layer"
|
||
|
|
module Log = (val Logs.src_log src : Logs.LOG)
|
||
|
|
|
||
|
|
module Make (S : Tcpip.Stack.V4V6) = struct
|
||
|
|
|
||
|
|
module IPM = struct
|
||
|
|
include Map.Make(struct
|
||
|
|
type t = Ipaddr.t * int
|
||
|
|
let compare (ip, p) (ip', p') = match Ipaddr.compare ip ip' with
|
||
|
|
| 0 -> compare p p'
|
||
|
|
| x -> x
|
||
|
|
end)
|
||
|
|
let find k t = try Some (find k t) with Not_found -> None
|
||
|
|
end
|
||
|
|
|
||
|
|
module U = S.UDP
|
||
|
|
module T = S.TCP
|
||
|
|
|
||
|
|
type f = {
|
||
|
|
flow : T.flow ;
|
||
|
|
mutable linger : Cstruct.t ;
|
||
|
|
}
|
||
|
|
|
||
|
|
let of_flow flow = { flow ; linger = Cstruct.empty }
|
||
|
|
|
||
|
|
let flow { flow ; _ } = flow
|
||
|
|
|
||
|
|
let rec read_exactly f length =
|
||
|
|
let dst_ip, dst_port = T.dst f.flow in
|
||
|
|
if Cstruct.length f.linger >= length then
|
||
|
|
let a, b = Cstruct.split f.linger length in
|
||
|
|
f.linger <- b ;
|
||
|
|
Lwt.return (Ok a)
|
||
|
|
else
|
||
|
|
T.read f.flow >>= function
|
||
|
|
| Ok `Eof ->
|
||
|
|
Log.debug (fun m -> m "end of file on flow %a:%d" Ipaddr.pp dst_ip dst_port) ;
|
||
|
|
T.close f.flow >>= fun () ->
|
||
|
|
Lwt.return (Error ())
|
||
|
|
| Error e ->
|
||
|
|
Log.err (fun m -> m "error %a reading flow %a:%d" T.pp_error e Ipaddr.pp dst_ip dst_port) ;
|
||
|
|
T.close f.flow >>= fun () ->
|
||
|
|
Lwt.return (Error ())
|
||
|
|
| Ok (`Data b) ->
|
||
|
|
f.linger <- Cstruct.append f.linger b ;
|
||
|
|
read_exactly f length
|
||
|
|
|
||
|
|
let send_udp stack src_port dst dst_port data =
|
||
|
|
Log.debug (fun m -> m "udp: sending %d bytes from %d to %a:%d"
|
||
|
|
(Cstruct.length data) src_port Ipaddr.pp dst dst_port) ;
|
||
|
|
U.write ~src_port ~dst ~dst_port (S.udp stack) data >|= function
|
||
|
|
| Error e -> Log.warn (fun m -> m "udp: failure %a while sending from %d to %a:%d"
|
||
|
|
U.pp_error e src_port Ipaddr.pp dst dst_port)
|
||
|
|
| Ok () -> ()
|
||
|
|
|
||
|
|
let send_tcp flow answer =
|
||
|
|
let dst_ip, dst_port = T.dst flow in
|
||
|
|
Log.debug (fun m -> m "tcp: sending %d bytes to %a:%d" (Cstruct.length answer) Ipaddr.pp dst_ip dst_port) ;
|
||
|
|
let len = Cstruct.create 2 in
|
||
|
|
Cstruct.BE.set_uint16 len 0 (Cstruct.length answer) ;
|
||
|
|
T.write flow (Cstruct.append len answer) >>= function
|
||
|
|
| Ok () -> Lwt.return (Ok ())
|
||
|
|
| Error e ->
|
||
|
|
Log.err (fun m -> m "tcp: error %a while writing to %a:%d" T.pp_write_error e Ipaddr.pp dst_ip dst_port) ;
|
||
|
|
T.close flow >|= fun () ->
|
||
|
|
Error ()
|
||
|
|
|
||
|
|
let send_tcp_multiple flow datas =
|
||
|
|
Lwt_list.fold_left_s (fun acc d ->
|
||
|
|
match acc with
|
||
|
|
| Error () -> Lwt.return (Error ())
|
||
|
|
| Ok () -> send_tcp flow d)
|
||
|
|
(Ok ()) datas
|
||
|
|
|
||
|
|
let read_tcp flow =
|
||
|
|
read_exactly flow 2 >>= function
|
||
|
|
| Error () -> Lwt.return (Error ())
|
||
|
|
| Ok l ->
|
||
|
|
let len = Cstruct.BE.get_uint16 l 0 in
|
||
|
|
read_exactly flow len
|
||
|
|
end
|