2009-03-13 Xan Lopez <xlopez@igalia.com>
[WebKit-https.git] / PlanetWebKit / planet / planet / timeoutsocket.py
1
2 ####
3 # Copyright 2000,2001 by Timothy O'Malley <timo@alum.mit.edu>
4
5 #                All Rights Reserved
6
7 # Permission to use, copy, modify, and distribute this software
8 # and its documentation for any purpose and without fee is hereby
9 # granted, provided that the above copyright notice appear in all
10 # copies and that both that copyright notice and this permission
11 # notice appear in supporting documentation, and that the name of
12 # Timothy O'Malley  not be used in advertising or publicity
13 # pertaining to distribution of the software without specific, written
14 # prior permission. 
15
16 # Timothy O'Malley DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
17 # SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
18 # AND FITNESS, IN NO EVENT SHALL Timothy O'Malley BE LIABLE FOR
19 # ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
20 # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
21 # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
22 # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
23 # PERFORMANCE OF THIS SOFTWARE. 
24 #
25 ####
26
27 """Timeout Socket
28
29 This module enables a timeout mechanism on all TCP connections.  It
30 does this by inserting a shim into the socket module.  After this module
31 has been imported, all socket creation goes through this shim.  As a
32 result, every TCP connection will support a timeout.
33
34 The beauty of this method is that it immediately and transparently
35 enables the entire python library to support timeouts on TCP sockets.
36 As an example, if you wanted to SMTP connections to have a 20 second
37 timeout:
38
39     import timeoutsocket
40     import smtplib
41     timeoutsocket.setDefaultSocketTimeout(20)
42
43
44 The timeout applies to the socket functions that normally block on
45 execution:  read, write, connect, and accept.  If any of these 
46 operations exceeds the specified timeout, the exception Timeout
47 will be raised.
48
49 The default timeout value is set to None.  As a result, importing
50 this module does not change the default behavior of a socket.  The
51 timeout mechanism only activates when the timeout has been set to
52 a numeric value.  (This behavior mimics the behavior of the
53 select.select() function.)
54
55 This module implements two classes: TimeoutSocket and TimeoutFile.
56
57 The TimeoutSocket class defines a socket-like object that attempts to
58 avoid the condition where a socket may block indefinitely.  The
59 TimeoutSocket class raises a Timeout exception whenever the
60 current operation delays too long. 
61
62 The TimeoutFile class defines a file-like object that uses the TimeoutSocket
63 class.  When the makefile() method of TimeoutSocket is called, it returns
64 an instance of a TimeoutFile.
65
66 Each of these objects adds two methods to manage the timeout value:
67
68     get_timeout()   -->  returns the timeout of the socket or file
69     set_timeout()   -->  sets the timeout of the socket or file
70
71
72 As an example, one might use the timeout feature to create httplib
73 connections that will timeout after 30 seconds:
74
75     import timeoutsocket
76     import httplib
77     H = httplib.HTTP("www.python.org")
78     H.sock.set_timeout(30)
79
80 Note:  When used in this manner, the connect() routine may still
81 block because it happens before the timeout is set.  To avoid
82 this, use the 'timeoutsocket.setDefaultSocketTimeout()' function.
83
84 Good Luck!
85
86 """
87
88 __version__ = "$Revision: 1.1.1.1 $"
89 __author__  = "Timothy O'Malley <timo@alum.mit.edu>"
90
91 #
92 # Imports
93 #
94 import select, string
95 import socket
96 if not hasattr(socket, "_no_timeoutsocket"):
97     _socket = socket.socket
98 else:
99     _socket = socket._no_timeoutsocket
100
101
102 #
103 # Set up constants to test for Connected and Blocking operations.
104 # We delete 'os' and 'errno' to keep our namespace clean(er).
105 # Thanks to Alex Martelli and G. Li for the Windows error codes.
106 #
107 import os
108 if os.name == "nt":
109     _IsConnected = ( 10022, 10056 )
110     _ConnectBusy = ( 10035, )
111     _AcceptBusy  = ( 10035, )
112 else:
113     import errno
114     _IsConnected = ( errno.EISCONN, )
115     _ConnectBusy = ( errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK )
116     _AcceptBusy  = ( errno.EAGAIN, errno.EWOULDBLOCK )
117     del errno
118 del os
119
120
121 #
122 # Default timeout value for ALL TimeoutSockets
123 #
124 _DefaultTimeout = None
125 def setDefaultSocketTimeout(timeout):
126     global _DefaultTimeout
127     _DefaultTimeout = timeout
128 def getDefaultSocketTimeout():
129     return _DefaultTimeout
130
131 #
132 # Exceptions for socket errors and timeouts
133 #
134 Error = socket.error
135 class Timeout(Exception):
136     pass
137
138
139 #
140 # Factory function
141 #
142 from socket import AF_INET, SOCK_STREAM
143 def timeoutsocket(family=AF_INET, type=SOCK_STREAM, proto=None):
144     if family != AF_INET or type != SOCK_STREAM:
145         if proto:
146             return _socket(family, type, proto)
147         else:
148             return _socket(family, type)
149     return TimeoutSocket( _socket(family, type), _DefaultTimeout )
150 # end timeoutsocket
151
152 #
153 # The TimeoutSocket class definition
154 #
155 class TimeoutSocket:
156     """TimeoutSocket object
157     Implements a socket-like object that raises Timeout whenever
158     an operation takes too long.
159     The definition of 'too long' can be changed using the
160     set_timeout() method.
161     """
162
163     _copies = 0
164     _blocking = 1
165     
166     def __init__(self, sock, timeout):
167         self._sock     = sock
168         self._timeout  = timeout
169     # end __init__
170
171     def __getattr__(self, key):
172         return getattr(self._sock, key)
173     # end __getattr__
174
175     def get_timeout(self):
176         return self._timeout
177     # end set_timeout
178
179     def set_timeout(self, timeout=None):
180         self._timeout = timeout
181     # end set_timeout
182
183     def setblocking(self, blocking):
184         self._blocking = blocking
185         return self._sock.setblocking(blocking)
186     # end set_timeout
187
188     def connect_ex(self, addr):
189         errcode = 0
190         try:
191             self.connect(addr)
192         except Error, why:
193             errcode = why[0]
194         return errcode
195     # end connect_ex
196         
197     def connect(self, addr, port=None, dumbhack=None):
198         # In case we were called as connect(host, port)
199         if port != None:  addr = (addr, port)
200
201         # Shortcuts
202         sock    = self._sock
203         timeout = self._timeout
204         blocking = self._blocking
205
206         # First, make a non-blocking call to connect
207         try:
208             sock.setblocking(0)
209             sock.connect(addr)
210             sock.setblocking(blocking)
211             return
212         except Error, why:
213             # Set the socket's blocking mode back
214             sock.setblocking(blocking)
215             
216             # If we are not blocking, re-raise
217             if not blocking:
218                 raise
219             
220             # If we are already connected, then return success.
221             # If we got a genuine error, re-raise it.
222             errcode = why[0]
223             if dumbhack and errcode in _IsConnected:
224                 return
225             elif errcode not in _ConnectBusy:
226                 raise
227             
228         # Now, wait for the connect to happen
229         # ONLY if dumbhack indicates this is pass number one.
230         #   If select raises an error, we pass it on.
231         #   Is this the right behavior?
232         if not dumbhack:
233             r,w,e = select.select([], [sock], [], timeout)
234             if w:
235                 return self.connect(addr, dumbhack=1)
236
237         # If we get here, then we should raise Timeout
238         raise Timeout("Attempted connect to %s timed out." % str(addr) )
239     # end connect
240
241     def accept(self, dumbhack=None):
242         # Shortcuts
243         sock     = self._sock
244         timeout  = self._timeout
245         blocking = self._blocking
246
247         # First, make a non-blocking call to accept
248         #  If we get a valid result, then convert the
249         #  accept'ed socket into a TimeoutSocket.
250         # Be carefult about the blocking mode of ourselves.
251         try:
252             sock.setblocking(0)
253             newsock, addr = sock.accept()
254             sock.setblocking(blocking)
255             timeoutnewsock = self.__class__(newsock, timeout)
256             timeoutnewsock.setblocking(blocking)
257             return (timeoutnewsock, addr)
258         except Error, why:
259             # Set the socket's blocking mode back
260             sock.setblocking(blocking)
261
262             # If we are not supposed to block, then re-raise
263             if not blocking:
264                 raise
265             
266             # If we got a genuine error, re-raise it.
267             errcode = why[0]
268             if errcode not in _AcceptBusy:
269                 raise
270             
271         # Now, wait for the accept to happen
272         # ONLY if dumbhack indicates this is pass number one.
273         #   If select raises an error, we pass it on.
274         #   Is this the right behavior?
275         if not dumbhack:
276             r,w,e = select.select([sock], [], [], timeout)
277             if r:
278                 return self.accept(dumbhack=1)
279
280         # If we get here, then we should raise Timeout
281         raise Timeout("Attempted accept timed out.")
282     # end accept
283
284     def send(self, data, flags=0):
285         sock = self._sock
286         if self._blocking:
287             r,w,e = select.select([],[sock],[], self._timeout)
288             if not w:
289                 raise Timeout("Send timed out")
290         return sock.send(data, flags)
291     # end send
292
293     def recv(self, bufsize, flags=0):
294         sock = self._sock
295         if self._blocking:
296             r,w,e = select.select([sock], [], [], self._timeout)
297             if not r:
298                 raise Timeout("Recv timed out")
299         return sock.recv(bufsize, flags)
300     # end recv
301
302     def makefile(self, flags="r", bufsize=-1):
303         self._copies = self._copies +1
304         return TimeoutFile(self, flags, bufsize)
305     # end makefile
306
307     def close(self):
308         if self._copies <= 0:
309             self._sock.close()
310         else:
311             self._copies = self._copies -1
312     # end close
313
314 # end TimeoutSocket
315
316
317 class TimeoutFile:
318     """TimeoutFile object
319     Implements a file-like object on top of TimeoutSocket.
320     """
321     
322     def __init__(self, sock, mode="r", bufsize=4096):
323         self._sock          = sock
324         self._bufsize       = 4096
325         if bufsize > 0: self._bufsize = bufsize
326         if not hasattr(sock, "_inqueue"): self._sock._inqueue = ""
327
328     # end __init__
329
330     def __getattr__(self, key):
331         return getattr(self._sock, key)
332     # end __getattr__
333
334     def close(self):
335         self._sock.close()
336         self._sock = None
337     # end close
338     
339     def write(self, data):
340         self.send(data)
341     # end write
342
343     def read(self, size=-1):
344         _sock = self._sock
345         _bufsize = self._bufsize
346         while 1:
347             datalen = len(_sock._inqueue)
348             if datalen >= size >= 0:
349                 break
350             bufsize = _bufsize
351             if size > 0:
352                 bufsize = min(bufsize, size - datalen )
353             buf = self.recv(bufsize)
354             if not buf:
355                 break
356             _sock._inqueue = _sock._inqueue + buf
357         data = _sock._inqueue
358         _sock._inqueue = ""
359         if size > 0 and datalen > size:
360             _sock._inqueue = data[size:]
361             data = data[:size]
362         return data
363     # end read
364
365     def readline(self, size=-1):
366         _sock = self._sock
367         _bufsize = self._bufsize
368         while 1:
369             idx = string.find(_sock._inqueue, "\n")
370             if idx >= 0:
371                 break
372             datalen = len(_sock._inqueue)
373             if datalen >= size >= 0:
374                 break
375             bufsize = _bufsize
376             if size > 0:
377                 bufsize = min(bufsize, size - datalen )
378             buf = self.recv(bufsize)
379             if not buf:
380                 break
381             _sock._inqueue = _sock._inqueue + buf
382
383         data = _sock._inqueue
384         _sock._inqueue = ""
385         if idx >= 0:
386             idx = idx + 1
387             _sock._inqueue = data[idx:]
388             data = data[:idx]
389         elif size > 0 and datalen > size:
390             _sock._inqueue = data[size:]
391             data = data[:size]
392         return data
393     # end readline
394
395     def readlines(self, sizehint=-1):
396         result = []
397         data = self.read()
398         while data:
399             idx = string.find(data, "\n")
400             if idx >= 0:
401                 idx = idx + 1
402                 result.append( data[:idx] )
403                 data = data[idx:]
404             else:
405                 result.append( data )
406                 data = ""
407         return result
408     # end readlines
409
410     def flush(self):  pass
411
412 # end TimeoutFile
413
414
415 #
416 # Silently replace the socket() builtin function with
417 # our timeoutsocket() definition.
418 #
419 if not hasattr(socket, "_no_timeoutsocket"):
420     socket._no_timeoutsocket = socket.socket
421     socket.socket = timeoutsocket
422 del socket
423 socket = timeoutsocket
424 # Finis