sockettimeout.r2py

"""
<Author>
  Justin Cappos, Armon Dadgar
  This is a rewrite of the previous version by Richard Jordan

<Start Date>
  26 Aug 2009

<Description>
  A library that causes sockets to timeout if a recv / send call would
  block for more than an allotted amount of time.

"""
random = dy_import_module("random.r2py")

class SocketTimeoutError(Exception):
  """The socket timed out before receiving a response"""


class timeout_socket():
  """
  <Purpose>
    Provides a socket like object which supports custom timeouts
    for send() and recv().
  """

  # Initialize with the socket object and a default timeout
  def __init__(self,socket,timeout=10, checkintv=0.1):
    """
    <Purpose>
      Initializes a timeout socket object.

    <Arguments>
      socket:
              A socket like object to wrap. Must support send,recv,close.

      timeout:
              The default timeout for send() and recv().

      checkintv:
              How often socket operations (send,recv) should check if
              they can run. The smaller the interval the more time is
              spent busy waiting.
    """
    # Store the socket, timeout and check interval
    self.socket = socket
    self.timeout = timeout
    self.checkintv = checkintv


  # Allow changing the default timeout
  def settimeout(self,timeout=10):
    """
    <Purpose>
      Allows changing the default timeout interval.

    <Arguments>
      timeout:
              The new default timeout interval. Defaults to 10.
              Use 0 for no timeout. Given in seconds.

    """
    # Update
    self.timeout = timeout
  
  
  # Wrap close
  def close(self):
    """
    See socket.close()
    """
    return self.socket.close()


  # Provide a recv() implementation
  def recv(self,bytes,timeout=None):
    """
    <Purpose>
      Allows receiving data from the socket object with a custom timeout.

    <Arguments>
      bytes:
          The maximum amount of bytes to read

      timeout:
          (Optional) Defaults to the value given at initialization, or by settimeout.
          If provided, the socket operation will timeout after this amount of time (sec).
          Use 0 for no timeout.

    <Exceptions>
      As with socket.recv(), socket.willblock(). Additionally, SocketTimeoutError is
      raised if the operation times out.

    <Returns>
      The data received from the socket.
    """
    # Set the timeout if None
    if timeout is None:
      timeout = self.timeout

    # Get the start time
    starttime = getruntime()

    elapsed_time = 0

    while elapsed_time < timeout:
      try:
        data = self.socket.recv(bytes)
      except SocketWouldBlockError:
        sleep(self.checkintv)
        elapsed_time = getruntime() - starttime
      else:
        break
    else:
      raise SocketTimeoutError, "recv() timed out!!"

    return data


  # Provide a send() implementation
  def send(self,data,timeout=None):
    """
    <Purpose>
      Allows sending data with the socket object with a custom timeout.

    <Arguments>
      data:
          The data to send

      timeout:
          (Optional) Defaults to the value given at initialization, or by settimeout.
          If provided, the socket operation will timeout after this amount of time (sec).
          Use 0 for no timeout.

    <Exceptions>
      As with socket.send(), socket.willblock(). Additionally, SocketTimeoutError is
      raised if the operation times out.

    <Returns>
      The number of bytes sent.
    """
    # Set the timeout if None
    if timeout is None:
      timeout = self.timeout

    # Get the start time
    starttime = getruntime()

    elapsed_time = 0

    while elapsed_time < timeout:
      try:
        sentbytes = self.socket.send(data)
      except SocketWouldBlockError:
        sleep(self.checkintv)
        elapsed_time = getruntime() - starttime
      else:
        break
    else:
      raise SocketTimeoutError, "send() timed out!!"

    return sentbytes


