Read/write from both stdin/stdout and the sockHEADmaster

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

1 files changed, 15 insertions, 5 deletions

diff --git a/netcat.ml b/netcat.ml
index 53c9521..5fcaa28 100644
--- a/netcat.ml
+++ b/netcat.ml
@@ -1,12 +1,22 @@
 (** Connect to addr and port and start communicating **)
 let nc addr port =
     let rec discuss csock =
-        let buffer = Bytes.create 5 in
-        let n = Unix.recv csock buffer 0 (Bytes.length buffer) [] in
-        match n with
-        | 0 -> ()
-        | _ -> ignore (Unix.write Unix.stdout buffer 0 n); discuss csock
+        let aux input output =
+            let buffer = Bytes.create 256 in
+            let n = Unix.read input buffer 0 (Bytes.length buffer) in
+            match n with
+            | 0 -> ()
+            | _ -> ignore (Unix.write output buffer 0 n); discuss csock
+        in
+
+        match Unix.select [Unix.stdin; csock] [] [] (-1.) with
+        | (fds, _, _) -> (
+            match List.hd fds == Unix.stdin with
+            | true  -> aux Unix.stdin csock
+            | false -> aux csock Unix.stdout
+        )
     in
+
     let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
     Unix.connect sock (Unix.ADDR_INET (Unix.inet_addr_of_string addr, port));
     discuss sock