Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

connection.py 30KB

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