connection.py 31 KB

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