Added SafeDisconn for accept.

4 files changed, 63 insertions, 22 deletions

diff --git a/lib/pure/asyncdispatch.nim b/lib/pure/asyncdispatch.nim
index dea17d146..e1837b0ea 100644
--- a/lib/pure/asyncdispatch.nim
+++ b/lib/pure/asyncdispatch.nim
@@ -524,7 +524,10 @@ when defined(windows) or defined(nimdoc):
           if errcode == TOSErrorCode(-1):
             retFuture.complete()
           else:
-            retFuture.fail(newException(EOS, osErrorMsg(errcode)))
+            if flags.isDisconnectionError(errcode):
+              retFuture.complete()
+            else:
+              retFuture.fail(newException(EOS, osErrorMsg(errcode)))
     )
 
     let ret = WSASend(socket.TSocketHandle, addr dataBuf, 1, addr bytesReceived,
@@ -544,13 +547,19 @@ when defined(windows) or defined(nimdoc):
       # free ``ol``.
     return retFuture
 
-  proc acceptAddr*(socket: TAsyncFD): 
+  proc acceptAddr*(socket: TAsyncFD, flags = {TSocketFlags.SafeDisconn}):
       PFuture[tuple[address: string, client: TAsyncFD]] =
     ## Accepts a new connection. Returns a future containing the client socket
     ## corresponding to that connection and the remote address of the client.
     ## The future will complete when the connection is successfully accepted.
     ##
-    ## The resulting client socket is automatically registered to dispatcher.
+    ## The resulting client socket is automatically registered to the
+    ## dispatcher.
+    ##
+    ## The ``accept`` call may result in an error if the connecting socket
+    ## disconnects during the duration of the ``accept``. If the ``SafeDisconn``
+    ## flag is specified then this error will not be raised and instead
+    ## accept will be called again.
     verifyPresence(socket)
     var retFuture = newFuture[tuple[address: string, client: TAsyncFD]]("acceptAddr")
 
@@ -584,6 +593,18 @@ when defined(windows) or defined(nimdoc):
          client: clientSock.TAsyncFD)
       )
 
+    template failAccept(errcode): stmt =
+      if flags.isDisconnectionError(errcode):
+        var newAcceptFut = acceptAddr(socket, flags)
+        newAcceptFut.callback =
+          proc () =
+            if newAcceptFut.failed:
+              retFuture.fail(newAcceptFut.readError)
+            else:
+              retFuture.complete(newAcceptFut.read)
+      else:
+        retFuture.fail(newException(EOS, osErrorMsg(errcode)))
+
     var ol = PCustomOverlapped()
     GC_ref(ol)
     ol.data = TCompletionData(sock: socket, cb:
@@ -592,7 +613,7 @@ when defined(windows) or defined(nimdoc):
           if errcode == TOSErrorCode(-1):
             completeAccept()
           else:
-            retFuture.fail(newException(EOS, osErrorMsg(errcode)))
+            failAccept(errcode)
     )
 
     # http://msdn.microsoft.com/en-us/library/windows/desktop/ms737524%28v=vs.85%29.aspx
@@ -605,7 +626,7 @@ when defined(windows) or defined(nimdoc):
     if not ret:
       let err = osLastError()
       if err.int32 != ERROR_IO_PENDING:
-        retFuture.fail(newException(EOS, osErrorMsg(err)))
+        failAccept(err)
         GC_unref(ol)
     else:
       completeAccept()
@@ -749,7 +770,7 @@ else:
   
   proc connect*(socket: TAsyncFD, address: string, port: TPort,
     af = AF_INET): PFuture[void] =
-    var retFuture = newFuture[void]()
+    var retFuture = newFuture[void]("connect")
     
     proc cb(sock: TAsyncFD): bool =
       # We have connected.
@@ -784,7 +805,7 @@ else:
 
   proc recv*(socket: TAsyncFD, size: int,
              flags = {TSocketFlags.SafeDisconn}): PFuture[string] =