# Wrapper class for TCPServerSocket in emulcomm
class timeout_server_socket():
  """
  <Purpose>
    Provides a TCPSeverSocket like object which supports custom timeouts
    for getconnection().
  """

  def __init__(self, socket,timeout=10, checkintv=0.1):
    """
    <Purpose>
      Initializes a timeout TCPServerSocket object.

    <Arguments>
      socket:
              A socket like object to wrap. Must support getconnection, close.

      timeout:
              The default timeout for getconnection().

      checkintv:
              How often socket operations (getconnection) should check if
              they can run. The smaller the interval the more time is
              spent busy waiting.
    """
    self.socket = socket
    self.timeout = timeout
    self.checkintv = checkintv

  # Allow changing the default timeout
  def settimeout(self,timeout=10):
    """
    <Purpose>
      Allows changing the default timeout interval.

    <Arguments>
      timeout:
              The new default timeout interval. Defaults to 10.
              Use 0 for no timeout. Given in seconds.

    """
    # Update
    self.timeout = timeout

  def getconnection(self, timeout=None):
    """
    <Purpose>
      Allows accecpting new connection on TCPSeverSocket object with a custom 
      timeout.

    <Arguments>
      timeout:
          (Optional) Defaults to the value given at initialization, or by settimeout.
          If provided, the socket operation will timeout after this amount of time (sec).
          Use 0 for no timeout.

    <Exceptions>
      As with TCPSeverSocket.getconnection(). Additionally, SocketTimeoutError 
      is raised if the operation times out.

    <Returns>
      A 3-tuple consisting of remoteip, remoteport and client socket.
    """

    # Set the timeout if None
    if timeout is None:
      timeout = self.timeout

    # Get the start time.
    starttime = getruntime()

    elapsed_time = 0

    while elapsed_time < timeout:
      try:
        remoteip, remoteport, realsocketlikeobject = self.socket.getconnection()
      except SocketWouldBlockError:
        sleep(self.checkintv)
        elapsed_time = getruntime() - starttime
      else:
        break
    else:
      raise SocketTimeoutError, "getconnection() timed out!"

    thissocketlikeobject = timeout_socket(realsocketlikeobject, timeout)
    return (remoteip, remoteport, thissocketlikeobject)

  # Wrap close
  def close(self):
    """
    See socket.close()
    """
    self.socket.close()

    


def timeout_openconnection(desthost, destport, localip=None, localport=None, timeout=5):
  """
  <Purpose> 
    Wrapper for openconnection.   Very, very similar

  <Args>
    Similar to Repy's openconnection, but will try to automatically come 
    up with localip and localport if omitted.

  <Exception>
    Raises the same exceptions as openconnection.

  <Side Effects>
    Creates a socket object for the user

  <Returns>
    socket obj on success
  """
  # Initialize localip and localport (if not set by caller)
  if localip is None:
    localip = getmyip()

  if localport is None:
    # Get the ports available to our vessel (#1375)
    # (I'll try all of the nominally available ports, not only 
    # the ones that are currently unused, because port usage might 
    # change while we still set up things here.)
    available_resources, ignore, ignore = getresources()
    portrange = list(available_resources['connport'])
    # Shuffle the portrange to make it unlikely we ever reuse the 
    # same quintuple twice in a row (#1362)
    random.random_shuffle(portrange)
  else:
    portrange = [localport]

  for myport in portrange:
    try:
      realsocketlikeobject = openconnection(desthost, destport, localip, myport, timeout)
      # Heureka, we found a usable port. Continue outside of the for loop.
      break
    except (DuplicateTupleError, CleanupInProgressError, 
      AlreadyListeningError, AddressBindingError):
      # These are "legitimate" in the sense that there is something 
      # in our very program already using the port we just tried.
      pass
    except ResourceForbiddenError:
      # The port is not in the restrictions file. This can only 
      # happen if the user supplied us with the ``localport'' argument.
      raise
    except Exception:
      # E.g. RepyArgumentError due to using a FQDN instead of an 
      # IP address argument -- the user's fault, really.
      raise
  else:
    # Checked the whole port range without success
    raise AddressBindingError, "Could not find any usable source port in the range " + str(portrange) + " on IP address " + str(localip) + "."


  thissocketlikeobject = timeout_socket(realsocketlikeobject, timeout)
  return thissocketlikeobject





def timeout_listenforconnection(localip, localport, timeout=5):
  """
  <Purpose> 
    Wrapper for waitforconn.   Essentially does the same thing...

  <Args>
    Same as Repy waitforconn with the addition of a timeout argument.

  <Exceptions> 
    Same as Repy waitforconn

  <Side Effects>
    Sets up event listener which calls function on messages.

  <Returns>
    Handle to listener.
  """

  realsocketlikeobject = listenforconnection(localip, localport)

  thissocketlikeobject = timeout_server_socket(realsocketlikeobject, timeout)
  return thissocketlikeobject