Properly closing client connections, reworked client connection creation

Now the test script has two subsequent connections in it, for great justice.

chat.ml

@@ -177,9 +177,9 @@ let chat script proc =
       Unixqueue.run ues;
     with
       | Chat_match (got, expected) ->
-          raise (Failure ("Not matched: got " ^
+          raise (Failure ("Not matched: got \"" ^
                             (String.escaped got) ^
-                            ", expected " ^
+                            "\"\n  expected " ^
                             (string_of_chat_event expected)))
       | Chat_timeout evt ->
           raise (Failure ("Timeout waiting for " ^

client.ml

@@ -21,14 +21,8 @@ let by_nick = Hashtbl.create 25
 let error num args text =
   Error (Command.create (Some !(Irc.name)) num args (Some text))
 
-let close cli ues g fd =
-  Hashtbl.remove by_nick !(cli.nick);
-  Unix.close fd;
-  Unixqueue.remove_resource ues g (Unixqueue.Wait_in fd);
-  try
-    Unixqueue.remove_resource ues g (Unixqueue.Wait_out fd);
-  with Not_found ->
-    ()
+let close cli =
+  Iobuf.close cli.iobuf
 
 let write cli cmd =
   Iobuf.write cli.iobuf cmd
@@ -40,6 +34,9 @@ let reply cli num ?(args=[]) text =
                ([!(cli.nick)] @ args)
                (Some text))
 
+let handle_close cli () =
+  Hashtbl.remove by_nick !(cli.nick)
+
 let handle_command cli iobuf cmd =
   match (Command.as_tuple cmd) with
     | (None, "OPER", [name; password], None) ->
@@ -174,7 +171,7 @@ let rec handle_command_prereg (nick', username', realname', password') iobuf cmd
                          " " ^ Irc.version ^
                          " " ^ modes ^
                          " " ^ Channel.modes);
-      Iobuf.rebind iobuf (handle_command cli)
+      Iobuf.rebind iobuf (handle_command cli) (handle_close cli)
     with Error cmd ->
       Iobuf.write iobuf cmd
   in
@@ -194,7 +191,10 @@ let rec handle_command_prereg (nick', username', realname', password') iobuf cmd
                  username = username;
                  realname = realname}
     | _ ->
-        Iobuf.rebind iobuf (handle_command_prereg acc)
+        Iobuf.rebind iobuf (handle_command_prereg acc) ignore
+
+let handle_connection ues grp fd =
+  let command_handler = handle_command_prereg (None, None, None, None) in
+  let close_handler = ignore in
+    Iobuf.bind ues grp fd command_handler close_handler
     
-let create_command_handler () =
-  handle_command_prereg (None, None, None, None)

client.mli

@@ -1,5 +1,5 @@
 type t
 
 val write : t -> Command.t -> unit
-val create_command_handler : unit -> Iobuf.t -> Command.t -> unit
+val handle_connection : Unixqueue.event_system -> Unixqueue.group -> Unix.file_descr -> unit
   

iobuf.ml

@@ -8,7 +8,8 @@ type t = {ues: Unixqueue.event_system;
           unsent: string ref;
           ibuf: string;
           ibuf_len: int ref;
-          handle_command: (t -> Command.t -> unit) ref}
+          command_handler: (t -> Command.t -> unit) ref;
+	  close_handler: (unit -> unit) ref}
 
 let ibuf_max = 4096
 let max_outq = 50
@@ -16,25 +17,6 @@ let obuf_max = 4096
 
 let by_file_descr = Hashtbl.create 25
 
-let bind ues grp fd handle_command =
-  let (outq, unsent, ibuf, ibuf_len) =
-    (Queue.create (), ref "", String.create ibuf_max, ref 0)
-  in
-  let iobuf = {ues = ues;
-               grp = grp;
-               fd = fd;
-               outq = outq;
-               unsent = unsent;
-               ibuf = ibuf;
-               ibuf_len = ibuf_len;
-               handle_command = ref handle_command}
-  in
-    Hashtbl.replace by_file_descr fd iobuf;
-    Unixqueue.add_resource ues grp (Unixqueue.Wait_in fd, -.1.0)
-
-let rebind t handle_command =
-  t.handle_command := handle_command
-
 let write iobuf cmd =
   let was_empty = Queue.is_empty iobuf.outq in
     Queue.add cmd iobuf.outq;
@@ -43,6 +25,7 @@ let write iobuf cmd =
         iobuf.ues iobuf.grp (Unixqueue.Wait_out iobuf.fd, -.1.0)
 
 let close iobuf =