-    var retFuture = newFuture[string]()
+    var retFuture = newFuture[string]("recv")
     
     var readBuffer = newString(size)
 
@@ -815,7 +836,7 @@ else:
 
   proc send*(socket: TAsyncFD, data: string,
              flags = {TSocketFlags.SafeDisconn}): PFuture[void] =
-    var retFuture = newFuture[void]()
+    var retFuture = newFuture[void]("send")
     
     var written = 0
     
@@ -845,9 +866,10 @@ else:
     addWrite(socket, cb)
     return retFuture
 
-  proc acceptAddr*(socket: TAsyncFD): 
+  proc acceptAddr*(socket: TAsyncFD, flags = {TSocketFlags.SafeDisconn}):
       PFuture[tuple[address: string, client: TAsyncFD]] =
-    var retFuture = newFuture[tuple[address: string, client: TAsyncFD]]()
+    var retFuture = newFuture[tuple[address: string,
+        client: TAsyncFD]]("acceptAddr")
     proc cb(sock: TAsyncFD): bool =
       result = true
       var sockAddress: Tsockaddr_in
@@ -860,7 +882,10 @@ else:
         if lastError.int32 == EINTR:
           return false
         else:
-          retFuture.fail(newException(EOS, osErrorMsg(lastError)))
+          if flags.isDisconnectionError(lastError):
+            return false
+          else:
+            retFuture.fail(newException(EOS, osErrorMsg(lastError)))
       else:
         register(client.TAsyncFD)
         retFuture.complete(($inet_ntoa(sockAddress.sin_addr), client.TAsyncFD))
@@ -875,12 +900,13 @@ proc sleepAsync*(ms: int): PFuture[void] =
   p.timers.add((epochTime() + (ms / 1000), retFuture))
   return retFuture
 
-proc accept*(socket: TAsyncFD): PFuture[TAsyncFD] =
+proc accept*(socket: TAsyncFD,
+    flags = {TSocketFlags.SafeDisconn}): PFuture[TAsyncFD] =
   ## Accepts a new connection. Returns a future containing the client socket
   ## corresponding to that connection.
   ## The future will complete when the connection is successfully accepted.
   var retFut = newFuture[TAsyncFD]("accept")
-  var fut = acceptAddr(socket)
+  var fut = acceptAddr(socket, flags)
   fut.callback =
     proc (future: PFuture[tuple[address: string, client: TAsyncFD]]) =
       assert future.finished
