feat(portal_tcp): add tcp_stream function

2 files changed, 29 insertions, 1 deletions

diff --git a/portal/tcp/dune b/portal/tcp/dune
index 43366c0..40fa592 100644
--- a/portal/tcp/dune
+++ b/portal/tcp/dune
@@ -1,3 +1,3 @@
 (library
  (name portal_tcp)
- (libraries lwt markup markup-lwt))
+ (libraries lwt lwt.unix markup markup-lwt))
diff --git a/portal/tcp/portal_tcp.ml b/portal/tcp/portal_tcp.ml
index 3aa5e56..3ee3716 100644
--- a/portal/tcp/portal_tcp.ml
+++ b/portal/tcp/portal_tcp.ml
@@ -1,8 +1,36 @@
+open Lwt.Syntax
+open Lwt_unix
+
 (** [xmpp_port domain] is the port where [domain]'s XMPP server is hosted.
 
     Currently, it falls back to 5222 (always), but should use SRV records in the near
     future. *)
 let xmpp_port (_domain : string) : int = 5222
 
+(** [tcp_stream domain] is a (stream, push) tuple communicating with the XMPP server
+    hosted on [domain] via plaintext TCP. *)
+let tcp_stream (domain : string) =
+  let get_socket {ai_addr; ai_family; _} =
+    let sock = socket ai_family SOCK_STREAM 0
+    in let+ () = Lwt_unix.connect sock ai_addr
+       in sock
+  and port_number = xmpp_port domain |> string_of_int in
+  let* addrinfos = getaddrinfo domain port_number [AI_SOCKTYPE SOCK_STREAM]
+  in let+ sock = List.map get_socket addrinfos |> Lwt.pick
+     in let stream =
+          Lwt_stream.from (fun () ->
+              let bsize = 4096 in
+              let buffer = Bytes.create 4096 in
+              let* len = read sock buffer 0 bsize
+              in match len with
+                 | 0 -> Lwt.return_none
+                 | len -> Lwt.return_some (Bytes.sub_string buffer 0 len)
+            )
+        and push msg =
+          let none () = close sock
+          and some s () = write_string sock s 0 (String.length s) |> Lwt.map ignore
+          in Option.fold ~none ~some msg |> Lwt.async
+        in (stream, push)
+
 let connect (domain : string) =
   xmpp_port domain |> string_of_int |> print_endline