+  !(iobuf.close_handler) ();
   Hashtbl.remove by_file_descr iobuf.fd;
   Unix.close iobuf.fd;
   Unixqueue.remove_resource iobuf.ues iobuf.grp (Unixqueue.Wait_in iobuf.fd);
@@ -51,6 +34,13 @@ let close iobuf =
   with Not_found ->
     ()
 
+let handle_close fd =
+  try
+    let iobuf = Hashtbl.find by_file_descr fd in
+      close iobuf
+  with Not_found ->
+    ()
+
 let crlf = Str.regexp "\r?\n"
 
 let handle_input iobuf =
@@ -65,7 +55,7 @@ let handle_input iobuf =
 	  String.blit leftover 0 iobuf.ibuf 0 !(iobuf.ibuf_len)
       | line :: tl ->
 	  let parsed = Command.from_string line in
-	    !(iobuf.handle_command) iobuf parsed;
+	    !(iobuf.command_handler) iobuf parsed;
 	    loop tl
   in
     loop lines
@@ -113,5 +103,24 @@ let handle_event ues esys e =
     | Unixqueue.Extra exn ->
         print_endline "extra"
 
-let add_event_handler ues g =
-  Unixqueue.add_handler ues g handle_event
+let bind ues grp fd command_handler close_handler =
+  let (outq, unsent, ibuf, ibuf_len) =
+    (Queue.create (), ref "", String.create ibuf_max, ref 0)
+  in
+  let iobuf = {ues = ues;
+               grp = grp;
+               fd = fd;
+               outq = outq;
+               unsent = unsent;
+               ibuf = ibuf;
+               ibuf_len = ibuf_len;
+               command_handler = ref command_handler;
+	       close_handler = ref close_handler}
+  in
+    Hashtbl.replace by_file_descr fd iobuf;
+    Unixqueue.add_resource ues grp (Unixqueue.Wait_in fd, -.1.0);
+    Unixqueue.add_close_action ues grp (fd, handle_close)
+
+let rebind t command_handler close_handler =
+  t.command_handler := command_handler;
+  t.close_handler := close_handler

iobuf.mli

@@ -1,8 +1,7 @@
 type t
 
 val write : t -> Command.t -> unit
-val bind : Unixqueue.event_system -> Unixqueue.group -> Unix.file_descr -> (t -> Command.t -> unit) -> unit
-val rebind: t -> (t -> Command.t -> unit) -> unit
+val bind : Unixqueue.event_system -> Unixqueue.group -> Unix.file_descr -> (t -> Command.t -> unit) -> (unit -> unit) -> unit
+val rebind: t -> (t -> Command.t -> unit) -> (unit -> unit) -> unit
 val close: t -> unit
-val add_event_handler : Unixqueue.event_system -> Unixqueue.group -> unit
-
+val handle_event : Unixqueue.event_system -> Unixqueue.event Equeue.t -> Unixqueue.event -> unit

ircd.ml

@@ -27,13 +27,10 @@ let establish_server ues connection_handler addr =
 let main () =
   let ues = Unixqueue.create_unix_event_system () in
   let g = Unixqueue.new_group ues in
-  let handle_connection fd =
-    Iobuf.bind ues g fd (Client.create_command_handler ())
-  in
-    Iobuf.add_event_handler ues g;
+    Unixqueue.add_handler ues g Iobuf.handle_event;
     establish_server
       ues
-      handle_connection
+      Client.handle_connection
       (Unix.ADDR_INET (Unix.inet_addr_any, 7777));
     ues#run ()
 

tests.ml

@@ -6,8 +6,8 @@ open Irc
 let do_chat script () =
   let ircd_instance ues fd =
     let g = Unixqueue.new_group ues in
-      Iobuf.add_event_handler ues g;
-      Iobuf.bind ues g fd (Client.create_command_handler ())
+      Unixqueue.add_handler ues g Iobuf.handle_event;
+      Client.handle_connection ues g fd
   in
     chat script ircd_instance
 
@@ -70,6 +70,13 @@ let regression_tests =
 		      Send "ISON otherguy thirdguy\r\n";
 		      Recv ":testserver.test 303 nick :\r\n";
                     ]));
+
+      "Second connection" >::
+	(do_chat ((do_login "otherguy") @
+		    [
+		      Send "ISON nick otherguy\r\n";
+		      Recv ":testserver.test 303 otherguy :otherguy\r\n";
+		    ]));
     ]
 
 let _ =