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

Last change on this file since 1538 was 391, checked in by dmik, 12 years ago

python: Merge vendor 2.7.6 to trunk.
  • Property svn:eol-style set to native
File size: 14.3 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
7# All rights reserved.
8#
9# Redistribution and use in source and binary forms, with or without
10# modification, are permitted provided that the following conditions
11# are met:
12#
13# 1. Redistributions of source code must retain the above copyright
14# notice, this list of conditions and the following disclaimer.
15# 2. Redistributions in binary form must reproduce the above copyright
16# notice, this list of conditions and the following disclaimer in the
17# documentation and/or other materials provided with the distribution.
18# 3. Neither the name of author nor the names of any contributors may be
19# used to endorse or promote products derived from this software
20# without specific prior written permission.
21#
22# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
23# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
24# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
25# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
26# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
27# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
28# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
29# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
30# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
31# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
32# SUCH DAMAGE.
33#
34
35__all__ = [ 'Client', 'Listener', 'Pipe' ]
36
37import os
38import sys
39import socket
40import errno
41import time
42import tempfile
43import itertools
44
45import _multiprocessing
46from multiprocessing import current_process, AuthenticationError
47from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
48from multiprocessing.forking import duplicate, close
49
50
51#
52#
53#
54
55BUFSIZE = 8192
56# A very generous timeout when it comes to local connections...
57CONNECTION_TIMEOUT = 20.
58
59_mmap_counter = itertools.count()
60
61default_family = 'AF_INET'
62families = ['AF_INET']
63
64if hasattr(socket, 'AF_UNIX'):
65 default_family = 'AF_UNIX'
66 families += ['AF_UNIX']
67
68if sys.platform == 'win32':
69 default_family = 'AF_PIPE'
70 families += ['AF_PIPE']
71
72
73def _init_timeout(timeout=CONNECTION_TIMEOUT):
74 return time.time() + timeout
75
76def _check_timeout(t):
77 return time.time() > t
78
79#
80#
81#
82
83def arbitrary_address(family):
84 '''
85 Return an arbitrary free address for the given family
86 '''
87 if family == 'AF_INET':
88 return ('localhost', 0)
89 elif family == 'AF_UNIX':
90 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
91 elif family == 'AF_PIPE':
92 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
93 (os.getpid(), _mmap_counter.next()))
94 else:
95 raise ValueError('unrecognized family')
96
97
98def address_type(address):
99 '''
100 Return the types of the address
101
102 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
103 '''
104 if type(address) == tuple:
105 return 'AF_INET'
106 elif type(address) is str and address.startswith('\\\\'):
107 return 'AF_PIPE'
108 elif type(address) is str:
109 return 'AF_UNIX'
110 else:
111 raise ValueError('address type of %r unrecognized' % address)
112
113#
114# Public functions
115#
116
117class Listener(object):
118 '''
119 Returns a listener object.
120
121 This is a wrapper for a bound socket which is 'listening' for
122 connections, or for a Windows named pipe.
123 '''
124 def __init__(self, address=None, family=None, backlog=1, authkey=None):
125 family = family or (address and address_type(address)) \
126 or default_family
127 address = address or arbitrary_address(family)
128
129 if family == 'AF_PIPE':
130 self._listener = PipeListener(address, backlog)
131 else:
132 self._listener = SocketListener(address, family, backlog)
133
134 if authkey is not None and not isinstance(authkey, bytes):
135 raise TypeError, 'authkey should be a byte string'
136
137 self._authkey = authkey
138
139 def accept(self):
140 '''
141 Accept a connection on the bound socket or named pipe of `self`.
142
143 Returns a `Connection` object.
144 '''
145 c = self._listener.accept()
146 if self._authkey:
147 deliver_challenge(c, self._authkey)
148 answer_challenge(c, self._authkey)
149 return c
150
151 def close(self):
152 '''
153 Close the bound socket or named pipe of `self`.
154 '''
155 return self._listener.close()
156
157 address = property(lambda self: self._listener._address)
158 last_accepted = property(lambda self: self._listener._last_accepted)
159
160
161def Client(address, family=None, authkey=None):
162 '''
163 Returns a connection to the address of a `Listener`
164 '''
165 family = family or address_type(address)
166 if family == 'AF_PIPE':
167 c = PipeClient(address)
168 else:
169 c = SocketClient(address)
170
171 if authkey is not None and not isinstance(authkey, bytes):
172 raise TypeError, 'authkey should be a byte string'
173
174 if authkey is not None:
175 answer_challenge(c, authkey)
176 deliver_challenge(c, authkey)
177
178 return c
179
180
181if sys.platform != 'win32':
182
183 def Pipe(duplex=True):
184 '''
185 Returns pair of connection objects at either end of a pipe
186 '''
187 if duplex:
188 s1, s2 = socket.socketpair()
189 s1.setblocking(True)
190 s2.setblocking(True)
191 c1 = _multiprocessing.Connection(os.dup(s1.fileno()))
192 c2 = _multiprocessing.Connection(os.dup(s2.fileno()))
193 s1.close()
194 s2.close()
195 else:
196 fd1, fd2 = os.pipe()
197 c1 = _multiprocessing.Connection(fd1, writable=False)
198 c2 = _multiprocessing.Connection(fd2, readable=False)
199
200 return c1, c2
201
202else:
203 from _multiprocessing import win32
204
205 def Pipe(duplex=True):
206 '''
207 Returns pair of connection objects at either end of a pipe
208 '''
209 address = arbitrary_address('AF_PIPE')
210 if duplex:
211 openmode = win32.PIPE_ACCESS_DUPLEX
212 access = win32.GENERIC_READ | win32.GENERIC_WRITE
213 obsize, ibsize = BUFSIZE, BUFSIZE
214 else:
215 openmode = win32.PIPE_ACCESS_INBOUND
216 access = win32.GENERIC_WRITE
217 obsize, ibsize = 0, BUFSIZE
218
219 h1 = win32.CreateNamedPipe(
220 address, openmode,
221 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
222 win32.PIPE_WAIT,
223 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
224 )
225 h2 = win32.CreateFile(
226 address, access, 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
227 )
228 win32.SetNamedPipeHandleState(
229 h2, win32.PIPE_READMODE_MESSAGE, None, None
230 )
231
232 try:
233 win32.ConnectNamedPipe(h1, win32.NULL)
234 except WindowsError, e:
235 if e.args[0] != win32.ERROR_PIPE_CONNECTED:
236 raise
237
238 c1 = _multiprocessing.PipeConnection(h1, writable=duplex)
239 c2 = _multiprocessing.PipeConnection(h2, readable=duplex)
240
241 return c1, c2
242
243#
244# Definitions for connections based on sockets
245#
246
247class SocketListener(object):
248 '''
249 Representation of a socket which is bound to an address and listening
250 '''
251 def __init__(self, address, family, backlog=1):
252 self._socket = socket.socket(getattr(socket, family))
253 try:
254 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
255 self._socket.setblocking(True)
256 self._socket.bind(address)
257 self._socket.listen(backlog)
258 self._address = self._socket.getsockname()
259 except socket.error:
260 self._socket.close()
261 raise
262 self._family = family
263 self._last_accepted = None
264
265 if family == 'AF_UNIX':
266 self._unlink = Finalize(
267 self, os.unlink, args=(address,), exitpriority=0
268 )
269 else:
270 self._unlink = None
271
272 def accept(self):
273 while True:
274 try:
275 s, self._last_accepted = self._socket.accept()
276 except socket.error as e:
277 if e.args[0] != errno.EINTR:
278 raise
279 else:
280 break
281 s.setblocking(True)
282 fd = duplicate(s.fileno())
283 conn = _multiprocessing.Connection(fd)
284 s.close()
285 return conn
286
287 def close(self):
288 self._socket.close()
289 if self._unlink is not None:
290 self._unlink()
291
292
293def SocketClient(address):
294 '''
295 Return a connection object connected to the socket given by `address`
296 '''
297 family = getattr(socket, address_type(address))
298 t = _init_timeout()
299
300 while 1:
301 s = socket.socket(family)
302 s.setblocking(True)
303 try:
304 s.connect(address)
305 except socket.error, e:
306 s.close()
307 if e.args[0] != errno.ECONNREFUSED or _check_timeout(t):
308 debug('failed to connect to address %s', address)
309 raise
310 time.sleep(0.01)
311 else:
312 break
313 else:
314 raise
315
316 fd = duplicate(s.fileno())
317 conn = _multiprocessing.Connection(fd)
318 s.close()
319 return conn
320
321#
322# Definitions for connections based on named pipes
323#
324
325if sys.platform == 'win32':
326
327 class PipeListener(object):
328 '''
329 Representation of a named pipe
330 '''
331 def __init__(self, address, backlog=None):
332 self._address = address
333 handle = win32.CreateNamedPipe(
334 address, win32.PIPE_ACCESS_DUPLEX,
335 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
336 win32.PIPE_WAIT,
337 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
338 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
339 )
340 self._handle_queue = [handle]
341 self._last_accepted = None
342
343 sub_debug('listener created with address=%r', self._address)
344
345 self.close = Finalize(
346 self, PipeListener._finalize_pipe_listener,
347 args=(self._handle_queue, self._address), exitpriority=0
348 )
349
350 def accept(self):
351 newhandle = win32.CreateNamedPipe(
352 self._address, win32.PIPE_ACCESS_DUPLEX,
353 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
354 win32.PIPE_WAIT,
355 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
356 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
357 )
358 self._handle_queue.append(newhandle)
359 handle = self._handle_queue.pop(0)
360 try:
361 win32.ConnectNamedPipe(handle, win32.NULL)
362 except WindowsError, e:
363 # ERROR_NO_DATA can occur if a client has already connected,
364 # written data and then disconnected -- see Issue 14725.
365 if e.args[0] not in (win32.ERROR_PIPE_CONNECTED,
366 win32.ERROR_NO_DATA):
367 raise
368 return _multiprocessing.PipeConnection(handle)
369
370 @staticmethod
371 def _finalize_pipe_listener(queue, address):
372 sub_debug('closing listener with address=%r', address)
373 for handle in queue:
374 close(handle)
375
376 def PipeClient(address):
377 '''
378 Return a connection object connected to the pipe given by `address`
379 '''
380 t = _init_timeout()
381 while 1:
382 try:
383 win32.WaitNamedPipe(address, 1000)
384 h = win32.CreateFile(
385 address, win32.GENERIC_READ | win32.GENERIC_WRITE,
386 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
387 )
388 except WindowsError, e:
389 if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
390 win32.ERROR_PIPE_BUSY) or _check_timeout(t):
391 raise
392 else:
393 break
394 else:
395 raise
396
397 win32.SetNamedPipeHandleState(
398 h, win32.PIPE_READMODE_MESSAGE, None, None
399 )
400 return _multiprocessing.PipeConnection(h)
401
402#
403# Authentication stuff
404#
405
406MESSAGE_LENGTH = 20
407
408CHALLENGE = b'#CHALLENGE#'
409WELCOME = b'#WELCOME#'
410FAILURE = b'#FAILURE#'
411
412def deliver_challenge(connection, authkey):
413 import hmac
414 assert isinstance(authkey, bytes)
415 message = os.urandom(MESSAGE_LENGTH)
416 connection.send_bytes(CHALLENGE + message)
417 digest = hmac.new(authkey, message).digest()
418 response = connection.recv_bytes(256) # reject large message
419 if response == digest:
420 connection.send_bytes(WELCOME)
421 else:
422 connection.send_bytes(FAILURE)
423 raise AuthenticationError('digest received was wrong')
424
425def answer_challenge(connection, authkey):
426 import hmac
427 assert isinstance(authkey, bytes)
428 message = connection.recv_bytes(256) # reject large message
429 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
430 message = message[len(CHALLENGE):]
431 digest = hmac.new(authkey, message).digest()
432 connection.send_bytes(digest)
433 response = connection.recv_bytes(256) # reject large message
434 if response != WELCOME:
435 raise AuthenticationError('digest sent was rejected')
436
437#
438# Support for using xmlrpclib for serialization
439#
440
441class ConnectionWrapper(object):
442 def __init__(self, conn, dumps, loads):
443 self._conn = conn
444 self._dumps = dumps
445 self._loads = loads
446 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
447 obj = getattr(conn, attr)
448 setattr(self, attr, obj)
449 def send(self, obj):
450 s = self._dumps(obj)
451 self._conn.send_bytes(s)
452 def recv(self):
453 s = self._conn.recv_bytes()
454 return self._loads(s)
455
456def _xml_dumps(obj):
457 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf8')
458
459def _xml_loads(s):
460 (obj,), method = xmlrpclib.loads(s.decode('utf8'))
461 return obj
462
463class XmlListener(Listener):
464 def accept(self):
465 global xmlrpclib
466 import xmlrpclib
467 obj = Listener.accept(self)
468 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
469
470def XmlClient(*args, **kwds):
471 global xmlrpclib
472 import xmlrpclib
473 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Note: See TracBrowser for help on using the repository browser.