connection.py 32 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039
  1. # -*- coding: utf-8 -*-
  2. #
  3. # A higher level module for using sockets (or Windows named pipes)
  4. #
  5. # multiprocessing/connection.py
  6. #
  7. # Copyright (c) 2006-2008, R Oudkerk
  8. # Licensed to PSF under a Contributor Agreement.
  9. #
  10. from __future__ import absolute_import
  11. import errno
  12. import io
  13. import os
  14. import sys
  15. import socket
  16. import select
  17. import struct
  18. import tempfile
  19. import itertools
  20. from . import reduction
  21. from . import util
  22. from . import AuthenticationError, BufferTooShort
  23. from ._ext import _billiard
  24. from .compat import setblocking, send_offset
  25. from .five import monotonic
  26. from .reduction import ForkingPickler
  27. try:
  28. from .compat import _winapi
  29. except ImportError:
  30. if sys.platform == 'win32':
  31. raise
  32. _winapi = None
  33. else:
  34. if sys.platform == 'win32':
  35. WAIT_OBJECT_0 = _winapi.WAIT_OBJECT_0
  36. try:
  37. WAIT_ABANDONED_0 = _winapi.WAIT_ABANDONED_0
  38. except AttributeError:
  39. WAIT_ABANDONED_0 = 128 # noqa
  40. WAIT_TIMEOUT = _winapi.WAIT_TIMEOUT
  41. INFINITE = _winapi.INFINITE
  42. __all__ = ['Client', 'Listener', 'Pipe', 'wait']
  43. is_pypy = hasattr(sys, 'pypy_version_info')
  44. #
  45. #
  46. #
  47. BUFSIZE = 8192
  48. # A very generous timeout when it comes to local connections...
  49. CONNECTION_TIMEOUT = 20.
  50. _mmap_counter = itertools.count()
  51. default_family = 'AF_INET'
  52. families = ['AF_INET']
  53. if hasattr(socket, 'AF_UNIX'):
  54. default_family = 'AF_UNIX'
  55. families += ['AF_UNIX']
  56. if sys.platform == 'win32':
  57. default_family = 'AF_PIPE'
  58. families += ['AF_PIPE']
  59. def _init_timeout(timeout=CONNECTION_TIMEOUT):
  60. return monotonic() + timeout
  61. def _check_timeout(t):
  62. return monotonic() > t
  63. #
  64. #
  65. #
  66. def arbitrary_address(family):
  67. '''
  68. Return an arbitrary free address for the given family
  69. '''
  70. if family == 'AF_INET':
  71. return ('localhost', 0)
  72. elif family == 'AF_UNIX':
  73. return tempfile.mktemp(prefix='listener-', dir=util.get_temp_dir())
  74. elif family == 'AF_PIPE':
  75. return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
  76. (os.getpid(), next(_mmap_counter)), dir="")
  77. else:
  78. raise ValueError('unrecognized family')
  79. def _validate_family(family):
  80. '''
  81. Checks if the family is valid for the current environment.
  82. '''
  83. if sys.platform != 'win32' and family == 'AF_PIPE':
  84. raise ValueError('Family %s is not recognized.' % family)
  85. if sys.platform == 'win32' and family == 'AF_UNIX':
  86. # double check
  87. if not hasattr(socket, family):
  88. raise ValueError('Family %s is not recognized.' % family)
  89. def address_type(address):
  90. '''
  91. Return the types of the address
  92. This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
  93. '''
  94. if type(address) == tuple:
  95. return 'AF_INET'
  96. elif type(address) is str and address.startswith('\\\\'):
  97. return 'AF_PIPE'
  98. elif type(address) is str:
  99. return 'AF_UNIX'
  100. else:
  101. raise ValueError('address type of %r unrecognized' % address)
  102. #
  103. # Connection classes
  104. #
  105. class _SocketContainer(object):
  106. def __init__(self, sock):
  107. self.sock = sock
  108. class _ConnectionBase(object):
  109. _handle = None
  110. def __init__(self, handle, readable=True, writable=True):
  111. if isinstance(handle, _SocketContainer):
  112. self._socket = handle.sock # keep ref so not collected
  113. handle = handle.sock.fileno()
  114. handle = handle.__index__()
  115. if handle < 0:
  116. raise ValueError("invalid handle")
  117. if not readable and not writable:
  118. raise ValueError(
  119. "at least one of `readable` and `writable` must be True")
  120. self._handle = handle
  121. self._readable = readable
  122. self._writable = writable
  123. # XXX should we use util.Finalize instead of a __del__?
  124. def __del__(self):
  125. if self._handle is not None:
  126. self._close()
  127. def _check_closed(self):
  128. if self._handle is None:
  129. raise OSError("handle is closed")
  130. def _check_readable(self):
  131. if not self._readable:
  132. raise OSError("connection is write-only")
  133. def _check_writable(self):
  134. if not self._writable:
  135. raise OSError("connection is read-only")
  136. def _bad_message_length(self):
  137. if self._writable:
  138. self._readable = False
  139. else:
  140. self.close()
  141. raise OSError("bad message length")
  142. @property
  143. def closed(self):
  144. """True if the connection is closed"""
  145. return self._handle is None
  146. @property
  147. def readable(self):
  148. """True if the connection is readable"""
  149. return self._readable
  150. @property
  151. def writable(self):
  152. """True if the connection is writable"""
  153. return self._writable
  154. def fileno(self):
  155. """File descriptor or handle of the connection"""
  156. self._check_closed()
  157. return self._handle
  158. def close(self):
  159. """Close the connection"""
  160. if self._handle is not None:
  161. try:
  162. self._close()
  163. finally:
  164. self._handle = None
  165. def send_bytes(self, buf, offset=0, size=None):
  166. """Send the bytes data from a bytes-like object"""
  167. self._check_closed()
  168. self._check_writable()
  169. m = memoryview(buf)
  170. # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
  171. if m.itemsize > 1:
  172. m = memoryview(bytes(m))
  173. n = len(m)
  174. if offset < 0:
  175. raise ValueError("offset is negative")
  176. if n < offset:
  177. raise ValueError("buffer length < offset")
  178. if size is None:
  179. size = n - offset
  180. elif size < 0:
  181. raise ValueError("size is negative")
  182. elif offset + size > n:
  183. raise ValueError("buffer length < offset + size")
  184. self._send_bytes(m[offset:offset + size])
  185. def send(self, obj):
  186. """Send a (picklable) object"""
  187. self._check_closed()
  188. self._check_writable()
  189. self._send_bytes(ForkingPickler.dumps(obj))
  190. def recv_bytes(self, maxlength=None):
  191. """
  192. Receive bytes data as a bytes object.
  193. """
  194. self._check_closed()
  195. self._check_readable()
  196. if maxlength is not None and maxlength < 0:
  197. raise ValueError("negative maxlength")
  198. buf = self._recv_bytes(maxlength)
  199. if buf is None:
  200. self._bad_message_length()
  201. return buf.getvalue()
  202. def recv_bytes_into(self, buf, offset=0):
  203. """
  204. Receive bytes data into a writeable bytes-like object.
  205. Return the number of bytes read.
  206. """
  207. self._check_closed()
  208. self._check_readable()
  209. with memoryview(buf) as m:
  210. # Get bytesize of arbitrary buffer
  211. itemsize = m.itemsize
  212. bytesize = itemsize * len(m)
  213. if offset < 0:
  214. raise ValueError("negative offset")
  215. elif offset > bytesize:
  216. raise ValueError("offset too large")
  217. result = self._recv_bytes()
  218. size = result.tell()
  219. if bytesize < offset + size:
  220. raise BufferTooShort(result.getvalue())
  221. # Message can fit in dest
  222. result.seek(0)
  223. result.readinto(m[
  224. offset // itemsize:(offset + size) // itemsize
  225. ])
  226. return size
  227. def recv(self):
  228. """Receive a (picklable) object"""
  229. self._check_closed()
  230. self._check_readable()
  231. buf = self._recv_bytes()
  232. return ForkingPickler.loadbuf(buf)
  233. def poll(self, timeout=0.0):
  234. """Whether there is any input available to be read"""
  235. self._check_closed()
  236. self._check_readable()
  237. return self._poll(timeout)
  238. def __enter__(self):
  239. return self
  240. def __exit__(self, exc_type, exc_value, exc_tb):
  241. self.close()
  242. def send_offset(self, buf, offset):
  243. return send_offset(self.fileno(), buf, offset)
  244. def setblocking(self, blocking):
  245. setblocking(self.fileno(), blocking)
  246. if _winapi:
  247. class PipeConnection(_ConnectionBase):
  248. """
  249. Connection class based on a Windows named pipe.
  250. Overlapped I/O is used, so the handles must have been created
  251. with FILE_FLAG_OVERLAPPED.
  252. """
  253. _got_empty_message = False
  254. def _close(self, _CloseHandle=_winapi.CloseHandle):
  255. _CloseHandle(self._handle)
  256. def _send_bytes(self, buf):
  257. ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
  258. try:
  259. if err == _winapi.ERROR_IO_PENDING:
  260. waitres = _winapi.WaitForMultipleObjects(
  261. [ov.event], False, INFINITE)
  262. assert waitres == WAIT_OBJECT_0
  263. except:
  264. ov.cancel()
  265. raise
  266. finally:
  267. nwritten, err = ov.GetOverlappedResult(True)
  268. assert err == 0
  269. assert nwritten == len(buf)
  270. def _recv_bytes(self, maxsize=None):
  271. if self._got_empty_message:
  272. self._got_empty_message = False
  273. return io.BytesIO()
  274. else:
  275. bsize = 128 if maxsize is None else min(maxsize, 128)
  276. try:
  277. ov, err = _winapi.ReadFile(
  278. self._handle, bsize, overlapped=True,
  279. )
  280. try:
  281. if err == _winapi.ERROR_IO_PENDING:
  282. waitres = _winapi.WaitForMultipleObjects(
  283. [ov.event], False, INFINITE)
  284. assert waitres == WAIT_OBJECT_0
  285. except:
  286. ov.cancel()
  287. raise
  288. finally:
  289. nread, err = ov.GetOverlappedResult(True)
  290. if err == 0:
  291. f = io.BytesIO()
  292. f.write(ov.getbuffer())
  293. return f
  294. elif err == _winapi.ERROR_MORE_DATA:
  295. return self._get_more_data(ov, maxsize)
  296. except OSError as e:
  297. if e.winerror == _winapi.ERROR_BROKEN_PIPE:
  298. raise EOFError
  299. else:
  300. raise
  301. raise RuntimeError(
  302. "shouldn't get here; expected KeyboardInterrupt")
  303. def _poll(self, timeout):
  304. if (self._got_empty_message or
  305. _winapi.PeekNamedPipe(self._handle)[0] != 0):
  306. return True
  307. return bool(wait([self], timeout))
  308. def _get_more_data(self, ov, maxsize):
  309. buf = ov.getbuffer()
  310. f = io.BytesIO()
  311. f.write(buf)
  312. left = _winapi.PeekNamedPipe(self._handle)[1]
  313. assert left > 0
  314. if maxsize is not None and len(buf) + left > maxsize:
  315. self._bad_message_length()
  316. ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
  317. rbytes, err = ov.GetOverlappedResult(True)
  318. assert err == 0
  319. assert rbytes == left
  320. f.write(ov.getbuffer())
  321. return f
  322. class Connection(_ConnectionBase):
  323. """
  324. Connection class based on an arbitrary file descriptor (Unix only), or
  325. a socket handle (Windows).
  326. """
  327. if _winapi:
  328. def _close(self, _close=_billiard.closesocket):
  329. _close(self._handle)
  330. _write = _billiard.send
  331. _read = _billiard.recv
  332. else:
  333. def _close(self, _close=os.close):
  334. _close(self._handle)
  335. _write = os.write
  336. _read = os.read
  337. def _send(self, buf, write=_write):
  338. remaining = len(buf)
  339. while True:
  340. try:
  341. n = write(self._handle, buf)
  342. except (OSError, IOError, socket.error) as exc:
  343. if getattr(exc, 'errno', None) != errno.EINTR:
  344. raise
  345. else:
  346. remaining -= n
  347. if remaining == 0:
  348. break
  349. buf = buf[n:]
  350. def _recv(self, size, read=_read):
  351. buf = io.BytesIO()
  352. handle = self._handle
  353. remaining = size
  354. while remaining > 0:
  355. try:
  356. chunk = read(handle, remaining)
  357. except (OSError, IOError, socket.error) as exc:
  358. if getattr(exc, 'errno', None) != errno.EINTR:
  359. raise
  360. else:
  361. n = len(chunk)
  362. if n == 0:
  363. if remaining == size:
  364. raise EOFError
  365. else:
  366. raise OSError("got end of file during message")
  367. buf.write(chunk)
  368. remaining -= n
  369. return buf
  370. def _send_bytes(self, buf, memoryview=memoryview):
  371. n = len(buf)
  372. # For wire compatibility with 3.2 and lower
  373. header = struct.pack("!i", n)
  374. if n > 16384:
  375. # The payload is large so Nagle's algorithm won't be triggered
  376. # and we'd better avoid the cost of concatenation.
  377. self._send(header)
  378. self._send(buf)
  379. else:
  380. # Issue #20540: concatenate before sending, to avoid delays due
  381. # to Nagle's algorithm on a TCP socket.
  382. # Also note we want to avoid sending a 0-length buffer separately,
  383. # to avoid "broken pipe" errors if the other end closed the pipe.
  384. if isinstance(buf, memoryview):
  385. buf = buf.tobytes()
  386. self._send(header + buf)
  387. def _recv_bytes(self, maxsize=None):
  388. buf = self._recv(4)
  389. size, = struct.unpack("!i", buf.getvalue())
  390. if maxsize is not None and size > maxsize:
  391. return None
  392. return self._recv(size)
  393. def _poll(self, timeout):
  394. r = wait([self], timeout)
  395. return bool(r)
  396. #
  397. # Public functions
  398. #
  399. class Listener(object):
  400. '''
  401. Returns a listener object.
  402. This is a wrapper for a bound socket which is 'listening' for
  403. connections, or for a Windows named pipe.
  404. '''
  405. def __init__(self, address=None, family=None, backlog=1, authkey=None):
  406. family = (family or
  407. (address and address_type(address)) or default_family)
  408. address = address or arbitrary_address(family)
  409. _validate_family(family)
  410. if family == 'AF_PIPE':
  411. self._listener = PipeListener(address, backlog)
  412. else:
  413. self._listener = SocketListener(address, family, backlog)
  414. if authkey is not None and not isinstance(authkey, bytes):
  415. raise TypeError('authkey should be a byte string')
  416. self._authkey = authkey
  417. def accept(self):
  418. '''
  419. Accept a connection on the bound socket or named pipe of `self`.
  420. Returns a `Connection` object.
  421. '''
  422. if self._listener is None:
  423. raise OSError('listener is closed')
  424. c = self._listener.accept()
  425. if self._authkey:
  426. deliver_challenge(c, self._authkey)
  427. answer_challenge(c, self._authkey)
  428. return c
  429. def close(self):
  430. '''
  431. Close the bound socket or named pipe of `self`.
  432. '''
  433. listener = self._listener
  434. if listener is not None:
  435. self._listener = None
  436. listener.close()
  437. address = property(lambda self: self._listener._address)
  438. last_accepted = property(lambda self: self._listener._last_accepted)
  439. def __enter__(self):
  440. return self
  441. def __exit__(self, exc_type, exc_value, exc_tb):
  442. self.close()
  443. def Client(address, family=None, authkey=None):
  444. '''
  445. Returns a connection to the address of a `Listener`
  446. '''
  447. family = family or address_type(address)
  448. _validate_family(family)
  449. if family == 'AF_PIPE':
  450. c = PipeClient(address)
  451. else:
  452. c = SocketClient(address)
  453. if authkey is not None and not isinstance(authkey, bytes):
  454. raise TypeError('authkey should be a byte string')
  455. if authkey is not None:
  456. answer_challenge(c, authkey)
  457. deliver_challenge(c, authkey)
  458. return c
  459. def detach(sock):
  460. if hasattr(sock, 'detach'):
  461. return sock.detach()
  462. # older socket lib does not have detach. We'll keep a reference around
  463. # so that it does not get garbage collected.
  464. return _SocketContainer(sock)
  465. if sys.platform != 'win32':
  466. def Pipe(duplex=True, rnonblock=False, wnonblock=False):
  467. '''
  468. Returns pair of connection objects at either end of a pipe
  469. '''
  470. if duplex:
  471. s1, s2 = socket.socketpair()
  472. s1.setblocking(not rnonblock)
  473. s2.setblocking(not wnonblock)
  474. c1 = Connection(detach(s1))
  475. c2 = Connection(detach(s2))
  476. else:
  477. fd1, fd2 = os.pipe()
  478. if rnonblock:
  479. setblocking(fd1, 0)
  480. if wnonblock:
  481. setblocking(fd2, 0)
  482. c1 = Connection(fd1, writable=False)
  483. c2 = Connection(fd2, readable=False)
  484. return c1, c2
  485. else:
  486. def Pipe(duplex=True, rnonblock=False, wnonblock=False):
  487. '''
  488. Returns pair of connection objects at either end of a pipe
  489. '''
  490. assert not rnonblock, 'rnonblock not supported on windows'
  491. assert not wnonblock, 'wnonblock not supported on windows'
  492. address = arbitrary_address('AF_PIPE')
  493. if duplex:
  494. openmode = _winapi.PIPE_ACCESS_DUPLEX
  495. access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
  496. obsize, ibsize = BUFSIZE, BUFSIZE
  497. else:
  498. openmode = _winapi.PIPE_ACCESS_INBOUND
  499. access = _winapi.GENERIC_WRITE
  500. obsize, ibsize = 0, BUFSIZE
  501. h1 = _winapi.CreateNamedPipe(
  502. address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
  503. _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
  504. _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
  505. _winapi.PIPE_WAIT,
  506. 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER,
  507. # default security descriptor: the handle cannot be inherited
  508. _winapi.NULL
  509. )
  510. h2 = _winapi.CreateFile(
  511. address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
  512. _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
  513. )
  514. _winapi.SetNamedPipeHandleState(
  515. h2, _winapi.PIPE_READMODE_MESSAGE, None, None
  516. )
  517. overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
  518. _, err = overlapped.GetOverlappedResult(True)
  519. assert err == 0
  520. c1 = PipeConnection(h1, writable=duplex)
  521. c2 = PipeConnection(h2, readable=duplex)
  522. return c1, c2
  523. #
  524. # Definitions for connections based on sockets
  525. #
  526. class SocketListener(object):
  527. '''
  528. Representation of a socket which is bound to an address and listening
  529. '''
  530. def __init__(self, address, family, backlog=1):
  531. self._socket = socket.socket(getattr(socket, family))
  532. try:
  533. # SO_REUSEADDR has different semantics on Windows (issue #2550).
  534. if os.name == 'posix':
  535. self._socket.setsockopt(socket.SOL_SOCKET,
  536. socket.SO_REUSEADDR, 1)
  537. self._socket.setblocking(True)
  538. self._socket.bind(address)
  539. self._socket.listen(backlog)
  540. self._address = self._socket.getsockname()
  541. except OSError:
  542. self._socket.close()
  543. raise
  544. self._family = family
  545. self._last_accepted = None
  546. if family == 'AF_UNIX':
  547. self._unlink = util.Finalize(
  548. self, os.unlink, args=(address,), exitpriority=0
  549. )
  550. else:
  551. self._unlink = None
  552. def accept(self):
  553. while True:
  554. try:
  555. s, self._last_accepted = self._socket.accept()
  556. except (OSError, IOError, socket.error) as exc:
  557. if getattr(exc, 'errno', None) != errno.EINTR:
  558. raise
  559. else:
  560. break
  561. s.setblocking(True)
  562. return Connection(detach(s))
  563. def close(self):
  564. try:
  565. self._socket.close()
  566. finally:
  567. unlink = self._unlink
  568. if unlink is not None:
  569. self._unlink = None
  570. unlink()
  571. def SocketClient(address):
  572. '''
  573. Return a connection object connected to the socket given by `address`
  574. '''
  575. family = address_type(address)
  576. s = socket.socket(getattr(socket, family))
  577. s.setblocking(True)
  578. s.connect(address)
  579. return Connection(detach(s))
  580. #
  581. # Definitions for connections based on named pipes
  582. #
  583. if sys.platform == 'win32':
  584. class PipeListener(object):
  585. '''
  586. Representation of a named pipe
  587. '''
  588. def __init__(self, address, backlog=None):
  589. self._address = address
  590. self._handle_queue = [self._new_handle(first=True)]
  591. self._last_accepted = None
  592. util.sub_debug('listener created with address=%r', self._address)
  593. self.close = util.Finalize(
  594. self, PipeListener._finalize_pipe_listener,
  595. args=(self._handle_queue, self._address), exitpriority=0
  596. )
  597. def _new_handle(self, first=False):
  598. flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
  599. if first:
  600. flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
  601. return _winapi.CreateNamedPipe(
  602. self._address, flags,
  603. _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
  604. _winapi.PIPE_WAIT,
  605. _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  606. _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
  607. )
  608. def accept(self):
  609. self._handle_queue.append(self._new_handle())
  610. handle = self._handle_queue.pop(0)
  611. try:
  612. ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
  613. except OSError as e:
  614. if e.winerror != _winapi.ERROR_NO_DATA:
  615. raise
  616. # ERROR_NO_DATA can occur if a client has already connected,
  617. # written data and then disconnected -- see Issue 14725.
  618. else:
  619. try:
  620. _winapi.WaitForMultipleObjects(
  621. [ov.event], False, INFINITE)
  622. except:
  623. ov.cancel()
  624. _winapi.CloseHandle(handle)
  625. raise
  626. finally:
  627. _, err = ov.GetOverlappedResult(True)
  628. assert err == 0
  629. return PipeConnection(handle)
  630. @staticmethod
  631. def _finalize_pipe_listener(queue, address):
  632. util.sub_debug('closing listener with address=%r', address)
  633. for handle in queue:
  634. _winapi.CloseHandle(handle)
  635. def PipeClient(address, _ignore=(_winapi.ERROR_SEM_TIMEOUT,
  636. _winapi.ERROR_PIPE_BUSY)):
  637. '''
  638. Return a connection object connected to the pipe given by `address`
  639. '''
  640. t = _init_timeout()
  641. while 1:
  642. try:
  643. _winapi.WaitNamedPipe(address, 1000)
  644. h = _winapi.CreateFile(
  645. address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
  646. 0, _winapi.NULL, _winapi.OPEN_EXISTING,
  647. _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
  648. )
  649. except OSError as e:
  650. if e.winerror not in _ignore or _check_timeout(t):
  651. raise
  652. else:
  653. break
  654. else:
  655. raise
  656. _winapi.SetNamedPipeHandleState(
  657. h, _winapi.PIPE_READMODE_MESSAGE, None, None
  658. )
  659. return PipeConnection(h)
  660. #
  661. # Authentication stuff
  662. #
  663. MESSAGE_LENGTH = 20
  664. CHALLENGE = b'#CHALLENGE#'
  665. WELCOME = b'#WELCOME#'
  666. FAILURE = b'#FAILURE#'
  667. def deliver_challenge(connection, authkey):
  668. import hmac
  669. assert isinstance(authkey, bytes)
  670. message = os.urandom(MESSAGE_LENGTH)
  671. connection.send_bytes(CHALLENGE + message)
  672. digest = hmac.new(authkey, message, 'md5').digest()
  673. response = connection.recv_bytes(256) # reject large message
  674. if response == digest:
  675. connection.send_bytes(WELCOME)
  676. else:
  677. connection.send_bytes(FAILURE)
  678. raise AuthenticationError('digest received was wrong')
  679. def answer_challenge(connection, authkey):
  680. import hmac
  681. assert isinstance(authkey, bytes)
  682. message = connection.recv_bytes(256) # reject large message
  683. assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
  684. message = message[len(CHALLENGE):]
  685. digest = hmac.new(authkey, message, 'md5').digest()
  686. connection.send_bytes(digest)
  687. response = connection.recv_bytes(256) # reject large message
  688. if response != WELCOME:
  689. raise AuthenticationError('digest sent was rejected')
  690. #
  691. # Support for using xmlrpclib for serialization
  692. #
  693. class ConnectionWrapper(object):
  694. def __init__(self, conn, dumps, loads):
  695. self._conn = conn
  696. self._dumps = dumps
  697. self._loads = loads
  698. for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
  699. obj = getattr(conn, attr)
  700. setattr(self, attr, obj)
  701. def send(self, obj):
  702. s = self._dumps(obj)
  703. self._conn.send_bytes(s)
  704. def recv(self):
  705. s = self._conn.recv_bytes()
  706. return self._loads(s)
  707. def _xml_dumps(obj):
  708. o = xmlrpclib.dumps((obj, ), None, None, None, 1) # noqa
  709. return o.encode('utf-8')
  710. def _xml_loads(s):
  711. (obj,), method = xmlrpclib.loads(s.decode('utf-8')) # noqa
  712. return obj
  713. class XmlListener(Listener):
  714. def accept(self):
  715. global xmlrpclib
  716. import xmlrpc.client as xmlrpclib # noqa
  717. obj = Listener.accept(self)
  718. return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
  719. def XmlClient(*args, **kwds):
  720. global xmlrpclib
  721. import xmlrpc.client as xmlrpclib # noqa
  722. return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
  723. #
  724. # Wait
  725. #
  726. if sys.platform == 'win32':
  727. def _exhaustive_wait(handles, timeout):
  728. # Return ALL handles which are currently signaled. (Only
  729. # returning the first signaled might create starvation issues.)
  730. L = list(handles)
  731. ready = []
  732. while L:
  733. res = _winapi.WaitForMultipleObjects(L, False, timeout)
  734. if res == WAIT_TIMEOUT:
  735. break
  736. elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
  737. res -= WAIT_OBJECT_0
  738. elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
  739. res -= WAIT_ABANDONED_0
  740. else:
  741. raise RuntimeError('Should not get here')
  742. ready.append(L[res])
  743. L = L[res + 1:]
  744. timeout = 0
  745. return ready
  746. _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
  747. def wait(object_list, timeout=None):
  748. '''
  749. Wait till an object in object_list is ready/readable.
  750. Returns list of those objects in object_list which are ready/readable.
  751. '''
  752. if timeout is None:
  753. timeout = INFINITE
  754. elif timeout < 0:
  755. timeout = 0
  756. else:
  757. timeout = int(timeout * 1000 + 0.5)
  758. object_list = list(object_list)
  759. waithandle_to_obj = {}
  760. ov_list = []
  761. ready_objects = set()
  762. ready_handles = set()
  763. try:
  764. for o in object_list:
  765. try:
  766. fileno = getattr(o, 'fileno')
  767. except AttributeError:
  768. waithandle_to_obj[o.__index__()] = o
  769. else:
  770. # start an overlapped read of length zero
  771. try:
  772. ov, err = _winapi.ReadFile(fileno(), 0, True)
  773. except OSError as e:
  774. err = e.winerror
  775. if err not in _ready_errors:
  776. raise
  777. if err == _winapi.ERROR_IO_PENDING:
  778. ov_list.append(ov)
  779. waithandle_to_obj[ov.event] = o
  780. else:
  781. # If o.fileno() is an overlapped pipe handle and
  782. # err == 0 then there is a zero length message
  783. # in the pipe, but it HAS NOT been consumed.
  784. ready_objects.add(o)
  785. timeout = 0
  786. ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
  787. finally:
  788. # request that overlapped reads stop
  789. for ov in ov_list:
  790. ov.cancel()
  791. # wait for all overlapped reads to stop
  792. for ov in ov_list:
  793. try:
  794. _, err = ov.GetOverlappedResult(True)
  795. except OSError as e:
  796. err = e.winerror
  797. if err not in _ready_errors:
  798. raise
  799. if err != _winapi.ERROR_OPERATION_ABORTED:
  800. o = waithandle_to_obj[ov.event]
  801. ready_objects.add(o)
  802. if err == 0:
  803. # If o.fileno() is an overlapped pipe handle then
  804. # a zero length message HAS been consumed.
  805. if hasattr(o, '_got_empty_message'):
  806. o._got_empty_message = True
  807. ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
  808. return [p for p in object_list if p in ready_objects]
  809. else:
  810. if hasattr(select, 'poll'):
  811. def _poll(fds, timeout):
  812. if timeout is not None:
  813. timeout = int(timeout * 1000) # timeout is in milliseconds
  814. fd_map = {}
  815. pollster = select.poll()
  816. for fd in fds:
  817. pollster.register(fd, select.POLLIN)
  818. if hasattr(fd, 'fileno'):
  819. fd_map[fd.fileno()] = fd
  820. else:
  821. fd_map[fd] = fd
  822. ls = []
  823. for fd, event in pollster.poll(timeout):
  824. if event & select.POLLNVAL:
  825. raise ValueError('invalid file descriptor %i' % fd)
  826. ls.append(fd_map[fd])
  827. return ls
  828. else:
  829. def _poll(fds, timeout): # noqa
  830. return select.select(fds, [], [], timeout)[0]
  831. def wait(object_list, timeout=None): # noqa
  832. '''
  833. Wait till an object in object_list is ready/readable.
  834. Returns list of those objects in object_list which are ready/readable.
  835. '''
  836. if timeout is not None:
  837. if timeout <= 0:
  838. return _poll(object_list, 0)
  839. else:
  840. deadline = monotonic() + timeout
  841. while True:
  842. try:
  843. return _poll(object_list, timeout)
  844. except (OSError, IOError, socket.error) as e:
  845. if e.errno != errno.EINTR:
  846. raise
  847. if timeout is not None:
  848. timeout = deadline - monotonic()
  849. #
  850. # Make connection and socket objects sharable if possible
  851. #
  852. if sys.platform == 'win32':
  853. def reduce_connection(conn):
  854. handle = conn.fileno()
  855. with socket.fromfd(handle, socket.AF_INET, socket.SOCK_STREAM) as s:
  856. from . import resource_sharer
  857. ds = resource_sharer.DupSocket(s)
  858. return rebuild_connection, (ds, conn.readable, conn.writable)
  859. def rebuild_connection(ds, readable, writable):
  860. sock = ds.detach()
  861. return Connection(detach(sock), readable, writable)
  862. reduction.register(Connection, reduce_connection)
  863. def reduce_pipe_connection(conn):
  864. access = ((_winapi.FILE_GENERIC_READ if conn.readable else 0) |
  865. (_winapi.FILE_GENERIC_WRITE if conn.writable else 0))
  866. dh = reduction.DupHandle(conn.fileno(), access)
  867. return rebuild_pipe_connection, (dh, conn.readable, conn.writable)
  868. def rebuild_pipe_connection(dh, readable, writable):
  869. return PipeConnection(detach(dh), readable, writable)
  870. reduction.register(PipeConnection, reduce_pipe_connection)
  871. else:
  872. def reduce_connection(conn):
  873. df = reduction.DupFd(conn.fileno())
  874. return rebuild_connection, (df, conn.readable, conn.writable)
  875. def rebuild_connection(df, readable, writable):
  876. return Connection(detach(df), readable, writable)
  877. reduction.register(Connection, reduce_connection)