blob: ca0c9731a7baf28b51129f7403ca748d79a9c3d8 [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
107
108def address_type(address):
109 '''
110 Return the types of the address
111
112 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
113 '''
114 if type(address) == tuple:
115 return 'AF_INET'
116 elif type(address) is str and address.startswith('\\\\'):
117 return 'AF_PIPE'
118 elif type(address) is str:
119 return 'AF_UNIX'
120 else:
121 raise ValueError('address type of %r unrecognized' % address)
122
123#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200124# Connection classes
125#
126
127class _ConnectionBase:
128 _handle = None
129
130 def __init__(self, handle, readable=True, writable=True):
131 handle = handle.__index__()
132 if handle < 0:
133 raise ValueError("invalid handle")
134 if not readable and not writable:
135 raise ValueError(
136 "at least one of `readable` and `writable` must be True")
137 self._handle = handle
138 self._readable = readable
139 self._writable = writable
140
Antoine Pitrou60001202011-07-09 01:03:46 +0200141 # XXX should we use util.Finalize instead of a __del__?
142
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200143 def __del__(self):
144 if self._handle is not None:
145 self._close()
146
147 def _check_closed(self):
148 if self._handle is None:
149 raise IOError("handle is closed")
150
151 def _check_readable(self):
152 if not self._readable:
153 raise IOError("connection is write-only")
154
155 def _check_writable(self):
156 if not self._writable:
157 raise IOError("connection is read-only")
158
159 def _bad_message_length(self):
160 if self._writable:
161 self._readable = False
162 else:
163 self.close()
164 raise IOError("bad message length")
165
166 @property
167 def closed(self):
168 """True if the connection is closed"""
169 return self._handle is None
170
171 @property
172 def readable(self):
173 """True if the connection is readable"""
174 return self._readable
175
176 @property
177 def writable(self):
178 """True if the connection is writable"""
179 return self._writable
180
181 def fileno(self):
182 """File descriptor or handle of the connection"""
183 self._check_closed()
184 return self._handle
185
186 def close(self):
187 """Close the connection"""
188 if self._handle is not None:
189 try:
190 self._close()
191 finally:
192 self._handle = None
193
194 def send_bytes(self, buf, offset=0, size=None):
195 """Send the bytes data from a bytes-like object"""
196 self._check_closed()
197 self._check_writable()
198 m = memoryview(buf)
199 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
200 if m.itemsize > 1:
201 m = memoryview(bytes(m))
202 n = len(m)
203 if offset < 0:
204 raise ValueError("offset is negative")
205 if n < offset:
206 raise ValueError("buffer length < offset")
207 if size is None:
208 size = n - offset
209 elif size < 0:
210 raise ValueError("size is negative")
211 elif offset + size > n:
212 raise ValueError("buffer length < offset + size")
213 self._send_bytes(m[offset:offset + size])
214
215 def send(self, obj):
216 """Send a (picklable) object"""
217 self._check_closed()
218 self._check_writable()
219 buf = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
220 self._send_bytes(memoryview(buf))
221
222 def recv_bytes(self, maxlength=None):
223 """
224 Receive bytes data as a bytes object.
225 """
226 self._check_closed()
227 self._check_readable()
228 if maxlength is not None and maxlength < 0:
229 raise ValueError("negative maxlength")
230 buf = self._recv_bytes(maxlength)
231 if buf is None:
232 self._bad_message_length()
233 return buf.getvalue()
234
235 def recv_bytes_into(self, buf, offset=0):
236 """
237 Receive bytes data into a writeable buffer-like object.
238 Return the number of bytes read.
239 """
240 self._check_closed()
241 self._check_readable()
242 with memoryview(buf) as m:
243 # Get bytesize of arbitrary buffer
244 itemsize = m.itemsize
245 bytesize = itemsize * len(m)
246 if offset < 0:
247 raise ValueError("negative offset")
248 elif offset > bytesize:
249 raise ValueError("offset too large")
250 result = self._recv_bytes()
251 size = result.tell()
252 if bytesize < offset + size:
253 raise BufferTooShort(result.getvalue())
254 # Message can fit in dest
255 result.seek(0)
256 result.readinto(m[offset // itemsize :
257 (offset + size) // itemsize])
258 return size
259
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100260 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200261 """Receive a (picklable) object"""
262 self._check_closed()
263 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100264 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200265 return pickle.loads(buf.getbuffer())
266
267 def poll(self, timeout=0.0):
268 """Whether there is any input available to be read"""
269 self._check_closed()
270 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200271 return self._poll(timeout)
272
273
274if win32:
275
276 class PipeConnection(_ConnectionBase):
277 """
278 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200279 Overlapped I/O is used, so the handles must have been created
280 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200281 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100282 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200283
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200284 def _close(self, _CloseHandle=win32.CloseHandle):
285 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200286
287 def _send_bytes(self, buf):
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100288 ov, err = win32.WriteFile(self._handle, buf, overlapped=True)
289 try:
290 if err == win32.ERROR_IO_PENDING:
291 waitres = win32.WaitForMultipleObjects(
292 [ov.event], False, INFINITE)
293 assert waitres == WAIT_OBJECT_0
294 except:
295 ov.cancel()
296 raise
297 finally:
298 nwritten, err = ov.GetOverlappedResult(True)
299 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200300 assert nwritten == len(buf)
301
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100302 def _recv_bytes(self, maxsize=None):
303 if self._got_empty_message:
304 self._got_empty_message = False
305 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200306 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100307 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200308 try:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100309 ov, err = win32.ReadFile(self._handle, bsize,
310 overlapped=True)
311 try:
312 if err == win32.ERROR_IO_PENDING:
313 waitres = win32.WaitForMultipleObjects(
314 [ov.event], False, INFINITE)
315 assert waitres == WAIT_OBJECT_0
316 except:
317 ov.cancel()
318 raise
319 finally:
320 nread, err = ov.GetOverlappedResult(True)
321 if err == 0:
322 f = io.BytesIO()
323 f.write(ov.getbuffer())
324 return f
325 elif err == win32.ERROR_MORE_DATA:
326 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200327 except IOError as e:
Antoine Pitrou6b4883d2011-10-12 02:54:14 +0200328 if e.winerror == win32.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200329 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100330 else:
331 raise
332 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200333
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100334 def _poll(self, timeout):
335 if (self._got_empty_message or
336 win32.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200337 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100338 if timeout < 0:
339 timeout = None
340 return bool(wait([self], timeout))
341
342 def _get_more_data(self, ov, maxsize):
343 buf = ov.getbuffer()
344 f = io.BytesIO()
345 f.write(buf)
346 left = win32.PeekNamedPipe(self._handle)[1]
347 assert left > 0
348 if maxsize is not None and len(buf) + left > maxsize:
349 self._bad_message_length()
350 ov, err = win32.ReadFile(self._handle, left, overlapped=True)
351 rbytes, err = ov.GetOverlappedResult(True)
352 assert err == 0
353 assert rbytes == left
354 f.write(ov.getbuffer())
355 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200356
357
358class Connection(_ConnectionBase):
359 """
360 Connection class based on an arbitrary file descriptor (Unix only), or
361 a socket handle (Windows).
362 """
363
364 if win32:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200365 def _close(self, _close=win32.closesocket):
366 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200367 _write = win32.send
368 _read = win32.recv
369 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200370 def _close(self, _close=os.close):
371 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200372 _write = os.write
373 _read = os.read
374
375 def _send(self, buf, write=_write):
376 remaining = len(buf)
377 while True:
378 n = write(self._handle, buf)
379 remaining -= n
380 if remaining == 0:
381 break
382 buf = buf[n:]
383
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100384 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200385 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200386 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200387 remaining = size
388 while remaining > 0:
Antoine Pitroudd696492011-06-08 17:21:55 +0200389 chunk = read(handle, remaining)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200390 n = len(chunk)
391 if n == 0:
392 if remaining == size:
393 raise EOFError
394 else:
395 raise IOError("got end of file during message")
396 buf.write(chunk)
397 remaining -= n
398 return buf
399
400 def _send_bytes(self, buf):
401 # For wire compatibility with 3.2 and lower
402 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200403 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200404 # The condition is necessary to avoid "broken pipe" errors
405 # when sending a 0-length buffer if the other end closed the pipe.
406 if n > 0:
407 self._send(buf)
408
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100409 def _recv_bytes(self, maxsize=None):
410 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200411 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200412 if maxsize is not None and size > maxsize:
413 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100414 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200415
416 def _poll(self, timeout):
Antoine Pitroudd696492011-06-08 17:21:55 +0200417 if timeout < 0.0:
418 timeout = None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100419 r = wait([self._handle], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200420 return bool(r)
421
422
423#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000424# Public functions
425#
426
427class Listener(object):
428 '''
429 Returns a listener object.
430
431 This is a wrapper for a bound socket which is 'listening' for
432 connections, or for a Windows named pipe.
433 '''
434 def __init__(self, address=None, family=None, backlog=1, authkey=None):
435 family = family or (address and address_type(address)) \
436 or default_family
437 address = address or arbitrary_address(family)
438
439 if family == 'AF_PIPE':
440 self._listener = PipeListener(address, backlog)
441 else:
442 self._listener = SocketListener(address, family, backlog)
443
444 if authkey is not None and not isinstance(authkey, bytes):
445 raise TypeError('authkey should be a byte string')
446
447 self._authkey = authkey
448
449 def accept(self):
450 '''
451 Accept a connection on the bound socket or named pipe of `self`.
452
453 Returns a `Connection` object.
454 '''
455 c = self._listener.accept()
456 if self._authkey:
457 deliver_challenge(c, self._authkey)
458 answer_challenge(c, self._authkey)
459 return c
460
461 def close(self):
462 '''
463 Close the bound socket or named pipe of `self`.
464 '''
465 return self._listener.close()
466
467 address = property(lambda self: self._listener._address)
468 last_accepted = property(lambda self: self._listener._last_accepted)
469
470
471def Client(address, family=None, authkey=None):
472 '''
473 Returns a connection to the address of a `Listener`
474 '''
475 family = family or address_type(address)
476 if family == 'AF_PIPE':
477 c = PipeClient(address)
478 else:
479 c = SocketClient(address)
480
481 if authkey is not None and not isinstance(authkey, bytes):
482 raise TypeError('authkey should be a byte string')
483
484 if authkey is not None:
485 answer_challenge(c, authkey)
486 deliver_challenge(c, authkey)
487
488 return c
489
490
491if sys.platform != 'win32':
492
493 def Pipe(duplex=True):
494 '''
495 Returns pair of connection objects at either end of a pipe
496 '''
497 if duplex:
498 s1, s2 = socket.socketpair()
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200499 c1 = Connection(s1.detach())
500 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000501 else:
502 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200503 c1 = Connection(fd1, writable=False)
504 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000505
506 return c1, c2
507
508else:
509
Benjamin Petersone711caf2008-06-11 16:44:04 +0000510 def Pipe(duplex=True):
511 '''
512 Returns pair of connection objects at either end of a pipe
513 '''
514 address = arbitrary_address('AF_PIPE')
515 if duplex:
516 openmode = win32.PIPE_ACCESS_DUPLEX
517 access = win32.GENERIC_READ | win32.GENERIC_WRITE
518 obsize, ibsize = BUFSIZE, BUFSIZE
519 else:
520 openmode = win32.PIPE_ACCESS_INBOUND
521 access = win32.GENERIC_WRITE
522 obsize, ibsize = 0, BUFSIZE
523
524 h1 = win32.CreateNamedPipe(
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100525 address, openmode | win32.FILE_FLAG_OVERLAPPED |
526 win32.FILE_FLAG_FIRST_PIPE_INSTANCE,
Benjamin Petersone711caf2008-06-11 16:44:04 +0000527 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
528 win32.PIPE_WAIT,
529 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
530 )
531 h2 = win32.CreateFile(
Antoine Pitroudd696492011-06-08 17:21:55 +0200532 address, access, 0, win32.NULL, win32.OPEN_EXISTING,
533 win32.FILE_FLAG_OVERLAPPED, win32.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000534 )
535 win32.SetNamedPipeHandleState(
536 h2, win32.PIPE_READMODE_MESSAGE, None, None
537 )
538
Antoine Pitroudd696492011-06-08 17:21:55 +0200539 overlapped = win32.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100540 _, err = overlapped.GetOverlappedResult(True)
541 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000542
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200543 c1 = PipeConnection(h1, writable=duplex)
544 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000545
546 return c1, c2
547
548#
549# Definitions for connections based on sockets
550#
551
552class SocketListener(object):
553 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000554 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000555 '''
556 def __init__(self, address, family, backlog=1):
557 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100558 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100559 # SO_REUSEADDR has different semantics on Windows (issue #2550).
560 if os.name == 'posix':
561 self._socket.setsockopt(socket.SOL_SOCKET,
562 socket.SO_REUSEADDR, 1)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100563 self._socket.bind(address)
564 self._socket.listen(backlog)
565 self._address = self._socket.getsockname()
566 except OSError:
567 self._socket.close()
568 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000569 self._family = family
570 self._last_accepted = None
571
Benjamin Petersone711caf2008-06-11 16:44:04 +0000572 if family == 'AF_UNIX':
573 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000574 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000575 )
576 else:
577 self._unlink = None
578
579 def accept(self):
580 s, self._last_accepted = self._socket.accept()
581 fd = duplicate(s.fileno())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200582 conn = Connection(fd)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000583 s.close()
584 return conn
585
586 def close(self):
587 self._socket.close()
588 if self._unlink is not None:
589 self._unlink()
590
591
592def SocketClient(address):
593 '''
594 Return a connection object connected to the socket given by `address`
595 '''
596 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000597 with socket.socket( getattr(socket, family) ) as s:
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100598 s.connect(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000599 fd = duplicate(s.fileno())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200600 conn = Connection(fd)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000601 return conn
602
603#
604# Definitions for connections based on named pipes
605#
606
607if sys.platform == 'win32':
608
609 class PipeListener(object):
610 '''
611 Representation of a named pipe
612 '''
613 def __init__(self, address, backlog=None):
614 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100615 self._handle_queue = [self._new_handle(first=True)]
616
Benjamin Petersone711caf2008-06-11 16:44:04 +0000617 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000618 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000619 self.close = Finalize(
620 self, PipeListener._finalize_pipe_listener,
621 args=(self._handle_queue, self._address), exitpriority=0
622 )
623
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100624 def _new_handle(self, first=False):
625 flags = win32.PIPE_ACCESS_DUPLEX | win32.FILE_FLAG_OVERLAPPED
626 if first:
627 flags |= win32.FILE_FLAG_FIRST_PIPE_INSTANCE
628 return win32.CreateNamedPipe(
629 self._address, flags,
Benjamin Petersone711caf2008-06-11 16:44:04 +0000630 win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
631 win32.PIPE_WAIT,
632 win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
633 win32.NMPWAIT_WAIT_FOREVER, win32.NULL
634 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100635
636 def accept(self):
637 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000638 handle = self._handle_queue.pop(0)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100639 ov = win32.ConnectNamedPipe(handle, overlapped=True)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000640 try:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100641 res = win32.WaitForMultipleObjects([ov.event], False, INFINITE)
642 except:
643 ov.cancel()
644 win32.CloseHandle(handle)
645 raise
646 finally:
647 _, err = ov.GetOverlappedResult(True)
648 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200649 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000650
651 @staticmethod
652 def _finalize_pipe_listener(queue, address):
653 sub_debug('closing listener with address=%r', address)
654 for handle in queue:
655 close(handle)
656
657 def PipeClient(address):
658 '''
659 Return a connection object connected to the pipe given by `address`
660 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000661 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000662 while 1:
663 try:
664 win32.WaitNamedPipe(address, 1000)
665 h = win32.CreateFile(
666 address, win32.GENERIC_READ | win32.GENERIC_WRITE,
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100667 0, win32.NULL, win32.OPEN_EXISTING,
668 win32.FILE_FLAG_OVERLAPPED, win32.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000669 )
670 except WindowsError as e:
Antoine Pitrou6b4883d2011-10-12 02:54:14 +0200671 if e.winerror not in (win32.ERROR_SEM_TIMEOUT,
672 win32.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000673 raise
674 else:
675 break
676 else:
677 raise
678
679 win32.SetNamedPipeHandleState(
680 h, win32.PIPE_READMODE_MESSAGE, None, None
681 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200682 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000683
684#
685# Authentication stuff
686#
687
688MESSAGE_LENGTH = 20
689
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000690CHALLENGE = b'#CHALLENGE#'
691WELCOME = b'#WELCOME#'
692FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000693
694def deliver_challenge(connection, authkey):
695 import hmac
696 assert isinstance(authkey, bytes)
697 message = os.urandom(MESSAGE_LENGTH)
698 connection.send_bytes(CHALLENGE + message)
699 digest = hmac.new(authkey, message).digest()
700 response = connection.recv_bytes(256) # reject large message
701 if response == digest:
702 connection.send_bytes(WELCOME)
703 else:
704 connection.send_bytes(FAILURE)
705 raise AuthenticationError('digest received was wrong')
706
707def answer_challenge(connection, authkey):
708 import hmac
709 assert isinstance(authkey, bytes)
710 message = connection.recv_bytes(256) # reject large message
711 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
712 message = message[len(CHALLENGE):]
713 digest = hmac.new(authkey, message).digest()
714 connection.send_bytes(digest)
715 response = connection.recv_bytes(256) # reject large message
716 if response != WELCOME:
717 raise AuthenticationError('digest sent was rejected')
718
719#
720# Support for using xmlrpclib for serialization
721#
722
723class ConnectionWrapper(object):
724 def __init__(self, conn, dumps, loads):
725 self._conn = conn
726 self._dumps = dumps
727 self._loads = loads
728 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
729 obj = getattr(conn, attr)
730 setattr(self, attr, obj)
731 def send(self, obj):
732 s = self._dumps(obj)
733 self._conn.send_bytes(s)
734 def recv(self):
735 s = self._conn.recv_bytes()
736 return self._loads(s)
737
738def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000739 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000740
741def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000742 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000743 return obj
744
745class XmlListener(Listener):
746 def accept(self):
747 global xmlrpclib
748 import xmlrpc.client as xmlrpclib
749 obj = Listener.accept(self)
750 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
751
752def XmlClient(*args, **kwds):
753 global xmlrpclib
754 import xmlrpc.client as xmlrpclib
755 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200756
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100757#
758# Wait
759#
760
761if sys.platform == 'win32':
762
763 def _exhaustive_wait(handles, timeout):
764 # Return ALL handles which are currently signalled. (Only
765 # returning the first signalled might create starvation issues.)
766 L = list(handles)
767 ready = []
768 while L:
769 res = win32.WaitForMultipleObjects(L, False, timeout)
770 if res == WAIT_TIMEOUT:
771 break
772 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
773 res -= WAIT_OBJECT_0
774 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
775 res -= WAIT_ABANDONED_0
776 else:
777 raise RuntimeError('Should not get here')
778 ready.append(L[res])
779 L = L[res+1:]
780 timeout = 0
781 return ready
782
783 _ready_errors = {win32.ERROR_BROKEN_PIPE, win32.ERROR_NETNAME_DELETED}
784
785 def wait(object_list, timeout=None):
786 '''
787 Wait till an object in object_list is ready/readable.
788
789 Returns list of those objects in object_list which are ready/readable.
790 '''
791 if timeout is None:
792 timeout = INFINITE
793 elif timeout < 0:
794 timeout = 0
795 else:
796 timeout = int(timeout * 1000 + 0.5)
797
798 object_list = list(object_list)
799 waithandle_to_obj = {}
800 ov_list = []
801 ready_objects = set()
802 ready_handles = set()
803
804 try:
805 for o in object_list:
806 try:
807 fileno = getattr(o, 'fileno')
808 except AttributeError:
809 waithandle_to_obj[o.__index__()] = o
810 else:
811 # start an overlapped read of length zero
812 try:
813 ov, err = win32.ReadFile(fileno(), 0, True)
814 except OSError as e:
815 err = e.winerror
816 if err not in _ready_errors:
817 raise
818 if err == win32.ERROR_IO_PENDING:
819 ov_list.append(ov)
820 waithandle_to_obj[ov.event] = o
821 else:
822 # If o.fileno() is an overlapped pipe handle and
823 # err == 0 then there is a zero length message
824 # in the pipe, but it HAS NOT been consumed.
825 ready_objects.add(o)
826 timeout = 0
827
828 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
829 finally:
830 # request that overlapped reads stop
831 for ov in ov_list:
832 ov.cancel()
833
834 # wait for all overlapped reads to stop
835 for ov in ov_list:
836 try:
837 _, err = ov.GetOverlappedResult(True)
838 except OSError as e:
839 err = e.winerror
840 if err not in _ready_errors:
841 raise
842 if err != win32.ERROR_OPERATION_ABORTED:
843 o = waithandle_to_obj[ov.event]
844 ready_objects.add(o)
845 if err == 0:
846 # If o.fileno() is an overlapped pipe handle then
847 # a zero length message HAS been consumed.
848 if hasattr(o, '_got_empty_message'):
849 o._got_empty_message = True
850
851 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
852 return [o for o in object_list if o in ready_objects]
853
854else:
855
856 def wait(object_list, timeout=None):
857 '''
858 Wait till an object in object_list is ready/readable.
859
860 Returns list of those objects in object_list which are ready/readable.
861 '''
862 if timeout is not None:
863 if timeout <= 0:
864 return select.select(object_list, [], [], 0)[0]
865 else:
866 deadline = time.time() + timeout
867 while True:
868 try:
869 return select.select(object_list, [], [], timeout)[0]
870 except OSError as e:
871 if e.errno != errno.EINTR:
872 raise
873 if timeout is not None:
874 timeout = deadline - time.time()
875
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200876
877# Late import because of circular import
878from multiprocessing.forking import duplicate, close