blob: 954e90156359f6cf9cdf2ab62bbf9becd248173c [file] [log] [blame]
Benjamin Petersone711caf2008-06-11 16:44:04 +00001#
2# A higher level module for using sockets (or Windows named pipes)
3#
4# multiprocessing/connection.py
5#
R. David Murray3fc969a2010-12-14 01:38:16 +00006# 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.
Benjamin Petersone711caf2008-06-11 16:44:04 +000033#
34
Antoine Pitroubdb1cf12012-03-05 19:28:37 +010035__all__ = [ 'Client', 'Listener', 'Pipe', 'wait' ]
Benjamin Petersone711caf2008-06-11 16:44:04 +000036
Antoine Pitrou87cf2202011-05-09 17:04:27 +020037import io
Benjamin Petersone711caf2008-06-11 16:44:04 +000038import os
39import sys
Antoine Pitrou87cf2202011-05-09 17:04:27 +020040import pickle
41import select
Benjamin Petersone711caf2008-06-11 16:44:04 +000042import socket
Antoine Pitrou87cf2202011-05-09 17:04:27 +020043import struct
Georg Brandl6aa2d1f2008-08-12 08:35:52 +000044import errno
Benjamin Petersone711caf2008-06-11 16:44:04 +000045import time
46import tempfile
47import itertools
48
49import _multiprocessing
Antoine Pitrou87cf2202011-05-09 17:04:27 +020050from multiprocessing import current_process, AuthenticationError, BufferTooShort
Antoine Pitroudd696492011-06-08 17:21:55 +020051from multiprocessing.util import (
52 get_temp_dir, Finalize, sub_debug, debug, _eintr_retry)
Antoine Pitrou87cf2202011-05-09 17:04:27 +020053try:
54 from _multiprocessing import win32
Antoine Pitroudd696492011-06-08 17:21:55 +020055 from _subprocess import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Antoine Pitrou87cf2202011-05-09 17:04:27 +020056except ImportError:
57 if sys.platform == 'win32':
58 raise
59 win32 = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000060
61#
62#
63#
64
65BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000066# A very generous timeout when it comes to local connections...
67CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000068
69_mmap_counter = itertools.count()
70
71default_family = 'AF_INET'
72families = ['AF_INET']
73
74if hasattr(socket, 'AF_UNIX'):
75 default_family = 'AF_UNIX'
76 families += ['AF_UNIX']
77
78if sys.platform == 'win32':
79 default_family = 'AF_PIPE'
80 families += ['AF_PIPE']
81
Antoine Pitrou45d61a32009-11-13 22:35:18 +000082
83def _init_timeout(timeout=CONNECTION_TIMEOUT):
84 return time.time() + timeout
85
86def _check_timeout(t):
87 return time.time() > t
88
Benjamin Petersone711caf2008-06-11 16:44:04 +000089#
90#
91#
92
93def arbitrary_address(family):
94 '''
95 Return an arbitrary free address for the given family
96 '''
97 if family == 'AF_INET':
98 return ('localhost', 0)
99 elif family == 'AF_UNIX':
100 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
101 elif family == 'AF_PIPE':
102 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
103 (os.getpid(), next(_mmap_counter)))
104 else:
105 raise ValueError('unrecognized family')
106
Antoine Pitrou709176f2012-04-01 17:19:09 +0200107def _validate_family(family):
108 '''
109 Checks if the family is valid for the current environment.
110 '''
111 if sys.platform != 'win32' and family == 'AF_PIPE':
112 raise ValueError('Family %s is not recognized.' % family)
113
Benjamin Petersone711caf2008-06-11 16:44:04 +0000114
115def address_type(address):
116 '''
117 Return the types of the address
118
119 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
120 '''
121 if type(address) == tuple:
122 return 'AF_INET'
123 elif type(address) is str and address.startswith('\\\\'):
124 return 'AF_PIPE'
125 elif type(address) is str:
126 return 'AF_UNIX'
127 else:
128 raise ValueError('address type of %r unrecognized' % address)
129
130#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200131# Connection classes
132#
133
134class _ConnectionBase:
135 _handle = None
136
137 def __init__(self, handle, readable=True, writable=True):
138 handle = handle.__index__()
139 if handle < 0:
140 raise ValueError("invalid handle")
141 if not readable and not writable:
142 raise ValueError(
143 "at least one of `readable` and `writable` must be True")
144 self._handle = handle
145 self._readable = readable
146 self._writable = writable
147
Antoine Pitrou60001202011-07-09 01:03:46 +0200148 # XXX should we use util.Finalize instead of a __del__?
149
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200150 def __del__(self):
151 if self._handle is not None:
152 self._close()
153
154 def _check_closed(self):
155 if self._handle is None:
156 raise IOError("handle is closed")
157
158 def _check_readable(self):
159 if not self._readable:
160 raise IOError("connection is write-only")
161
162 def _check_writable(self):
163 if not self._writable:
164 raise IOError("connection is read-only")
165
166 def _bad_message_length(self):
167 if self._writable:
168 self._readable = False
169 else:
170 self.close()
171 raise IOError("bad message length")
172
173 @property
174 def closed(self):
175 """True if the connection is closed"""
176 return self._handle is None
177
178 @property
179 def readable(self):
180 """True if the connection is readable"""
181 return self._readable
182
183 @property
184 def writable(self):
185 """True if the connection is writable"""
186 return self._writable
187
188 def fileno(self):
189 """File descriptor or handle of the connection"""
190 self._check_closed()
191 return self._handle
192
193 def close(self):
194 """Close the connection"""
195 if self._handle is not None:
196 try:
197 self._close()
198 finally:
199 self._handle = None
200
201 def send_bytes(self, buf, offset=0, size=None):
202 """Send the bytes data from a bytes-like object"""
203 self._check_closed()
204 self._check_writable()
205 m = memoryview(buf)
206 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
207 if m.itemsize > 1:
208 m = memoryview(bytes(m))
209 n = len(m)
210 if offset < 0:
211 raise ValueError("offset is negative")
212 if n < offset:
213 raise ValueError("buffer length < offset")
214 if size is None:
215 size = n - offset
216 elif size < 0:
217 raise ValueError("size is negative")
218 elif offset + size > n:
219 raise ValueError("buffer length < offset + size")
220 self._send_bytes(m[offset:offset + size])
221
222 def send(self, obj):
223 """Send a (picklable) object"""
224 self._check_closed()
225 self._check_writable()
226 buf = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
227 self._send_bytes(memoryview(buf))
228
229 def recv_bytes(self, maxlength=None):
230 """
231 Receive bytes data as a bytes object.
232 """
233 self._check_closed()
234 self._check_readable()
235 if maxlength is not None and maxlength < 0:
236 raise ValueError("negative maxlength")
237 buf = self._recv_bytes(maxlength)
238 if buf is None:
239 self._bad_message_length()
240 return buf.getvalue()
241
242 def recv_bytes_into(self, buf, offset=0):
243 """
244 Receive bytes data into a writeable buffer-like object.
245 Return the number of bytes read.
246 """
247 self._check_closed()
248 self._check_readable()
249 with memoryview(buf) as m:
250 # Get bytesize of arbitrary buffer
251 itemsize = m.itemsize
252 bytesize = itemsize * len(m)
253 if offset < 0:
254 raise ValueError("negative offset")
255 elif offset > bytesize:
256 raise ValueError("offset too large")
257 result = self._recv_bytes()
258 size = result.tell()
259 if bytesize < offset + size:
260 raise BufferTooShort(result.getvalue())
261 # Message can fit in dest
262 result.seek(0)
263 result.readinto(m[offset // itemsize :
264 (offset + size) // itemsize])
265 return size
266
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100267 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200268 """Receive a (picklable) object"""
269 self._check_closed()
270 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100271 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200272 return pickle.loads(buf.getbuffer())
273
274 def poll(self, timeout=0.0):
275 """Whether there is any input available to be read"""
276 self._check_closed()
277 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200278 return self._poll(timeout)
279
280
281if win32:
282
283 class PipeConnection(_ConnectionBase):
284 """
285 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200286 Overlapped I/O is used, so the handles must have been created
287 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200288 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100289 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200290
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200291 def _close(self, _CloseHandle=win32.CloseHandle):
292 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200293
294 def _send_bytes(self, buf):
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100295 ov, err = win32.WriteFile(self._handle, buf, overlapped=True)
296 try:
297 if err == win32.ERROR_IO_PENDING:
298 waitres = win32.WaitForMultipleObjects(
299 [ov.event], False, INFINITE)
300 assert waitres == WAIT_OBJECT_0
301 except:
302 ov.cancel()
303 raise
304 finally:
305 nwritten, err = ov.GetOverlappedResult(True)
306 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200307 assert nwritten == len(buf)
308
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100309 def _recv_bytes(self, maxsize=None):
310 if self._got_empty_message:
311 self._got_empty_message = False
312 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200313 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100314 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200315 try:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100316 ov, err = win32.ReadFile(self._handle, bsize,
317 overlapped=True)
318 try:
319 if err == win32.ERROR_IO_PENDING:
320 waitres = win32.WaitForMultipleObjects(
321 [ov.event], False, INFINITE)
322 assert waitres == WAIT_OBJECT_0
323 except:
324 ov.cancel()
325 raise
326 finally:
327 nread, err = ov.GetOverlappedResult(True)
328 if err == 0:
329 f = io.BytesIO()
330 f.write(ov.getbuffer())
331 return f
332 elif err == win32.ERROR_MORE_DATA:
333 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200334 except IOError as e:
Antoine Pitrou6b4883d2011-10-12 02:54:14 +0200335 if e.winerror == win32.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200336 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100337 else:
338 raise
339 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200340
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100341 def _poll(self, timeout):
342 if (self._got_empty_message or
343 win32.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200344 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100345 if timeout < 0:
346 timeout = None
347 return bool(wait([self], timeout))
348
349 def _get_more_data(self, ov, maxsize):
350 buf = ov.getbuffer()
351 f = io.BytesIO()
352 f.write(buf)
353 left = win32.PeekNamedPipe(self._handle)[1]
354 assert left > 0
355 if maxsize is not None and len(buf) + left > maxsize:
356 self._bad_message_length()
357 ov, err = win32.ReadFile(self._handle, left, overlapped=True)
358 rbytes, err = ov.GetOverlappedResult(True)
359 assert err == 0
360 assert rbytes == left
361 f.write(ov.getbuffer())
362 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200363
364
365class Connection(_ConnectionBase):
366 """
367 Connection class based on an arbitrary file descriptor (Unix only), or
368 a socket handle (Windows).
369 """
370
371 if win32:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200372 def _close(self, _close=win32.closesocket):
373 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200374 _write = win32.send
375 _read = win32.recv
376 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200377 def _close(self, _close=os.close):
378 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200379 _write = os.write
380 _read = os.read
381
382 def _send(self, buf, write=_write):
383 remaining = len(buf)
384 while True:
385 n = write(self._handle, buf)
386 remaining -= n
387 if remaining == 0:
388 break
389 buf = buf[n:]
390
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100391 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200392 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200393 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200394 remaining = size
395 while remaining > 0:
Antoine Pitroudd696492011-06-08 17:21:55 +0200396 chunk = read(handle, remaining)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200397 n = len(chunk)
398 if n == 0:
399 if remaining == size:
400 raise EOFError
401 else:
402 raise IOError("got end of file during message")
403 buf.write(chunk)
404 remaining -= n
405 return buf
406
407 def _send_bytes(self, buf):
408 # For wire compatibility with 3.2 and lower
409 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200410 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200411 # The condition is necessary to avoid "broken pipe" errors
412 # when sending a 0-length buffer if the other end closed the pipe.
413 if n > 0:
414 self._send(buf)
415
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100416 def _recv_bytes(self, maxsize=None):
417 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200418 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200419 if maxsize is not None and size > maxsize:
420 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100421 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200422
423 def _poll(self, timeout):
Antoine Pitroudd696492011-06-08 17:21:55 +0200424 if timeout < 0.0:
425 timeout = None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100426 r = wait([self._handle], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200427 return bool(r)
428
429
430#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000431# Public functions
432#
433
434class Listener(object):
435 '''
436 Returns a listener object.
437
438 This is a wrapper for a bound socket which is 'listening' for
439 connections, or for a Windows named pipe.
440 '''
441 def __init__(self, address=None, family=None, backlog=1, authkey=None):
442 family = family or (address and address_type(address)) \
443 or default_family
444 address = address or arbitrary_address(family)
445
Antoine Pitrou709176f2012-04-01 17:19:09 +0200446 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000447 if family == 'AF_PIPE':
448 self._listener = PipeListener(address, backlog)
449 else:
450 self._listener = SocketListener(address, family, backlog)
451
452 if authkey is not None and not isinstance(authkey, bytes):
453 raise TypeError('authkey should be a byte string')
454
455 self._authkey = authkey
456
457 def accept(self):
458 '''
459 Accept a connection on the bound socket or named pipe of `self`.
460
461 Returns a `Connection` object.
462 '''
463 c = self._listener.accept()
464 if self._authkey:
465 deliver_challenge(c, self._authkey)
466 answer_challenge(c, self._authkey)
467 return c
468
469 def close(self):
470 '''
471 Close the bound socket or named pipe of `self`.
472 '''
473 return self._listener.close()
474
475 address = property(lambda self: self._listener._address)
476 last_accepted = property(lambda self: self._listener._last_accepted)
477
478
479def Client(address, family=None, authkey=None):
480 '''
481 Returns a connection to the address of a `Listener`
482 '''
483 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200484 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000485 if family == 'AF_PIPE':
486 c = PipeClient(address)
487 else:
488 c = SocketClient(address)
489
490 if authkey is not None and not isinstance(authkey, bytes):
491 raise TypeError('authkey should be a byte string')
492
493 if authkey is not None:
494 answer_challenge(c, authkey)
495 deliver_challenge(c, authkey)
496
497 return c
498
499
500if sys.platform != 'win32':
501
502 def Pipe(duplex=True):
503 '''
504 Returns pair of connection objects at either end of a pipe
505 '''
506 if duplex:
507 s1, s2 = socket.socketpair()
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200508 c1 = Connection(s1.detach())
509 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000510 else:
511 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200512 c1 = Connection(fd1, writable=False)
513 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000514
515 return c1, c2
516
517else:
518
Benjamin Petersone711caf2008-06-11 16:44:04 +0000519 def Pipe(duplex=True):
520 '''
521 Returns pair of connection objects at either end of a pipe
522 '''
523 address = arbitrary_address('AF_PIPE')
524 if duplex:
525 openmode = win32.PIPE_ACCESS_DUPLEX
526 access = win32.GENERIC_READ | win32.GENERIC_WRITE
527 obsize, ibsize = BUFSIZE, BUFSIZE
528 else:
529 openmode = win32.PIPE_ACCESS_INBOUND
530 access = win32.GENERIC_WRITE
531 obsize, ibsize = 0, BUFSIZE
532
533 h1 = win32.CreateNamedPipe(
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100534 address, openmode | win32.FILE_FLAG_OVERLAPPED |
535 win32.FILE_FLAG_FIRST_PIPE_INSTANCE,
Benjamin Petersone711caf2008-06-11 16:44:04 +0000536 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
537 win32.PIPE_WAIT,
538 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
539 )
540 h2 = win32.CreateFile(
Antoine Pitroudd696492011-06-08 17:21:55 +0200541 address, access, 0, win32.NULL, win32.OPEN_EXISTING,
542 win32.FILE_FLAG_OVERLAPPED, win32.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000543 )
544 win32.SetNamedPipeHandleState(
545 h2, win32.PIPE_READMODE_MESSAGE, None, None
546 )
547
Antoine Pitroudd696492011-06-08 17:21:55 +0200548 overlapped = win32.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100549 _, err = overlapped.GetOverlappedResult(True)
550 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000551
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200552 c1 = PipeConnection(h1, writable=duplex)
553 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000554
555 return c1, c2
556
557#
558# Definitions for connections based on sockets
559#
560
561class SocketListener(object):
562 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000563 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000564 '''
565 def __init__(self, address, family, backlog=1):
566 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100567 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100568 # SO_REUSEADDR has different semantics on Windows (issue #2550).
569 if os.name == 'posix':
570 self._socket.setsockopt(socket.SOL_SOCKET,
571 socket.SO_REUSEADDR, 1)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100572 self._socket.bind(address)
573 self._socket.listen(backlog)
574 self._address = self._socket.getsockname()
575 except OSError:
576 self._socket.close()
577 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000578 self._family = family
579 self._last_accepted = None
580
Benjamin Petersone711caf2008-06-11 16:44:04 +0000581 if family == 'AF_UNIX':
582 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000583 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000584 )
585 else:
586 self._unlink = None
587
588 def accept(self):
589 s, self._last_accepted = self._socket.accept()
590 fd = duplicate(s.fileno())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200591 conn = Connection(fd)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000592 s.close()
593 return conn
594
595 def close(self):
596 self._socket.close()
597 if self._unlink is not None:
598 self._unlink()
599
600
601def SocketClient(address):
602 '''
603 Return a connection object connected to the socket given by `address`
604 '''
605 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000606 with socket.socket( getattr(socket, family) ) as s:
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100607 s.connect(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000608 fd = duplicate(s.fileno())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200609 conn = Connection(fd)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000610 return conn
611
612#
613# Definitions for connections based on named pipes
614#
615
616if sys.platform == 'win32':
617
618 class PipeListener(object):
619 '''
620 Representation of a named pipe
621 '''
622 def __init__(self, address, backlog=None):
623 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100624 self._handle_queue = [self._new_handle(first=True)]
625
Benjamin Petersone711caf2008-06-11 16:44:04 +0000626 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000627 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000628 self.close = Finalize(
629 self, PipeListener._finalize_pipe_listener,
630 args=(self._handle_queue, self._address), exitpriority=0
631 )
632
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100633 def _new_handle(self, first=False):
634 flags = win32.PIPE_ACCESS_DUPLEX | win32.FILE_FLAG_OVERLAPPED
635 if first:
636 flags |= win32.FILE_FLAG_FIRST_PIPE_INSTANCE
637 return win32.CreateNamedPipe(
638 self._address, flags,
Benjamin Petersone711caf2008-06-11 16:44:04 +0000639 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
640 win32.PIPE_WAIT,
641 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
642 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
643 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100644
645 def accept(self):
646 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000647 handle = self._handle_queue.pop(0)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100648 ov = win32.ConnectNamedPipe(handle, overlapped=True)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000649 try:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100650 res = win32.WaitForMultipleObjects([ov.event], False, INFINITE)
651 except:
652 ov.cancel()
653 win32.CloseHandle(handle)
654 raise
655 finally:
656 _, err = ov.GetOverlappedResult(True)
657 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200658 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000659
660 @staticmethod
661 def _finalize_pipe_listener(queue, address):
662 sub_debug('closing listener with address=%r', address)
663 for handle in queue:
664 close(handle)
665
666 def PipeClient(address):
667 '''
668 Return a connection object connected to the pipe given by `address`
669 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000670 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000671 while 1:
672 try:
673 win32.WaitNamedPipe(address, 1000)
674 h = win32.CreateFile(
675 address, win32.GENERIC_READ | win32.GENERIC_WRITE,
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100676 0, win32.NULL, win32.OPEN_EXISTING,
677 win32.FILE_FLAG_OVERLAPPED, win32.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000678 )
679 except WindowsError as e:
Antoine Pitrou6b4883d2011-10-12 02:54:14 +0200680 if e.winerror not in (win32.ERROR_SEM_TIMEOUT,
681 win32.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000682 raise
683 else:
684 break
685 else:
686 raise
687
688 win32.SetNamedPipeHandleState(
689 h, win32.PIPE_READMODE_MESSAGE, None, None
690 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200691 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000692
693#
694# Authentication stuff
695#
696
697MESSAGE_LENGTH = 20
698
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000699CHALLENGE = b'#CHALLENGE#'
700WELCOME = b'#WELCOME#'
701FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000702
703def deliver_challenge(connection, authkey):
704 import hmac
705 assert isinstance(authkey, bytes)
706 message = os.urandom(MESSAGE_LENGTH)
707 connection.send_bytes(CHALLENGE + message)
708 digest = hmac.new(authkey, message).digest()
709 response = connection.recv_bytes(256) # reject large message
710 if response == digest:
711 connection.send_bytes(WELCOME)
712 else:
713 connection.send_bytes(FAILURE)
714 raise AuthenticationError('digest received was wrong')
715
716def answer_challenge(connection, authkey):
717 import hmac
718 assert isinstance(authkey, bytes)
719 message = connection.recv_bytes(256) # reject large message
720 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
721 message = message[len(CHALLENGE):]
722 digest = hmac.new(authkey, message).digest()
723 connection.send_bytes(digest)
724 response = connection.recv_bytes(256) # reject large message
725 if response != WELCOME:
726 raise AuthenticationError('digest sent was rejected')
727
728#
729# Support for using xmlrpclib for serialization
730#
731
732class ConnectionWrapper(object):
733 def __init__(self, conn, dumps, loads):
734 self._conn = conn
735 self._dumps = dumps
736 self._loads = loads
737 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
738 obj = getattr(conn, attr)
739 setattr(self, attr, obj)
740 def send(self, obj):
741 s = self._dumps(obj)
742 self._conn.send_bytes(s)
743 def recv(self):
744 s = self._conn.recv_bytes()
745 return self._loads(s)
746
747def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000748 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000749
750def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000751 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000752 return obj
753
754class XmlListener(Listener):
755 def accept(self):
756 global xmlrpclib
757 import xmlrpc.client as xmlrpclib
758 obj = Listener.accept(self)
759 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
760
761def XmlClient(*args, **kwds):
762 global xmlrpclib
763 import xmlrpc.client as xmlrpclib
764 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200765
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100766#
767# Wait
768#
769
770if sys.platform == 'win32':
771
772 def _exhaustive_wait(handles, timeout):
773 # Return ALL handles which are currently signalled. (Only
774 # returning the first signalled might create starvation issues.)
775 L = list(handles)
776 ready = []
777 while L:
778 res = win32.WaitForMultipleObjects(L, False, timeout)
779 if res == WAIT_TIMEOUT:
780 break
781 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
782 res -= WAIT_OBJECT_0
783 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
784 res -= WAIT_ABANDONED_0
785 else:
786 raise RuntimeError('Should not get here')
787 ready.append(L[res])
788 L = L[res+1:]
789 timeout = 0
790 return ready
791
792 _ready_errors = {win32.ERROR_BROKEN_PIPE, win32.ERROR_NETNAME_DELETED}
793
794 def wait(object_list, timeout=None):
795 '''
796 Wait till an object in object_list is ready/readable.
797
798 Returns list of those objects in object_list which are ready/readable.
799 '''
800 if timeout is None:
801 timeout = INFINITE
802 elif timeout < 0:
803 timeout = 0
804 else:
805 timeout = int(timeout * 1000 + 0.5)
806
807 object_list = list(object_list)
808 waithandle_to_obj = {}
809 ov_list = []
810 ready_objects = set()
811 ready_handles = set()
812
813 try:
814 for o in object_list:
815 try:
816 fileno = getattr(o, 'fileno')
817 except AttributeError:
818 waithandle_to_obj[o.__index__()] = o
819 else:
820 # start an overlapped read of length zero
821 try:
822 ov, err = win32.ReadFile(fileno(), 0, True)
823 except OSError as e:
824 err = e.winerror
825 if err not in _ready_errors:
826 raise
827 if err == win32.ERROR_IO_PENDING:
828 ov_list.append(ov)
829 waithandle_to_obj[ov.event] = o
830 else:
831 # If o.fileno() is an overlapped pipe handle and
832 # err == 0 then there is a zero length message
833 # in the pipe, but it HAS NOT been consumed.
834 ready_objects.add(o)
835 timeout = 0
836
837 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
838 finally:
839 # request that overlapped reads stop
840 for ov in ov_list:
841 ov.cancel()
842
843 # wait for all overlapped reads to stop
844 for ov in ov_list:
845 try:
846 _, err = ov.GetOverlappedResult(True)
847 except OSError as e:
848 err = e.winerror
849 if err not in _ready_errors:
850 raise
851 if err != win32.ERROR_OPERATION_ABORTED:
852 o = waithandle_to_obj[ov.event]
853 ready_objects.add(o)
854 if err == 0:
855 # If o.fileno() is an overlapped pipe handle then
856 # a zero length message HAS been consumed.
857 if hasattr(o, '_got_empty_message'):
858 o._got_empty_message = True
859
860 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
861 return [o for o in object_list if o in ready_objects]
862
863else:
864
865 def wait(object_list, timeout=None):
866 '''
867 Wait till an object in object_list is ready/readable.
868
869 Returns list of those objects in object_list which are ready/readable.
870 '''
871 if timeout is not None:
872 if timeout <= 0:
873 return select.select(object_list, [], [], 0)[0]
874 else:
875 deadline = time.time() + timeout
876 while True:
877 try:
878 return select.select(object_list, [], [], timeout)[0]
879 except OSError as e:
880 if e.errno != errno.EINTR:
881 raise
882 if timeout is not None:
883 timeout = deadline - time.time()
884
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200885
886# Late import because of circular import
887from multiprocessing.forking import duplicate, close