source: python/trunk/Lib/multiprocessing/connection.py@ 388

Last change on this file since 388 was 2, checked in by Yuri Dario, 15 years ago

Initial import for vendor code.
  • Property svn:eol-style set to native
File size: 12.2 KB
Line 
1#
2# A higher level module for using sockets (or Windows named pipes)
3#
4# multiprocessing/connection.py
5#
6# Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
7#
8
9__all__ = [ 'Client', 'Listener', 'Pipe' ]
10
11import os
12import sys
13import socket
14import errno
15import time
16import tempfile
17import itertools
18
19import _multiprocessing
20from multiprocessing import current_process, AuthenticationError
21from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
22from multiprocessing.forking import duplicate, close
23
24
25#
26#
27#
28
29BUFSIZE = 8192
30# A very generous timeout when it comes to local connections...
31CONNECTION_TIMEOUT = 20.
32
33_mmap_counter = itertools.count()
34
35default_family = 'AF_INET'
36families = ['AF_INET']
37
38if hasattr(socket, 'AF_UNIX'):
39 default_family = 'AF_UNIX'
40 families += ['AF_UNIX']
41
42if sys.platform == 'win32':
43 default_family = 'AF_PIPE'
44 families += ['AF_PIPE']
45
46
47def _init_timeout(timeout=CONNECTION_TIMEOUT):
48 return time.time() + timeout
49
50def _check_timeout(t):
51 return time.time() > t
52
53#
54#
55#
56
57def arbitrary_address(family):
58 '''
59 Return an arbitrary free address for the given family
60 '''
61 if family == 'AF_INET':
62 return ('localhost', 0)
63 elif family == 'AF_UNIX':
64 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
65 elif family == 'AF_PIPE':
66 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
67 (os.getpid(), _mmap_counter.next()))
68 else:
69 raise ValueError('unrecognized family')
70
71
72def address_type(address):
73 '''
74 Return the types of the address
75
76 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
77 '''
78 if type(address) == tuple:
79 return 'AF_INET'
80 elif type(address) is str and address.startswith('\\\\'):
81 return 'AF_PIPE'
82 elif type(address) is str:
83 return 'AF_UNIX'
84 else:
85 raise ValueError('address type of %r unrecognized' % address)
86
87#
88# Public functions
89#
90
91class Listener(object):
92 '''
93 Returns a listener object.
94
95 This is a wrapper for a bound socket which is 'listening' for
96 connections, or for a Windows named pipe.
97 '''
98 def __init__(self, address=None, family=None, backlog=1, authkey=None):
99 family = family or (address and address_type(address)) \
100 or default_family
101 address = address or arbitrary_address(family)
102
103 if family == 'AF_PIPE':
104 self._listener = PipeListener(address, backlog)
105 else:
106 self._listener = SocketListener(address, family, backlog)
107
108 if authkey is not None and not isinstance(authkey, bytes):
109 raise TypeError, 'authkey should be a byte string'
110
111 self._authkey = authkey
112
113 def accept(self):
114 '''
115 Accept a connection on the bound socket or named pipe of `self`.
116
117 Returns a `Connection` object.
118 '''
119 c = self._listener.accept()
120 if self._authkey:
121 deliver_challenge(c, self._authkey)
122 answer_challenge(c, self._authkey)
123 return c
124
125 def close(self):
126 '''
127 Close the bound socket or named pipe of `self`.
128 '''
129 return self._listener.close()
130
131 address = property(lambda self: self._listener._address)
132 last_accepted = property(lambda self: self._listener._last_accepted)
133
134
135def Client(address, family=None, authkey=None):
136 '''
137 Returns a connection to the address of a `Listener`
138 '''
139 family = family or address_type(address)
140 if family == 'AF_PIPE':
141 c = PipeClient(address)
142 else:
143 c = SocketClient(address)
144
145 if authkey is not None and not isinstance(authkey, bytes):
146 raise TypeError, 'authkey should be a byte string'
147
148 if authkey is not None:
149 answer_challenge(c, authkey)
150 deliver_challenge(c, authkey)
151
152 return c
153
154
155if sys.platform != 'win32':
156
157 def Pipe(duplex=True):
158 '''
159 Returns pair of connection objects at either end of a pipe
160 '''
161 if duplex:
162 s1, s2 = socket.socketpair()
163 c1 = _multiprocessing.Connection(os.dup(s1.fileno()))
164 c2 = _multiprocessing.Connection(os.dup(s2.fileno()))
165 s1.close()
166 s2.close()
167 else:
168 fd1, fd2 = os.pipe()
169 c1 = _multiprocessing.Connection(fd1, writable=False)
170 c2 = _multiprocessing.Connection(fd2, readable=False)
171
172 return c1, c2
173
174else:
175
176 from ._multiprocessing import win32
177
178 def Pipe(duplex=True):
179 '''
180 Returns pair of connection objects at either end of a pipe
181 '''
182 address = arbitrary_address('AF_PIPE')
183 if duplex:
184 openmode = win32.PIPE_ACCESS_DUPLEX
185 access = win32.GENERIC_READ | win32.GENERIC_WRITE
186 obsize, ibsize = BUFSIZE, BUFSIZE
187 else:
188 openmode = win32.PIPE_ACCESS_INBOUND
189 access = win32.GENERIC_WRITE
190 obsize, ibsize = 0, BUFSIZE
191
192 h1 = win32.CreateNamedPipe(
193 address, openmode,
194 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
195 win32.PIPE_WAIT,
196 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
197 )
198 h2 = win32.CreateFile(
199 address, access, 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
200 )
201 win32.SetNamedPipeHandleState(
202 h2, win32.PIPE_READMODE_MESSAGE, None, None
203 )
204
205 try:
206 win32.ConnectNamedPipe(h1, win32.NULL)
207 except WindowsError, e:
208 if e.args[0] != win32.ERROR_PIPE_CONNECTED:
209 raise
210
211 c1 = _multiprocessing.PipeConnection(h1, writable=duplex)
212 c2 = _multiprocessing.PipeConnection(h2, readable=duplex)
213
214 return c1, c2
215
216#
217# Definitions for connections based on sockets
218#
219
220class SocketListener(object):
221 '''
222 Representation of a socket which is bound to an address and listening
223 '''
224 def __init__(self, address, family, backlog=1):
225 self._socket = socket.socket(getattr(socket, family))
226 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
227 self._socket.bind(address)
228 self._socket.listen(backlog)
229 self._address = self._socket.getsockname()
230 self._family = family
231 self._last_accepted = None
232
233 if family == 'AF_UNIX':
234 self._unlink = Finalize(
235 self, os.unlink, args=(address,), exitpriority=0
236 )
237 else:
238 self._unlink = None
239
240 def accept(self):
241 s, self._last_accepted = self._socket.accept()
242 fd = duplicate(s.fileno())
243 conn = _multiprocessing.Connection(fd)
244 s.close()
245 return conn
246
247 def close(self):
248 self._socket.close()
249 if self._unlink is not None:
250 self._unlink()
251
252
253def SocketClient(address):
254 '''
255 Return a connection object connected to the socket given by `address`
256 '''
257 family = address_type(address)
258 s = socket.socket( getattr(socket, family) )
259 t = _init_timeout()
260
261 while 1:
262 try:
263 s.connect(address)
264 except socket.error, e:
265 if e.args[0] != errno.ECONNREFUSED or _check_timeout(t):
266 debug('failed to connect to address %s', address)
267 raise
268 time.sleep(0.01)
269 else:
270 break
271 else:
272 raise
273
274 fd = duplicate(s.fileno())
275 conn = _multiprocessing.Connection(fd)
276 s.close()
277 return conn
278
279#
280# Definitions for connections based on named pipes
281#
282
283if sys.platform == 'win32':
284
285 class PipeListener(object):
286 '''
287 Representation of a named pipe
288 '''
289 def __init__(self, address, backlog=None):
290 self._address = address
291 handle = win32.CreateNamedPipe(
292 address, win32.PIPE_ACCESS_DUPLEX,
293 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
294 win32.PIPE_WAIT,
295 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
296 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
297 )
298 self._handle_queue = [handle]
299 self._last_accepted = None
300
301 sub_debug('listener created with address=%r', self._address)
302
303 self.close = Finalize(
304 self, PipeListener._finalize_pipe_listener,
305 args=(self._handle_queue, self._address), exitpriority=0
306 )
307
308 def accept(self):
309 newhandle = win32.CreateNamedPipe(
310 self._address, win32.PIPE_ACCESS_DUPLEX,
311 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
312 win32.PIPE_WAIT,
313 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
314 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
315 )
316 self._handle_queue.append(newhandle)
317 handle = self._handle_queue.pop(0)
318 try:
319 win32.ConnectNamedPipe(handle, win32.NULL)
320 except WindowsError, e:
321 if e.args[0] != win32.ERROR_PIPE_CONNECTED:
322 raise
323 return _multiprocessing.PipeConnection(handle)
324
325 @staticmethod
326 def _finalize_pipe_listener(queue, address):
327 sub_debug('closing listener with address=%r', address)
328 for handle in queue:
329 close(handle)
330
331 def PipeClient(address):
332 '''
333 Return a connection object connected to the pipe given by `address`
334 '''
335 t = _init_timeout()
336 while 1:
337 try:
338 win32.WaitNamedPipe(address, 1000)
339 h = win32.CreateFile(
340 address, win32.GENERIC_READ | win32.GENERIC_WRITE,
341 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
342 )
343 except WindowsError, e:
344 if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
345 win32.ERROR_PIPE_BUSY) or _check_timeout(t):
346 raise
347 else:
348 break
349 else:
350 raise
351
352 win32.SetNamedPipeHandleState(
353 h, win32.PIPE_READMODE_MESSAGE, None, None
354 )
355 return _multiprocessing.PipeConnection(h)
356
357#
358# Authentication stuff
359#
360
361MESSAGE_LENGTH = 20
362
363CHALLENGE = b'#CHALLENGE#'
364WELCOME = b'#WELCOME#'
365FAILURE = b'#FAILURE#'
366
367def deliver_challenge(connection, authkey):
368 import hmac
369 assert isinstance(authkey, bytes)
370 message = os.urandom(MESSAGE_LENGTH)
371 connection.send_bytes(CHALLENGE + message)
372 digest = hmac.new(authkey, message).digest()
373 response = connection.recv_bytes(256) # reject large message
374 if response == digest:
375 connection.send_bytes(WELCOME)
376 else:
377 connection.send_bytes(FAILURE)
378 raise AuthenticationError('digest received was wrong')
379
380def answer_challenge(connection, authkey):
381 import hmac
382 assert isinstance(authkey, bytes)
383 message = connection.recv_bytes(256) # reject large message
384 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
385 message = message[len(CHALLENGE):]
386 digest = hmac.new(authkey, message).digest()
387 connection.send_bytes(digest)
388 response = connection.recv_bytes(256) # reject large message
389 if response != WELCOME:
390 raise AuthenticationError('digest sent was rejected')
391
392#
393# Support for using xmlrpclib for serialization
394#
395
396class ConnectionWrapper(object):
397 def __init__(self, conn, dumps, loads):
398 self._conn = conn
399 self._dumps = dumps
400 self._loads = loads
401 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
402 obj = getattr(conn, attr)
403 setattr(self, attr, obj)
404 def send(self, obj):
405 s = self._dumps(obj)
406 self._conn.send_bytes(s)
407 def recv(self):
408 s = self._conn.recv_bytes()
409 return self._loads(s)
410
411def _xml_dumps(obj):
412 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf8')
413
414def _xml_loads(s):
415 (obj,), method = xmlrpclib.loads(s.decode('utf8'))
416 return obj
417
418class XmlListener(Listener):
419 def accept(self):
420 global xmlrpclib
421 import xmlrpclib
422 obj = Listener.accept(self)
423 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
424
425def XmlClient(*args, **kwds):
426 global xmlrpclib
427 import xmlrpclib
428 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Note: See TracBrowser for help on using the repository browser.