Fixed a problem with message parsing in the IRC module. Introduced timeout to procs in the sockets module.

2 files changed, 87 insertions, 6 deletions

diff --git a/lib/pure/irc.nim b/lib/pure/irc.nim
index 8946a9d8f..81dff024e 100644
--- a/lib/pure/irc.nim
+++ b/lib/pure/irc.nim
@@ -142,6 +142,8 @@ proc parseMessage(msg: string): TIRCEvent =
     inc(i) # Skip `:`
     var nick = ""
     i.inc msg.parseUntil(nick, {'!', ' '}, i)
+    result.nick = ""
+    result.serverName = ""
     if msg[i] == '!':
       result.nick = nick
       inc(i) # Skip `!`
@@ -237,7 +239,8 @@ proc processLine(irc: var TIRC, line: string): TIRCEvent =
     result = parseMessage(line)
     # Get the origin
     result.origin = result.params[0]
-    if result.origin == irc.nick: result.origin = result.nick
+    if result.origin == irc.nick and
+       result.nick != "": result.origin = result.nick
 
     if result.cmd == MError:
       irc.close()
@@ -386,6 +389,7 @@ proc asyncIRC*(address: string, port: TPort = 6667.TPort,
   result.messageBuffer = @[]
   result.handleEvent = ircEvent
   result.userArg = userArg
+  result.lineBuffer = ""
 
 proc register*(d: PDispatcher, irc: PAsyncIRC) =
   ## Registers ``irc`` with dispatcher ``d``.
diff --git a/lib/pure/sockets.nim b/lib/pure/sockets.nim
index 5721d79fe..fb7d337ee 100755
--- a/lib/pure/sockets.nim
+++ b/lib/pure/sockets.nim
@@ -8,8 +8,12 @@
 #
 
 ## This module implements a simple portable type-safe sockets layer.
+##
+## Most procedures raise EOS on error.
+
 
 import os, parseutils
+from times import epochTime
 
 when defined(Windows):
   import winlean
@@ -59,6 +63,8 @@ type
   TRecvLineResult* = enum ## result for recvLineAsync
     RecvFullLine, RecvPartialLine, RecvDisconnected, RecvFail
 
+  ETimeout* = object of ESynch
+
 const
   InvalidSocket* = TSocket(-1'i32) ## invalid socket number
 
@@ -377,12 +383,14 @@ proc connect*(socket: TSocket, name: string, port = TPort(0),
   ## host name. If ``name`` is a host name, this function will try each IP
   ## of that host name. ``htons`` is already performed on ``port`` so you must
   ## not do it.
+  
   var hints: TAddrInfo
   var aiList: ptr TAddrInfo = nil
   hints.ai_family = toInt(af)
   hints.ai_socktype = toInt(SOCK_STREAM)
   hints.ai_protocol = toInt(IPPROTO_TCP)
   gaiNim(name, port, hints, aiList)
+  
   # try all possibilities:
   var success = false
   var it = aiList
@@ -445,7 +453,6 @@ proc connectAsync*(socket: TSocket, name: string, port = TPort(0),
   freeaddrinfo(aiList)
   if not success: OSError()
 
-
 proc timeValFromMilliseconds(timeout = 500): TTimeVal =
   if timeout != -1:
     var seconds = timeout div 1000
@@ -545,7 +552,33 @@ proc select*(readfds: var seq[TSocket], timeout = 500): int =
     result = int(select(cint(m+1), addr(rd), nil, nil, nil))
   
   pruneSocketSet(readfds, (rd))
+  
+proc recv*(socket: TSocket, data: pointer, size: int): int =
+  ## receives data from a socket
+  result = recv(cint(socket), data, size, 0'i32)
 
+template waitFor(): stmt =
+  if timeout - int(waited * 1000.0) < 1:
+    raise newException(ETimedout, "Call to recv() timed out.")
+  var s = @[socket]
+  var startTime = epochTime()
+  if select(s, timeout - int(waited * 1000.0)) != 1:
+    raise newException(ETimedout, "Call to recv() timed out.")
+  waited += (epochTime() - startTime)
+
+proc recv*(socket: TSocket, data: var string, size: int, timeout: int): int =
+  ## overload with a ``timeout`` parameter in miliseconds.
+  var waited = 0.0 # number of seconds already waited  
+  
+  var read = 0
+  while read < size:
+    waitFor()
+    result = recv(cint(socket), addr(data[read]), 1, 0'i32)
+    if result < 0:
+      return
+    inc(read)
+  
+  result = read
 
 proc recvLine*(socket: TSocket, line: var TaintedString): bool =
   ## returns false if no further data is available. `Line` must be initialized
@@ -567,6 +600,29 @@ proc recvLine*(socket: TSocket, line: var TaintedString): bool =
     elif c == '\L': return true
     add(line.string, c)
 
+proc recvLine*(socket: TSocket, line: var TaintedString, timeout: int): bool =
+  ## variant with a ``timeout`` parameter, the timeout parameter specifies
+  ## how many miliseconds to wait for data.
+  
+  var waited = 0.0 # number of seconds already waited
+  
+  setLen(line.string, 0)
+  while true:
+    var c: char
+    waitFor()
+    var n = recv(cint(socket), addr(c), 1, 0'i32)
+    if n < 0: return
+    elif n == 0: return true
+    if c == '\r':
+      waitFor()
+      n = recv(cint(socket), addr(c), 1, MSG_PEEK)
+      if n > 0 and c == '\L':
+        discard recv(cint(socket), addr(c), 1, 0'i32)
+      elif n <= 0: return false
+      return true
+    elif c == '\L': return true
+    add(line.string, c)
+
 proc recvLineAsync*(socket: TSocket, line: var TaintedString): TRecvLineResult =
   ## similar to ``recvLine`` but for non-blocking sockets.
   ## The values of the returned enum should be pretty self explanatory:
@@ -592,10 +648,6 @@ proc recvLineAsync*(socket: TSocket, line: var TaintedString): TRecvLineResult =
     elif c == '\L': return RecvFullLine
     add(line.string, c)
 
-proc recv*(socket: TSocket, data: pointer, size: int): int =
-  ## receives data from a socket
-  result = recv(cint(socket), data, size, 0'i32)
-
 proc recv*(socket: TSocket): TaintedString =
   ## receives all the data from the socket.
   ## Socket errors will result in an ``EOS`` error.
@@ -625,6 +677,16 @@ proc recv*(socket: TSocket): TaintedString =
       add(result.string, buf)
       if bytesRead != bufSize-1: break
 
+proc recvTimeout*(socket: TSocket, timeout: int): TaintedString =
+  ## overloaded variant to support a ``timeout`` parameter, the ``timeout``
+  ## parameter specifies the amount of miliseconds to wait for data on the
+  ## socket.
+  var s = @[socket]
+  if s.select(timeout) != 1:
+    raise newException(ETimedout, "Call to recv() timed out.")
+  
+  return socket.recv
+
 proc recvAsync*(socket: TSocket, s: var TaintedString): bool =
   ## receives all the data from a non-blocking socket. If socket is non-blocking 
   ## and there are no messages available, `False` will be returned.
@@ -723,6 +785,21 @@ proc setBlocking*(s: TSocket, blocking: bool) =
       if fcntl(cint(s), F_SETFL, mode) == -1:
         OSError()
 
+proc connect*(socket: TSocket, timeout: int, name: string, port = TPort(0),
+             af: TDomain = AF_INET) =
+  ## Overload for ``connect`` to support timeouts. The ``timeout`` parameter 
+  ## specifies the time in miliseconds of how long to wait for a connection
+  ## to be made.
+  ##
+  ## **Warning:** If ``socket`` is non-blocking and timeout is not ``-1`` then
+  ## this function may set blocking mode on ``socket`` to true.
+  socket.setBlocking(true)
+  
+  socket.connectAsync(name, port, af)
+  var s: seq[TSocket] = @[socket]
+  if selectWrite(s, timeout) != 1:
+    raise newException(ETimedout, "Call to connect() timed out.")
+
 when defined(Windows):
   var wsa: TWSADATA
   if WSAStartup(0x0101'i16, wsa) != 0: OSError()