blob: 3a61e5ee42a93e5a91b98dec64c68fee3485ff03 [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:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020054 import _winapi
55 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Antoine Pitrou87cf2202011-05-09 17:04:27 +020056except ImportError:
57 if sys.platform == 'win32':
58 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020059 _winapi = 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
Antoine Pitrou6d20cba2012-04-03 20:12:23 +0200114 if sys.platform == 'win32' and family == 'AF_UNIX':
115 # double check
116 if not hasattr(socket, family):
117 raise ValueError('Family %s is not recognized.' % family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000118
119def address_type(address):
120 '''
121 Return the types of the address
122
123 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
124 '''
125 if type(address) == tuple:
126 return 'AF_INET'
127 elif type(address) is str and address.startswith('\\\\'):
128 return 'AF_PIPE'
129 elif type(address) is str:
130 return 'AF_UNIX'
131 else:
132 raise ValueError('address type of %r unrecognized' % address)
133
134#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200135# Connection classes
136#
137
138class _ConnectionBase:
139 _handle = None
140
141 def __init__(self, handle, readable=True, writable=True):
142 handle = handle.__index__()
143 if handle < 0:
144 raise ValueError("invalid handle")
145 if not readable and not writable:
146 raise ValueError(
147 "at least one of `readable` and `writable` must be True")
148 self._handle = handle
149 self._readable = readable
150 self._writable = writable
151
Antoine Pitrou60001202011-07-09 01:03:46 +0200152 # XXX should we use util.Finalize instead of a __del__?
153
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200154 def __del__(self):
155 if self._handle is not None:
156 self._close()
157
158 def _check_closed(self):
159 if self._handle is None:
160 raise IOError("handle is closed")
161
162 def _check_readable(self):
163 if not self._readable:
164 raise IOError("connection is write-only")
165
166 def _check_writable(self):
167 if not self._writable:
168 raise IOError("connection is read-only")
169
170 def _bad_message_length(self):
171 if self._writable:
172 self._readable = False
173 else:
174 self.close()
175 raise IOError("bad message length")
176
177 @property
178 def closed(self):
179 """True if the connection is closed"""
180 return self._handle is None
181
182 @property
183 def readable(self):
184 """True if the connection is readable"""
185 return self._readable
186
187 @property
188 def writable(self):
189 """True if the connection is writable"""
190 return self._writable
191
192 def fileno(self):
193 """File descriptor or handle of the connection"""
194 self._check_closed()
195 return self._handle
196
197 def close(self):
198 """Close the connection"""
199 if self._handle is not None:
200 try:
201 self._close()
202 finally:
203 self._handle = None
204
205 def send_bytes(self, buf, offset=0, size=None):
206 """Send the bytes data from a bytes-like object"""
207 self._check_closed()
208 self._check_writable()
209 m = memoryview(buf)
210 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
211 if m.itemsize > 1:
212 m = memoryview(bytes(m))
213 n = len(m)
214 if offset < 0:
215 raise ValueError("offset is negative")
216 if n < offset:
217 raise ValueError("buffer length < offset")
218 if size is None:
219 size = n - offset
220 elif size < 0:
221 raise ValueError("size is negative")
222 elif offset + size > n:
223 raise ValueError("buffer length < offset + size")
224 self._send_bytes(m[offset:offset + size])
225
226 def send(self, obj):
227 """Send a (picklable) object"""
228 self._check_closed()
229 self._check_writable()
230 buf = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
231 self._send_bytes(memoryview(buf))
232
233 def recv_bytes(self, maxlength=None):
234 """
235 Receive bytes data as a bytes object.
236 """
237 self._check_closed()
238 self._check_readable()
239 if maxlength is not None and maxlength < 0:
240 raise ValueError("negative maxlength")
241 buf = self._recv_bytes(maxlength)
242 if buf is None:
243 self._bad_message_length()
244 return buf.getvalue()
245
246 def recv_bytes_into(self, buf, offset=0):
247 """
248 Receive bytes data into a writeable buffer-like object.
249 Return the number of bytes read.
250 """
251 self._check_closed()
252 self._check_readable()
253 with memoryview(buf) as m:
254 # Get bytesize of arbitrary buffer
255 itemsize = m.itemsize
256 bytesize = itemsize * len(m)
257 if offset < 0:
258 raise ValueError("negative offset")
259 elif offset > bytesize:
260 raise ValueError("offset too large")
261 result = self._recv_bytes()
262 size = result.tell()
263 if bytesize < offset + size:
264 raise BufferTooShort(result.getvalue())
265 # Message can fit in dest
266 result.seek(0)
267 result.readinto(m[offset // itemsize :
268 (offset + size) // itemsize])
269 return size
270
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100271 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200272 """Receive a (picklable) object"""
273 self._check_closed()
274 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100275 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200276 return pickle.loads(buf.getbuffer())
277
278 def poll(self, timeout=0.0):
279 """Whether there is any input available to be read"""
280 self._check_closed()
281 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200282 return self._poll(timeout)
283
284
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200285if _winapi:
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200286
287 class PipeConnection(_ConnectionBase):
288 """
289 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200290 Overlapped I/O is used, so the handles must have been created
291 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200292 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100293 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200294
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200295 def _close(self, _CloseHandle=_winapi.CloseHandle):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200296 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200297
298 def _send_bytes(self, buf):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200299 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100300 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200301 if err == _winapi.ERROR_IO_PENDING:
302 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100303 [ov.event], False, INFINITE)
304 assert waitres == WAIT_OBJECT_0
305 except:
306 ov.cancel()
307 raise
308 finally:
309 nwritten, err = ov.GetOverlappedResult(True)
310 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200311 assert nwritten == len(buf)
312
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100313 def _recv_bytes(self, maxsize=None):
314 if self._got_empty_message:
315 self._got_empty_message = False
316 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200317 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100318 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200319 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200320 ov, err = _winapi.ReadFile(self._handle, bsize,
321 overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100322 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200323 if err == _winapi.ERROR_IO_PENDING:
324 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100325 [ov.event], False, INFINITE)
326 assert waitres == WAIT_OBJECT_0
327 except:
328 ov.cancel()
329 raise
330 finally:
331 nread, err = ov.GetOverlappedResult(True)
332 if err == 0:
333 f = io.BytesIO()
334 f.write(ov.getbuffer())
335 return f
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200336 elif err == _winapi.ERROR_MORE_DATA:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100337 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200338 except IOError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200339 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200340 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100341 else:
342 raise
343 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200344
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100345 def _poll(self, timeout):
346 if (self._got_empty_message or
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200347 _winapi.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200348 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100349 if timeout < 0:
350 timeout = None
351 return bool(wait([self], timeout))
352
353 def _get_more_data(self, ov, maxsize):
354 buf = ov.getbuffer()
355 f = io.BytesIO()
356 f.write(buf)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200357 left = _winapi.PeekNamedPipe(self._handle)[1]
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100358 assert left > 0
359 if maxsize is not None and len(buf) + left > maxsize:
360 self._bad_message_length()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200361 ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100362 rbytes, err = ov.GetOverlappedResult(True)
363 assert err == 0
364 assert rbytes == left
365 f.write(ov.getbuffer())
366 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200367
368
369class Connection(_ConnectionBase):
370 """
371 Connection class based on an arbitrary file descriptor (Unix only), or
372 a socket handle (Windows).
373 """
374
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200375 if _winapi:
376 def _close(self, _close=_multiprocessing.closesocket):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200377 _close(self._handle)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200378 _write = _multiprocessing.send
379 _read = _multiprocessing.recv
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200380 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200381 def _close(self, _close=os.close):
382 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200383 _write = os.write
384 _read = os.read
385
386 def _send(self, buf, write=_write):
387 remaining = len(buf)
388 while True:
389 n = write(self._handle, buf)
390 remaining -= n
391 if remaining == 0:
392 break
393 buf = buf[n:]
394
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100395 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200396 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200397 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200398 remaining = size
399 while remaining > 0:
Antoine Pitroudd696492011-06-08 17:21:55 +0200400 chunk = read(handle, remaining)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200401 n = len(chunk)
402 if n == 0:
403 if remaining == size:
404 raise EOFError
405 else:
406 raise IOError("got end of file during message")
407 buf.write(chunk)
408 remaining -= n
409 return buf
410
411 def _send_bytes(self, buf):
412 # For wire compatibility with 3.2 and lower
413 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200414 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200415 # The condition is necessary to avoid "broken pipe" errors
416 # when sending a 0-length buffer if the other end closed the pipe.
417 if n > 0:
418 self._send(buf)
419
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100420 def _recv_bytes(self, maxsize=None):
421 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200422 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200423 if maxsize is not None and size > maxsize:
424 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100425 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200426
427 def _poll(self, timeout):
Antoine Pitroudd696492011-06-08 17:21:55 +0200428 if timeout < 0.0:
429 timeout = None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100430 r = wait([self._handle], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200431 return bool(r)
432
433
434#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000435# Public functions
436#
437
438class Listener(object):
439 '''
440 Returns a listener object.
441
442 This is a wrapper for a bound socket which is 'listening' for
443 connections, or for a Windows named pipe.
444 '''
445 def __init__(self, address=None, family=None, backlog=1, authkey=None):
446 family = family or (address and address_type(address)) \
447 or default_family
448 address = address or arbitrary_address(family)
449
Antoine Pitrou709176f2012-04-01 17:19:09 +0200450 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000451 if family == 'AF_PIPE':
452 self._listener = PipeListener(address, backlog)
453 else:
454 self._listener = SocketListener(address, family, backlog)
455
456 if authkey is not None and not isinstance(authkey, bytes):
457 raise TypeError('authkey should be a byte string')
458
459 self._authkey = authkey
460
461 def accept(self):
462 '''
463 Accept a connection on the bound socket or named pipe of `self`.
464
465 Returns a `Connection` object.
466 '''
467 c = self._listener.accept()
468 if self._authkey:
469 deliver_challenge(c, self._authkey)
470 answer_challenge(c, self._authkey)
471 return c
472
473 def close(self):
474 '''
475 Close the bound socket or named pipe of `self`.
476 '''
477 return self._listener.close()
478
479 address = property(lambda self: self._listener._address)
480 last_accepted = property(lambda self: self._listener._last_accepted)
481
482
483def Client(address, family=None, authkey=None):
484 '''
485 Returns a connection to the address of a `Listener`
486 '''
487 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200488 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000489 if family == 'AF_PIPE':
490 c = PipeClient(address)
491 else:
492 c = SocketClient(address)
493
494 if authkey is not None and not isinstance(authkey, bytes):
495 raise TypeError('authkey should be a byte string')
496
497 if authkey is not None:
498 answer_challenge(c, authkey)
499 deliver_challenge(c, authkey)
500
501 return c
502
503
504if sys.platform != 'win32':
505
506 def Pipe(duplex=True):
507 '''
508 Returns pair of connection objects at either end of a pipe
509 '''
510 if duplex:
511 s1, s2 = socket.socketpair()
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200512 c1 = Connection(s1.detach())
513 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000514 else:
515 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200516 c1 = Connection(fd1, writable=False)
517 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000518
519 return c1, c2
520
521else:
522
Benjamin Petersone711caf2008-06-11 16:44:04 +0000523 def Pipe(duplex=True):
524 '''
525 Returns pair of connection objects at either end of a pipe
526 '''
527 address = arbitrary_address('AF_PIPE')
528 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200529 openmode = _winapi.PIPE_ACCESS_DUPLEX
530 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000531 obsize, ibsize = BUFSIZE, BUFSIZE
532 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200533 openmode = _winapi.PIPE_ACCESS_INBOUND
534 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000535 obsize, ibsize = 0, BUFSIZE
536
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200537 h1 = _winapi.CreateNamedPipe(
538 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
539 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
540 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
541 _winapi.PIPE_WAIT,
542 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000543 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200544 h2 = _winapi.CreateFile(
545 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
546 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000547 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200548 _winapi.SetNamedPipeHandleState(
549 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000550 )
551
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200552 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100553 _, err = overlapped.GetOverlappedResult(True)
554 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000555
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200556 c1 = PipeConnection(h1, writable=duplex)
557 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000558
559 return c1, c2
560
561#
562# Definitions for connections based on sockets
563#
564
565class SocketListener(object):
566 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000567 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000568 '''
569 def __init__(self, address, family, backlog=1):
570 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100571 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100572 # SO_REUSEADDR has different semantics on Windows (issue #2550).
573 if os.name == 'posix':
574 self._socket.setsockopt(socket.SOL_SOCKET,
575 socket.SO_REUSEADDR, 1)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100576 self._socket.bind(address)
577 self._socket.listen(backlog)
578 self._address = self._socket.getsockname()
579 except OSError:
580 self._socket.close()
581 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000582 self._family = family
583 self._last_accepted = None
584
Benjamin Petersone711caf2008-06-11 16:44:04 +0000585 if family == 'AF_UNIX':
586 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000587 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000588 )
589 else:
590 self._unlink = None
591
592 def accept(self):
593 s, self._last_accepted = self._socket.accept()
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200594 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000595
596 def close(self):
597 self._socket.close()
598 if self._unlink is not None:
599 self._unlink()
600
601
602def SocketClient(address):
603 '''
604 Return a connection object connected to the socket given by `address`
605 '''
606 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000607 with socket.socket( getattr(socket, family) ) as s:
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100608 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200609 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000610
611#
612# Definitions for connections based on named pipes
613#
614
615if sys.platform == 'win32':
616
617 class PipeListener(object):
618 '''
619 Representation of a named pipe
620 '''
621 def __init__(self, address, backlog=None):
622 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100623 self._handle_queue = [self._new_handle(first=True)]
624
Benjamin Petersone711caf2008-06-11 16:44:04 +0000625 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000626 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000627 self.close = Finalize(
628 self, PipeListener._finalize_pipe_listener,
629 args=(self._handle_queue, self._address), exitpriority=0
630 )
631
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100632 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200633 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100634 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200635 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
636 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100637 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200638 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
639 _winapi.PIPE_WAIT,
640 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
641 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000642 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100643
644 def accept(self):
645 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000646 handle = self._handle_queue.pop(0)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200647 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000648 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200649 res = _winapi.WaitForMultipleObjects([ov.event], False, INFINITE)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100650 except:
651 ov.cancel()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200652 _winapi.CloseHandle(handle)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100653 raise
654 finally:
655 _, err = ov.GetOverlappedResult(True)
656 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200657 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000658
659 @staticmethod
660 def _finalize_pipe_listener(queue, address):
661 sub_debug('closing listener with address=%r', address)
662 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200663 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000664
665 def PipeClient(address):
666 '''
667 Return a connection object connected to the pipe given by `address`
668 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000669 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000670 while 1:
671 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200672 _winapi.WaitNamedPipe(address, 1000)
673 h = _winapi.CreateFile(
674 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
675 0, _winapi.NULL, _winapi.OPEN_EXISTING,
676 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000677 )
678 except WindowsError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200679 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
680 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000681 raise
682 else:
683 break
684 else:
685 raise
686
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200687 _winapi.SetNamedPipeHandleState(
688 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000689 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200690 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000691
692#
693# Authentication stuff
694#
695
696MESSAGE_LENGTH = 20
697
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000698CHALLENGE = b'#CHALLENGE#'
699WELCOME = b'#WELCOME#'
700FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000701
702def deliver_challenge(connection, authkey):
703 import hmac
704 assert isinstance(authkey, bytes)
705 message = os.urandom(MESSAGE_LENGTH)
706 connection.send_bytes(CHALLENGE + message)
707 digest = hmac.new(authkey, message).digest()
708 response = connection.recv_bytes(256) # reject large message
709 if response == digest:
710 connection.send_bytes(WELCOME)
711 else:
712 connection.send_bytes(FAILURE)
713 raise AuthenticationError('digest received was wrong')
714
715def answer_challenge(connection, authkey):
716 import hmac
717 assert isinstance(authkey, bytes)
718 message = connection.recv_bytes(256) # reject large message
719 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
720 message = message[len(CHALLENGE):]
721 digest = hmac.new(authkey, message).digest()
722 connection.send_bytes(digest)
723 response = connection.recv_bytes(256) # reject large message
724 if response != WELCOME:
725 raise AuthenticationError('digest sent was rejected')
726
727#
728# Support for using xmlrpclib for serialization
729#
730
731class ConnectionWrapper(object):
732 def __init__(self, conn, dumps, loads):
733 self._conn = conn
734 self._dumps = dumps
735 self._loads = loads
736 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
737 obj = getattr(conn, attr)
738 setattr(self, attr, obj)
739 def send(self, obj):
740 s = self._dumps(obj)
741 self._conn.send_bytes(s)
742 def recv(self):
743 s = self._conn.recv_bytes()
744 return self._loads(s)
745
746def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000747 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000748
749def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000750 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000751 return obj
752
753class XmlListener(Listener):
754 def accept(self):
755 global xmlrpclib
756 import xmlrpc.client as xmlrpclib
757 obj = Listener.accept(self)
758 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
759
760def XmlClient(*args, **kwds):
761 global xmlrpclib
762 import xmlrpc.client as xmlrpclib
763 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200764
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100765#
766# Wait
767#
768
769if sys.platform == 'win32':
770
771 def _exhaustive_wait(handles, timeout):
772 # Return ALL handles which are currently signalled. (Only
773 # returning the first signalled might create starvation issues.)
774 L = list(handles)
775 ready = []
776 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200777 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100778 if res == WAIT_TIMEOUT:
779 break
780 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
781 res -= WAIT_OBJECT_0
782 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
783 res -= WAIT_ABANDONED_0
784 else:
785 raise RuntimeError('Should not get here')
786 ready.append(L[res])
787 L = L[res+1:]
788 timeout = 0
789 return ready
790
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200791 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100792
793 def wait(object_list, timeout=None):
794 '''
795 Wait till an object in object_list is ready/readable.
796
797 Returns list of those objects in object_list which are ready/readable.
798 '''
799 if timeout is None:
800 timeout = INFINITE
801 elif timeout < 0:
802 timeout = 0
803 else:
804 timeout = int(timeout * 1000 + 0.5)
805
806 object_list = list(object_list)
807 waithandle_to_obj = {}
808 ov_list = []
809 ready_objects = set()
810 ready_handles = set()
811
812 try:
813 for o in object_list:
814 try:
815 fileno = getattr(o, 'fileno')
816 except AttributeError:
817 waithandle_to_obj[o.__index__()] = o
818 else:
819 # start an overlapped read of length zero
820 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200821 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100822 except OSError as e:
823 err = e.winerror
824 if err not in _ready_errors:
825 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200826 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100827 ov_list.append(ov)
828 waithandle_to_obj[ov.event] = o
829 else:
830 # If o.fileno() is an overlapped pipe handle and
831 # err == 0 then there is a zero length message
832 # in the pipe, but it HAS NOT been consumed.
833 ready_objects.add(o)
834 timeout = 0
835
836 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
837 finally:
838 # request that overlapped reads stop
839 for ov in ov_list:
840 ov.cancel()
841
842 # wait for all overlapped reads to stop
843 for ov in ov_list:
844 try:
845 _, err = ov.GetOverlappedResult(True)
846 except OSError as e:
847 err = e.winerror
848 if err not in _ready_errors:
849 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200850 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100851 o = waithandle_to_obj[ov.event]
852 ready_objects.add(o)
853 if err == 0:
854 # If o.fileno() is an overlapped pipe handle then
855 # a zero length message HAS been consumed.
856 if hasattr(o, '_got_empty_message'):
857 o._got_empty_message = True
858
859 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
860 return [o for o in object_list if o in ready_objects]
861
862else:
863
864 def wait(object_list, timeout=None):
865 '''
866 Wait till an object in object_list is ready/readable.
867
868 Returns list of those objects in object_list which are ready/readable.
869 '''
870 if timeout is not None:
871 if timeout <= 0:
872 return select.select(object_list, [], [], 0)[0]
873 else:
874 deadline = time.time() + timeout
875 while True:
876 try:
877 return select.select(object_list, [], [], timeout)[0]
878 except OSError as e:
879 if e.errno != errno.EINTR:
880 raise
881 if timeout is not None:
882 timeout = deadline - time.time()