Merge pull request #19 from c-cube/wait-for-welcome

wait for `RPL_WELCOME` upon listening

examples/example2.ml

@@ -6,7 +6,6 @@ let host = ref "irc.freenode.net"
 let port = ref 6667
 let nick = ref "bobobobot"
 let channel = ref "#demo_irc"
-let sleep = ref 5.
 let message = "Hello, world!  This is a test from ocaml-irc-client"
 
 let string_list_to_string string_list =
@@ -28,15 +27,16 @@ let callback connection result =
     Lwt_io.printl e
 
 let lwt_main =
+  Lwt_io.printl "Connecting..."
+  >>= fun () ->
   C.connect_by_name ~server:!host ~port:!port ~nick:!nick ()
   >>= function
   | None -> Lwt_io.printl "could not find host"
   | Some connection ->
   Lwt_io.printl "Connected"
   >>= fun () ->
   let t = C.listen ~connection ~callback in
-  Lwt_unix.sleep !sleep
-  >>= fun () -> Lwt_io.printl "send join msg"
+  Lwt_io.printl "send join msg"
   >>= fun () -> C.send_join ~connection ~channel:!channel
   >>= fun () -> C.send_privmsg ~connection ~target:!channel ~message
   >>= fun () -> t (* wait for completion of t *)
@@ -46,7 +46,6 @@ let options = Arg.align
   [ "-host", Arg.Set_string host, " set remove server host name"
   ; "-port", Arg.Set_int port, " set remote server port"
   ; "-chan", Arg.Set_string channel, " channel to join"
-  ; "-sleep", Arg.Set_float sleep, " sleep before joining (in s)"
   ]
 
 let _ =

lib/irc_client.ml

@@ -1,6 +1,11 @@
 module Make(Io: Irc_transport.IO) = struct
   type connection_t = {
     sock: Io.file_descr;
+    buffer: Buffer.t;
+    read_length: int;
+    read_data: Bytes.t; (* for reading *)
+    lines: string Queue.t; (* lines read so far *)
+    mutable terminated: bool;
   }
 
   open Io
@@ -48,18 +53,62 @@ module Make(Io: Irc_transport.IO) = struct
     let msg = M.user ~username ~mode ~realname in
     send ~connection msg
 
+  let mk_connection_ sock =
+    let read_length = 1024 in
+    { sock = sock;
+      buffer=Buffer.create 128;
+      read_length;
+      read_data = Bytes.make read_length ' ';
+      lines = Queue.create ();
+      terminated = false;
+    }
+
+  let rec next_line_ ~connection:c : string option Io.t =
+    if c.terminated
+    then return None
+    else if Queue.length c.lines > 0
+    then return (Some (Queue.pop c.lines))
+    else (
+      (* Read some data into our string. *)
+      Io.read c.sock c.read_data 0 c.read_length
+      >>= function
+      | 0 ->
+        c.terminated <- true;
+        return None (* EOF from server - we have quit or been kicked. *)
+      | len ->
+        (* read some data, push lines into [c.lines] (if any) *)
+        let input = Bytes.sub_string c.read_data 0 len in
+        let lines = Irc_helpers.handle_input ~buffer:c.buffer ~input in
+        List.iter (fun l -> Queue.push l c.lines) lines;
+        next_line_ ~connection:c
+    )
+
+  let rec wait_for_welcome ~connection =
+    next_line_ ~connection
+    >>= function
+    | None -> return ()
+    | Some line ->
+      match M.parse line with
+        | `Ok {M.command = M.Other ("001", _); _} ->
+          (* we received "RPL_WELCOME", i.e. 001 *)
+          return ()
+        | _ -> wait_for_welcome ~connection
+
+
   let connect
       ?(username="irc-client") ?(mode=0) ?(realname="irc-client")
       ?password ~addr ~port ~nick () =
     Io.open_socket addr port >>= (fun sock ->
-      let connection = {sock = sock} in begin
+      let connection = mk_connection_ sock in
+      begin
         match password with
         | Some password -> send_pass ~connection ~password
         | None -> return ()
       end
-      >>= (fun () -> send_nick ~connection ~nick)
-      >>= (fun () -> send_user ~connection ~username ~mode ~realname)
-      >>= (fun () -> return connection))
+      >>= fun () -> send_nick ~connection ~nick
+      >>= fun () -> send_user ~connection ~username ~mode ~realname
+      >>= fun () -> wait_for_welcome ~connection
+      >>= fun () -> return connection)
 
   let connect_by_name
       ?(username="irc-client") ?(mode=0) ?(realname="irc-client")
@@ -69,32 +118,21 @@ module Make(Io: Irc_transport.IO) = struct
       | [] -> Io.return None
       | addr :: _ ->
         connect ~addr ~port ~username ~mode ~realname ~nick ?password ()
-        >>= (fun connection -> Io.return (Some connection)))
+        >>= fun connection -> Io.return (Some connection))
 
   let listen ~connection ~callback =
-    let read_length = 1024 in
-    let read_data = Bytes.make read_length ' ' in
-    let rec listen' ~buffer =
-      (* Read some data into our string. *)
-      Io.read connection.sock read_data 0 read_length
+    let rec listen' () =
+      next_line_ ~connection
       >>= function
-      | 0 -> return () (* EOF from server - we have quit or been kicked. *)
-      | len ->
-        let input = Bytes.sub_string read_data 0 len in
-        let lines = Irc_helpers.handle_input ~buffer ~input in
-        let _ = Io.iter
-          (fun line ->
-             match M.parse line with
-             | `Ok {M.command = M.PING message; _} ->
-               (* Handle pings without calling the callback. *)
-               send_pong ~connection ~message
-             | result ->
-               callback connection result
-
-          ) lines
-        in
-        listen' ~buffer
+      | None -> return ()
+      | Some line ->
+        begin match M.parse line with
+          | `Ok {M.command = M.PING message; _} ->
+            (* Handle pings without calling the callback. *)
+            send_pong ~connection ~message
+          | result -> callback connection result
+        end
+        >>= listen'
     in
-    let buffer = Buffer.create 256 in
-    listen' ~buffer
+    listen' ()
 end