VndbApi: Check for remote disconnect when sending/receiving data

1 files changed, 19 insertions, 7 deletions

diff --git a/vndbApi.ml b/vndbApi.ml
index f837b16..9565f25 100644
--- a/vndbApi.ml
+++ b/vndbApi.ml
@@ -97,18 +97,22 @@ class connection
       in
       let rec dorecv () =
         let cnt = recv s buf 0 (String.length buf) [] in
-        reply := !reply ^ (String.sub buf 0 cnt);
-        match String.contains !reply '\004' with
-        | true  -> finish ()
-        | false -> pollfunc s true dorecv
+        if cnt == 0 then self#disconnect else (
+          reply := !reply ^ (String.sub buf 0 cnt);
+          match String.contains !reply '\004' with
+          | true  -> finish ()
+          | false -> pollfunc s true dorecv
+        )
       in
       let dosend () = 
         (* assumes the command can be sent with a single call *)
-        let _ = send s
+        let cnt = send s
           (cmd ^ (String.make 1 (char_of_int 4)))
           0 (String.length cmd + 1) [] in
-        match linefunc with None -> () | Some f -> f true cmd;
-        pollfunc s true dorecv
+        if cnt == 0 then self#disconnect else (
+          match linefunc with None -> () | Some f -> f true cmd;
+          pollfunc s true dorecv
+        )
       in
       pollfunc s false dosend
     )
@@ -154,6 +158,14 @@ class connection
 
 
   method disconnect =
+    (* we'll have to clear the comamnd queue when disconnected, since we need to
+     * be able to issue the login command the first thing after a reconnect...
+     * It might be an idea to either special-case the login command or call all
+     * registered callback functions with an error.
+     * Behaviour is currently undefined when disconnect is called while a
+     * command is being handled. *)
+    Queue.clear cmd_queue;
+    cmd_running <- false;
     match s_sock with None -> () | Some s -> close s; s_sock <- None;
     match disconnectfunc with None -> () | Some f -> f ()