diff --git a/lib/pure/asyncnet.nim b/lib/pure/asyncnet.nim
index db6f80b06..b81c4b13a 100644
--- a/lib/pure/asyncnet.nim
+++ b/lib/pure/asyncnet.nim
@@ -135,13 +135,13 @@ proc send*(socket: PAsyncSocket, data: string,
   assert socket != nil
   result = send(socket.fd.TAsyncFD, data, flags)
 
-proc acceptAddr*(socket: PAsyncSocket): 
+proc acceptAddr*(socket: PAsyncSocket, flags = {TSocketFlags.SafeDisconn}):
       PFuture[tuple[address: string, client: PAsyncSocket]] =
   ## Accepts a new connection. Returns a future containing the client socket
   ## corresponding to that connection and the remote address of the client.
   ## The future will complete when the connection is successfully accepted.
   var retFuture = newFuture[tuple[address: string, client: PAsyncSocket]]("asyncnet.acceptAddr")
-  var fut = acceptAddr(socket.fd.TAsyncFD)
+  var fut = acceptAddr(socket.fd.TAsyncFD, flags)
   fut.callback =
     proc (future: PFuture[tuple[address: string, client: TAsyncFD]]) =
       assert future.finished
@@ -153,12 +153,13 @@ proc acceptAddr*(socket: PAsyncSocket):
         retFuture.complete(resultTup)
   return retFuture
 
-proc accept*(socket: PAsyncSocket): PFuture[PAsyncSocket] =
+proc accept*(socket: PAsyncSocket,
+    flags = {TSocketFlags.SafeDisconn}): PFuture[PAsyncSocket] =
   ## Accepts a new connection. Returns a future containing the client socket
   ## corresponding to that connection.
   ## The future will complete when the connection is successfully accepted.
   var retFut = newFuture[PAsyncSocket]("asyncnet.accept")
-  var fut = acceptAddr(socket)
+  var fut = acceptAddr(socket, flags)
   fut.callback =
     proc (future: PFuture[tuple[address: string, client: PAsyncSocket]]) =
       assert future.finished
diff --git a/lib/pure/net.nim b/lib/pure/net.nim
index 5f83b1dea..224087994 100644
--- a/lib/pure/net.nim
+++ b/lib/pure/net.nim
@@ -569,8 +569,8 @@ proc bindAddr*(socket: PSocket, port = TPort(0), address = "") {.
       osError(osLastError())
     dealloc(aiList)
 
-proc acceptAddr*(server: PSocket, client: var PSocket, address: var string) {.
-  tags: [FReadIO].} =
+proc acceptAddr*(server: PSocket, client: var PSocket, address: var string,
+                 flags = {TSocketFlags.SafeDisconn}) {.tags: [FReadIO].} =
   ## Blocks until a connection is being made from a client. When a connection
   ## is made sets ``client`` to the client socket and ``address`` to the address
   ## of the connecting client.
@@ -581,6 +581,11 @@ proc acceptAddr*(server: PSocket, client: var PSocket, address: var string) {.
   ##
   ## **Note**: ``client`` must be initialised (with ``new``), this function 
   ## makes no effort to initialise the ``client`` variable.
+  ##
+  ## The ``accept`` call may result in an error if the connecting socket
+  ## disconnects during the duration of the ``accept``. If the ``SafeDisconn``
+  ## flag is specified then this error will not be raised and instead
+  ## accept will be called again.
   assert(client != nil)
   var sockAddress: Tsockaddr_in
   var addrLen = sizeof(sockAddress).TSocklen
@@ -589,6 +594,8 @@ proc acceptAddr*(server: PSocket, client: var PSocket, address: var string) {.
   
   if sock == osInvalidSocket:
     let err = osLastError()
+    if flags.isDisconnectionError(err):
+      acceptAddr(server, client, address, flags)
     osError(err)
   else:
     client.fd = sock
@@ -658,15 +665,20 @@ when false: #defined(ssl):
       acceptAddrPlain(AcceptNoClient, AcceptSuccess):
         doHandshake()
 
-proc accept*(server: PSocket, client: var PSocket) {.tags: [FReadIO].} =
+proc accept*(server: PSocket, client: var PSocket,
+             flags = {TSocketFlags.SafeDisconn}) {.tags: [FReadIO].} =
   ## Equivalent to ``acceptAddr`` but doesn't return the address, only the
   ## socket.
   ## 
   ## **Note**: ``client`` must be initialised (with ``new``), this function
   ## makes no effort to initialise the ``client`` variable.
-  
+  ##
+  ## The ``accept`` call may result in an error if the connecting socket
+  ## disconnects during the duration of the ``accept``. If the ``SafeDisconn``
+  ## flag is specified then this error will not be raised and instead
+  ## accept will be called again.
   var addrDummy = ""
-  acceptAddr(server, client, addrDummy)
+  acceptAddr(server, client, addrDummy, flags)
 
 proc close*(socket: PSocket) =
   ## Closes a socket.
diff --git a/web/news.txt b/web/news.txt
index 20e241221..a4cbefd52 100644
--- a/web/news.txt
+++ b/web/news.txt
@@ -38,6 +38,8 @@ News
   - ``sequtils.distnct`` has been renamed to ``sequtils.deduplicate``.
   - Added ``algorithm.reversed``
   - Added ``uri.combine`` and ``uri.parseUri``.
+  - Some sockets procedures now support a ``SafeDisconn`` flag which causes 
+    them to handle disconnection errors and not raise them.
 
 2014-04-21 Version 0.9.4 released
 =================================