blob: 64d71bc361f8144701d9f2d3d1d412fb58643553 [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 Pitrou5438ed12012-04-24 22:56:57 +020053from multiprocessing.forking import ForkingPickler
Antoine Pitrou87cf2202011-05-09 17:04:27 +020054try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020055 import _winapi
56 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Antoine Pitrou87cf2202011-05-09 17:04:27 +020057except ImportError:
58 if sys.platform == 'win32':
59 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020060 _winapi = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000061
62#
63#
64#
65
66BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000067# A very generous timeout when it comes to local connections...
68CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000069
70_mmap_counter = itertools.count()
71
72default_family = 'AF_INET'
73families = ['AF_INET']
74
75if hasattr(socket, 'AF_UNIX'):
76 default_family = 'AF_UNIX'
77 families += ['AF_UNIX']
78
79if sys.platform == 'win32':
80 default_family = 'AF_PIPE'
81 families += ['AF_PIPE']
82
Antoine Pitrou45d61a32009-11-13 22:35:18 +000083
84def _init_timeout(timeout=CONNECTION_TIMEOUT):
85 return time.time() + timeout
86
87def _check_timeout(t):
88 return time.time() > t
89
Benjamin Petersone711caf2008-06-11 16:44:04 +000090#
91#
92#
93
94def arbitrary_address(family):
95 '''
96 Return an arbitrary free address for the given family
97 '''
98 if family == 'AF_INET':
99 return ('localhost', 0)
100 elif family == 'AF_UNIX':
101 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
102 elif family == 'AF_PIPE':
103 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
104 (os.getpid(), next(_mmap_counter)))
105 else:
106 raise ValueError('unrecognized family')
107
Antoine Pitrou709176f2012-04-01 17:19:09 +0200108def _validate_family(family):
109 '''
110 Checks if the family is valid for the current environment.
111 '''
112 if sys.platform != 'win32' and family == 'AF_PIPE':
113 raise ValueError('Family %s is not recognized.' % family)
114
Antoine Pitrou6d20cba2012-04-03 20:12:23 +0200115 if sys.platform == 'win32' and family == 'AF_UNIX':
116 # double check
117 if not hasattr(socket, family):
118 raise ValueError('Family %s is not recognized.' % family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000119
120def address_type(address):
121 '''
122 Return the types of the address
123
124 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
125 '''
126 if type(address) == tuple:
127 return 'AF_INET'
128 elif type(address) is str and address.startswith('\\\\'):
129 return 'AF_PIPE'
130 elif type(address) is str:
131 return 'AF_UNIX'
132 else:
133 raise ValueError('address type of %r unrecognized' % address)
134
135#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200136# Connection classes
137#
138
139class _ConnectionBase:
140 _handle = None
141
142 def __init__(self, handle, readable=True, writable=True):
143 handle = handle.__index__()
144 if handle < 0:
145 raise ValueError("invalid handle")
146 if not readable and not writable:
147 raise ValueError(
148 "at least one of `readable` and `writable` must be True")
149 self._handle = handle
150 self._readable = readable
151 self._writable = writable
152
Antoine Pitrou60001202011-07-09 01:03:46 +0200153 # XXX should we use util.Finalize instead of a __del__?
154
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200155 def __del__(self):
156 if self._handle is not None:
157 self._close()
158
159 def _check_closed(self):
160 if self._handle is None:
161 raise IOError("handle is closed")
162
163 def _check_readable(self):
164 if not self._readable:
165 raise IOError("connection is write-only")
166
167 def _check_writable(self):
168 if not self._writable:
169 raise IOError("connection is read-only")
170
171 def _bad_message_length(self):
172 if self._writable:
173 self._readable = False
174 else:
175 self.close()
176 raise IOError("bad message length")
177
178 @property
179 def closed(self):
180 """True if the connection is closed"""
181 return self._handle is None
182
183 @property
184 def readable(self):
185 """True if the connection is readable"""
186 return self._readable
187
188 @property
189 def writable(self):
190 """True if the connection is writable"""
191 return self._writable
192
193 def fileno(self):
194 """File descriptor or handle of the connection"""
195 self._check_closed()
196 return self._handle
197
198 def close(self):
199 """Close the connection"""
200 if self._handle is not None:
201 try:
202 self._close()
203 finally:
204 self._handle = None
205
206 def send_bytes(self, buf, offset=0, size=None):
207 """Send the bytes data from a bytes-like object"""
208 self._check_closed()
209 self._check_writable()
210 m = memoryview(buf)
211 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
212 if m.itemsize > 1:
213 m = memoryview(bytes(m))
214 n = len(m)
215 if offset < 0:
216 raise ValueError("offset is negative")
217 if n < offset:
218 raise ValueError("buffer length < offset")
219 if size is None:
220 size = n - offset
221 elif size < 0:
222 raise ValueError("size is negative")
223 elif offset + size > n:
224 raise ValueError("buffer length < offset + size")
225 self._send_bytes(m[offset:offset + size])
226
227 def send(self, obj):
228 """Send a (picklable) object"""
229 self._check_closed()
230 self._check_writable()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200231 buf = io.BytesIO()
232 ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj)
233 self._send_bytes(buf.getbuffer())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200234
235 def recv_bytes(self, maxlength=None):
236 """
237 Receive bytes data as a bytes object.
238 """
239 self._check_closed()
240 self._check_readable()
241 if maxlength is not None and maxlength < 0:
242 raise ValueError("negative maxlength")
243 buf = self._recv_bytes(maxlength)
244 if buf is None:
245 self._bad_message_length()
246 return buf.getvalue()
247
248 def recv_bytes_into(self, buf, offset=0):
249 """
250 Receive bytes data into a writeable buffer-like object.
251 Return the number of bytes read.
252 """
253 self._check_closed()
254 self._check_readable()
255 with memoryview(buf) as m:
256 # Get bytesize of arbitrary buffer
257 itemsize = m.itemsize
258 bytesize = itemsize * len(m)
259 if offset < 0:
260 raise ValueError("negative offset")
261 elif offset > bytesize:
262 raise ValueError("offset too large")
263 result = self._recv_bytes()
264 size = result.tell()
265 if bytesize < offset + size:
266 raise BufferTooShort(result.getvalue())
267 # Message can fit in dest
268 result.seek(0)
269 result.readinto(m[offset // itemsize :
270 (offset + size) // itemsize])
271 return size
272
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100273 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200274 """Receive a (picklable) object"""
275 self._check_closed()
276 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100277 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200278 return pickle.loads(buf.getbuffer())
279
280 def poll(self, timeout=0.0):
281 """Whether there is any input available to be read"""
282 self._check_closed()
283 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200284 return self._poll(timeout)
285
286
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200287if _winapi:
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200288
289 class PipeConnection(_ConnectionBase):
290 """
291 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200292 Overlapped I/O is used, so the handles must have been created
293 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200294 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100295 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200296
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200297 def _close(self, _CloseHandle=_winapi.CloseHandle):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200298 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200299
300 def _send_bytes(self, buf):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200301 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100302 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200303 if err == _winapi.ERROR_IO_PENDING:
304 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100305 [ov.event], False, INFINITE)
306 assert waitres == WAIT_OBJECT_0
307 except:
308 ov.cancel()
309 raise
310 finally:
311 nwritten, err = ov.GetOverlappedResult(True)
312 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200313 assert nwritten == len(buf)
314
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100315 def _recv_bytes(self, maxsize=None):
316 if self._got_empty_message:
317 self._got_empty_message = False
318 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200319 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100320 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200321 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200322 ov, err = _winapi.ReadFile(self._handle, bsize,
323 overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100324 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200325 if err == _winapi.ERROR_IO_PENDING:
326 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100327 [ov.event], False, INFINITE)
328 assert waitres == WAIT_OBJECT_0
329 except:
330 ov.cancel()
331 raise
332 finally:
333 nread, err = ov.GetOverlappedResult(True)
334 if err == 0:
335 f = io.BytesIO()
336 f.write(ov.getbuffer())
337 return f
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200338 elif err == _winapi.ERROR_MORE_DATA:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100339 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200340 except IOError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200341 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200342 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100343 else:
344 raise
345 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200346
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100347 def _poll(self, timeout):
348 if (self._got_empty_message or
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200349 _winapi.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200350 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100351 if timeout < 0:
352 timeout = None
353 return bool(wait([self], timeout))
354
355 def _get_more_data(self, ov, maxsize):
356 buf = ov.getbuffer()
357 f = io.BytesIO()
358 f.write(buf)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200359 left = _winapi.PeekNamedPipe(self._handle)[1]
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100360 assert left > 0
361 if maxsize is not None and len(buf) + left > maxsize:
362 self._bad_message_length()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200363 ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100364 rbytes, err = ov.GetOverlappedResult(True)
365 assert err == 0
366 assert rbytes == left
367 f.write(ov.getbuffer())
368 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200369
370
371class Connection(_ConnectionBase):
372 """
373 Connection class based on an arbitrary file descriptor (Unix only), or
374 a socket handle (Windows).
375 """
376
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200377 if _winapi:
378 def _close(self, _close=_multiprocessing.closesocket):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200379 _close(self._handle)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200380 _write = _multiprocessing.send
381 _read = _multiprocessing.recv
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200382 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200383 def _close(self, _close=os.close):
384 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200385 _write = os.write
386 _read = os.read
387
388 def _send(self, buf, write=_write):
389 remaining = len(buf)
390 while True:
391 n = write(self._handle, buf)
392 remaining -= n
393 if remaining == 0:
394 break
395 buf = buf[n:]
396
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100397 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200398 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200399 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200400 remaining = size
401 while remaining > 0:
Antoine Pitroudd696492011-06-08 17:21:55 +0200402 chunk = read(handle, remaining)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200403 n = len(chunk)
404 if n == 0:
405 if remaining == size:
406 raise EOFError
407 else:
408 raise IOError("got end of file during message")
409 buf.write(chunk)
410 remaining -= n
411 return buf
412
413 def _send_bytes(self, buf):
414 # For wire compatibility with 3.2 and lower
415 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200416 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200417 # The condition is necessary to avoid "broken pipe" errors
418 # when sending a 0-length buffer if the other end closed the pipe.
419 if n > 0:
420 self._send(buf)
421
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100422 def _recv_bytes(self, maxsize=None):
423 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200424 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200425 if maxsize is not None and size > maxsize:
426 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100427 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200428
429 def _poll(self, timeout):
Antoine Pitroudd696492011-06-08 17:21:55 +0200430 if timeout < 0.0:
431 timeout = None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100432 r = wait([self._handle], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200433 return bool(r)
434
435
436#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000437# Public functions
438#
439
440class Listener(object):
441 '''
442 Returns a listener object.
443
444 This is a wrapper for a bound socket which is 'listening' for
445 connections, or for a Windows named pipe.
446 '''
447 def __init__(self, address=None, family=None, backlog=1, authkey=None):
448 family = family or (address and address_type(address)) \
449 or default_family
450 address = address or arbitrary_address(family)
451
Antoine Pitrou709176f2012-04-01 17:19:09 +0200452 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000453 if family == 'AF_PIPE':
454 self._listener = PipeListener(address, backlog)
455 else:
456 self._listener = SocketListener(address, family, backlog)
457
458 if authkey is not None and not isinstance(authkey, bytes):
459 raise TypeError('authkey should be a byte string')
460
461 self._authkey = authkey
462
463 def accept(self):
464 '''
465 Accept a connection on the bound socket or named pipe of `self`.
466
467 Returns a `Connection` object.
468 '''
469 c = self._listener.accept()
470 if self._authkey:
471 deliver_challenge(c, self._authkey)
472 answer_challenge(c, self._authkey)
473 return c
474
475 def close(self):
476 '''
477 Close the bound socket or named pipe of `self`.
478 '''
479 return self._listener.close()
480
481 address = property(lambda self: self._listener._address)
482 last_accepted = property(lambda self: self._listener._last_accepted)
483
484
485def Client(address, family=None, authkey=None):
486 '''
487 Returns a connection to the address of a `Listener`
488 '''
489 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200490 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000491 if family == 'AF_PIPE':
492 c = PipeClient(address)
493 else:
494 c = SocketClient(address)
495
496 if authkey is not None and not isinstance(authkey, bytes):
497 raise TypeError('authkey should be a byte string')
498
499 if authkey is not None:
500 answer_challenge(c, authkey)
501 deliver_challenge(c, authkey)
502
503 return c
504
505
506if sys.platform != 'win32':
507
508 def Pipe(duplex=True):
509 '''
510 Returns pair of connection objects at either end of a pipe
511 '''
512 if duplex:
513 s1, s2 = socket.socketpair()
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200514 c1 = Connection(s1.detach())
515 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000516 else:
517 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200518 c1 = Connection(fd1, writable=False)
519 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000520
521 return c1, c2
522
523else:
524
Benjamin Petersone711caf2008-06-11 16:44:04 +0000525 def Pipe(duplex=True):
526 '''
527 Returns pair of connection objects at either end of a pipe
528 '''
529 address = arbitrary_address('AF_PIPE')
530 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200531 openmode = _winapi.PIPE_ACCESS_DUPLEX
532 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000533 obsize, ibsize = BUFSIZE, BUFSIZE
534 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200535 openmode = _winapi.PIPE_ACCESS_INBOUND
536 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000537 obsize, ibsize = 0, BUFSIZE
538
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200539 h1 = _winapi.CreateNamedPipe(
540 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
541 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
542 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
543 _winapi.PIPE_WAIT,
544 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000545 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200546 h2 = _winapi.CreateFile(
547 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
548 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000549 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200550 _winapi.SetNamedPipeHandleState(
551 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000552 )
553
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200554 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100555 _, err = overlapped.GetOverlappedResult(True)
556 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000557
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200558 c1 = PipeConnection(h1, writable=duplex)
559 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000560
561 return c1, c2
562
563#
564# Definitions for connections based on sockets
565#
566
567class SocketListener(object):
568 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000569 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000570 '''
571 def __init__(self, address, family, backlog=1):
572 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100573 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100574 # SO_REUSEADDR has different semantics on Windows (issue #2550).
575 if os.name == 'posix':
576 self._socket.setsockopt(socket.SOL_SOCKET,
577 socket.SO_REUSEADDR, 1)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100578 self._socket.bind(address)
579 self._socket.listen(backlog)
580 self._address = self._socket.getsockname()
581 except OSError:
582 self._socket.close()
583 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000584 self._family = family
585 self._last_accepted = None
586
Benjamin Petersone711caf2008-06-11 16:44:04 +0000587 if family == 'AF_UNIX':
588 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000589 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000590 )
591 else:
592 self._unlink = None
593
594 def accept(self):
595 s, self._last_accepted = self._socket.accept()
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200596 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000597
598 def close(self):
599 self._socket.close()
600 if self._unlink is not None:
601 self._unlink()
602
603
604def SocketClient(address):
605 '''
606 Return a connection object connected to the socket given by `address`
607 '''
608 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000609 with socket.socket( getattr(socket, family) ) as s:
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100610 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200611 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000612
613#
614# Definitions for connections based on named pipes
615#
616
617if sys.platform == 'win32':
618
619 class PipeListener(object):
620 '''
621 Representation of a named pipe
622 '''
623 def __init__(self, address, backlog=None):
624 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100625 self._handle_queue = [self._new_handle(first=True)]
626
Benjamin Petersone711caf2008-06-11 16:44:04 +0000627 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000628 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000629 self.close = Finalize(
630 self, PipeListener._finalize_pipe_listener,
631 args=(self._handle_queue, self._address), exitpriority=0
632 )
633
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100634 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200635 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100636 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200637 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
638 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100639 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200640 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
641 _winapi.PIPE_WAIT,
642 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
643 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000644 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100645
646 def accept(self):
647 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000648 handle = self._handle_queue.pop(0)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200649 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000650 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200651 res = _winapi.WaitForMultipleObjects([ov.event], False, INFINITE)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100652 except:
653 ov.cancel()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200654 _winapi.CloseHandle(handle)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100655 raise
656 finally:
657 _, err = ov.GetOverlappedResult(True)
658 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200659 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000660
661 @staticmethod
662 def _finalize_pipe_listener(queue, address):
663 sub_debug('closing listener with address=%r', address)
664 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200665 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000666
667 def PipeClient(address):
668 '''
669 Return a connection object connected to the pipe given by `address`
670 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000671 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000672 while 1:
673 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200674 _winapi.WaitNamedPipe(address, 1000)
675 h = _winapi.CreateFile(
676 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
677 0, _winapi.NULL, _winapi.OPEN_EXISTING,
678 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000679 )
680 except WindowsError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200681 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
682 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000683 raise
684 else:
685 break
686 else:
687 raise
688
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200689 _winapi.SetNamedPipeHandleState(
690 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000691 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200692 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000693
694#
695# Authentication stuff
696#
697
698MESSAGE_LENGTH = 20
699
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000700CHALLENGE = b'#CHALLENGE#'
701WELCOME = b'#WELCOME#'
702FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000703
704def deliver_challenge(connection, authkey):
705 import hmac
706 assert isinstance(authkey, bytes)
707 message = os.urandom(MESSAGE_LENGTH)
708 connection.send_bytes(CHALLENGE + message)
709 digest = hmac.new(authkey, message).digest()
710 response = connection.recv_bytes(256) # reject large message
711 if response == digest:
712 connection.send_bytes(WELCOME)
713 else:
714 connection.send_bytes(FAILURE)
715 raise AuthenticationError('digest received was wrong')
716
717def answer_challenge(connection, authkey):
718 import hmac
719 assert isinstance(authkey, bytes)
720 message = connection.recv_bytes(256) # reject large message
721 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
722 message = message[len(CHALLENGE):]
723 digest = hmac.new(authkey, message).digest()
724 connection.send_bytes(digest)
725 response = connection.recv_bytes(256) # reject large message
726 if response != WELCOME:
727 raise AuthenticationError('digest sent was rejected')
728
729#
730# Support for using xmlrpclib for serialization
731#
732
733class ConnectionWrapper(object):
734 def __init__(self, conn, dumps, loads):
735 self._conn = conn
736 self._dumps = dumps
737 self._loads = loads
738 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
739 obj = getattr(conn, attr)
740 setattr(self, attr, obj)
741 def send(self, obj):
742 s = self._dumps(obj)
743 self._conn.send_bytes(s)
744 def recv(self):
745 s = self._conn.recv_bytes()
746 return self._loads(s)
747
748def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000749 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000750
751def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000752 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000753 return obj
754
755class XmlListener(Listener):
756 def accept(self):
757 global xmlrpclib
758 import xmlrpc.client as xmlrpclib
759 obj = Listener.accept(self)
760 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
761
762def XmlClient(*args, **kwds):
763 global xmlrpclib
764 import xmlrpc.client as xmlrpclib
765 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200766
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100767#
768# Wait
769#
770
771if sys.platform == 'win32':
772
773 def _exhaustive_wait(handles, timeout):
774 # Return ALL handles which are currently signalled. (Only
775 # returning the first signalled might create starvation issues.)
776 L = list(handles)
777 ready = []
778 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200779 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100780 if res == WAIT_TIMEOUT:
781 break
782 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
783 res -= WAIT_OBJECT_0
784 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
785 res -= WAIT_ABANDONED_0
786 else:
787 raise RuntimeError('Should not get here')
788 ready.append(L[res])
789 L = L[res+1:]
790 timeout = 0
791 return ready
792
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200793 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100794
795 def wait(object_list, timeout=None):
796 '''
797 Wait till an object in object_list is ready/readable.
798
799 Returns list of those objects in object_list which are ready/readable.
800 '''
801 if timeout is None:
802 timeout = INFINITE
803 elif timeout < 0:
804 timeout = 0
805 else:
806 timeout = int(timeout * 1000 + 0.5)
807
808 object_list = list(object_list)
809 waithandle_to_obj = {}
810 ov_list = []
811 ready_objects = set()
812 ready_handles = set()
813
814 try:
815 for o in object_list:
816 try:
817 fileno = getattr(o, 'fileno')
818 except AttributeError:
819 waithandle_to_obj[o.__index__()] = o
820 else:
821 # start an overlapped read of length zero
822 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200823 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100824 except OSError as e:
825 err = e.winerror
826 if err not in _ready_errors:
827 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200828 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100829 ov_list.append(ov)
830 waithandle_to_obj[ov.event] = o
831 else:
832 # If o.fileno() is an overlapped pipe handle and
833 # err == 0 then there is a zero length message
834 # in the pipe, but it HAS NOT been consumed.
835 ready_objects.add(o)
836 timeout = 0
837
838 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
839 finally:
840 # request that overlapped reads stop
841 for ov in ov_list:
842 ov.cancel()
843
844 # wait for all overlapped reads to stop
845 for ov in ov_list:
846 try:
847 _, err = ov.GetOverlappedResult(True)
848 except OSError as e:
849 err = e.winerror
850 if err not in _ready_errors:
851 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200852 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100853 o = waithandle_to_obj[ov.event]
854 ready_objects.add(o)
855 if err == 0:
856 # If o.fileno() is an overlapped pipe handle then
857 # a zero length message HAS been consumed.
858 if hasattr(o, '_got_empty_message'):
859 o._got_empty_message = True
860
861 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
862 return [o for o in object_list if o in ready_objects]
863
864else:
865
866 def wait(object_list, timeout=None):
867 '''
868 Wait till an object in object_list is ready/readable.
869
870 Returns list of those objects in object_list which are ready/readable.
871 '''
872 if timeout is not None:
873 if timeout <= 0:
874 return select.select(object_list, [], [], 0)[0]
875 else:
876 deadline = time.time() + timeout
877 while True:
878 try:
879 return select.select(object_list, [], [], timeout)[0]
880 except OSError as e:
881 if e.errno != errno.EINTR:
882 raise
883 if timeout is not None:
884 timeout = deadline - time.time()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200885
886#
887# Make connection and socket objects sharable if possible
888#
889
890if sys.platform == 'win32':
891 from . import reduction
892 ForkingPickler.register(socket.socket, reduction.reduce_socket)
893 ForkingPickler.register(Connection, reduction.reduce_connection)
894 ForkingPickler.register(PipeConnection, reduction.reduce_pipe_connection)
895else:
896 try:
897 from . import reduction
898 except ImportError:
899 pass
900 else:
901 ForkingPickler.register(socket.socket, reduction.reduce_socket)
902 ForkingPickler.register(Connection, reduction.reduce_connection)