feat(portal_tcp): move tcp_stream push function to connect

1 files changed, 12 insertions, 11 deletions

diff --git a/portal/tcp/portal_tcp.ml b/portal/tcp/portal_tcp.ml
index 1a672dd..3722b63 100644
--- a/portal/tcp/portal_tcp.ml
+++ b/portal/tcp/portal_tcp.ml
@@ -7,9 +7,9 @@ open Lwt_unix
     future. *)
 let xmpp_port (_domain : string) : int = 5222
 
-(** [tcp_stream domain] is a (stream, push) tuple communicating with the XMPP server
+(** [tcp_stream domain] is a (stream, socket) tuple communicating with the XMPP server
     hosted on [domain] via plaintext TCP. *)
-let tcp_stream (domain : string) =
+let tcp_stream (domain : string) : (string Lwt_stream.t * file_descr) Lwt.t =
   let get_socket {ai_addr; ai_family; _} =
     let sock = socket ai_family SOCK_STREAM 0
     in let+ () = Lwt_unix.connect sock ai_addr
@@ -20,17 +20,18 @@ let tcp_stream (domain : string) =
      in let stream =
           Lwt_stream.from (fun () ->
               let bsize = 4096 in
-              let buffer = Bytes.create 4096 in
+              let buffer = Bytes.create bsize 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)
+                 | len -> Lwt.return_some (Bytes.sub_string buffer 0 len))
+        in (stream, sock)
 
 let connect (domain : string) =
-  tcp_stream domain
+  let+ tcp_stream, tcp_socket = tcp_stream domain
+  in let push msg =
+       let none () = close tcp_socket
+       and some s () =
+         write_string tcp_socket s 0 (String.length s) |> Lwt.map ignore
+       in Option.fold ~none ~some msg |> Lwt.async
+     in tcp_stream, push