connection.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477
  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. # All rights reserved.
  8. #
  9. # Redistribution and use in source and binary forms, with or without
  10. # modification, are permitted provided that the following conditions
  11. # are met:
  12. #
  13. # 1. Redistributions of source code must retain the above copyright
  14. # notice, this list of conditions and the following disclaimer.
  15. # 2. Redistributions in binary form must reproduce the above copyright
  16. # notice, this list of conditions and the following disclaimer in the
  17. # documentation and/or other materials provided with the distribution.
  18. # 3. Neither the name of author nor the names of any contributors may be
  19. # used to endorse or promote products derived from this software
  20. # without specific prior written permission.
  21. #
  22. # THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
  23. # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  24. # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  25. # ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  26. # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  27. # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  28. # OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  29. # HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  30. # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  31. # OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  32. # SUCH DAMAGE.
  33. #
  34. __all__ = [ 'Client', 'Listener', 'Pipe' ]
  35. import os
  36. import sys
  37. import socket
  38. import errno
  39. import time
  40. import tempfile
  41. import itertools
  42. import _multiprocessing
  43. from multiprocessing import current_process, AuthenticationError
  44. from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
  45. from multiprocessing.forking import duplicate, close
  46. #
  47. #
  48. #
  49. BUFSIZE = 8192
  50. # A very generous timeout when it comes to local connections...
  51. CONNECTION_TIMEOUT = 20.
  52. _mmap_counter = itertools.count()
  53. default_family = 'AF_INET'
  54. families = ['AF_INET']
  55. if hasattr(socket, 'AF_UNIX'):
  56. default_family = 'AF_UNIX'
  57. families += ['AF_UNIX']
  58. if sys.platform == 'win32':
  59. default_family = 'AF_PIPE'
  60. families += ['AF_PIPE']
  61. def _init_timeout(timeout=CONNECTION_TIMEOUT):
  62. return time.time() + timeout
  63. def _check_timeout(t):
  64. return time.time() > t
  65. #
  66. #
  67. #
  68. def arbitrary_address(family):
  69. '''
  70. Return an arbitrary free address for the given family
  71. '''
  72. if family == 'AF_INET':
  73. return ('localhost', 0)
  74. elif family == 'AF_UNIX':
  75. return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
  76. elif family == 'AF_PIPE':
  77. return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
  78. (os.getpid(), _mmap_counter.next()), dir="")
  79. else:
  80. raise ValueError('unrecognized family')
  81. def address_type(address):
  82. '''
  83. Return the types of the address
  84. This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
  85. '''
  86. if type(address) == tuple:
  87. return 'AF_INET'
  88. elif type(address) is str and address.startswith('\\\\'):
  89. return 'AF_PIPE'
  90. elif type(address) is str:
  91. return 'AF_UNIX'
  92. else:
  93. raise ValueError('address type of %r unrecognized' % address)
  94. #
  95. # Public functions
  96. #
  97. class Listener(object):
  98. '''
  99. Returns a listener object.
  100. This is a wrapper for a bound socket which is 'listening' for
  101. connections, or for a Windows named pipe.
  102. '''
  103. def __init__(self, address=None, family=None, backlog=1, authkey=None):
  104. family = family or (address and address_type(address)) \
  105. or default_family
  106. address = address or arbitrary_address(family)
  107. if family == 'AF_PIPE':
  108. self._listener = PipeListener(address, backlog)
  109. else:
  110. self._listener = SocketListener(address, family, backlog)
  111. if authkey is not None and not isinstance(authkey, bytes):
  112. raise TypeError, 'authkey should be a byte string'
  113. self._authkey = authkey
  114. def accept(self):
  115. '''
  116. Accept a connection on the bound socket or named pipe of `self`.
  117. Returns a `Connection` object.
  118. '''
  119. c = self._listener.accept()
  120. if self._authkey:
  121. deliver_challenge(c, self._authkey)
  122. answer_challenge(c, self._authkey)
  123. return c
  124. def close(self):
  125. '''
  126. Close the bound socket or named pipe of `self`.
  127. '''
  128. return self._listener.close()
  129. address = property(lambda self: self._listener._address)
  130. last_accepted = property(lambda self: self._listener._last_accepted)
  131. def Client(address, family=None, authkey=None):
  132. '''
  133. Returns a connection to the address of a `Listener`
  134. '''
  135. family = family or address_type(address)
  136. if family == 'AF_PIPE':
  137. c = PipeClient(address)
  138. else:
  139. c = SocketClient(address)
  140. if authkey is not None and not isinstance(authkey, bytes):
  141. raise TypeError, 'authkey should be a byte string'
  142. if authkey is not None:
  143. answer_challenge(c, authkey)
  144. deliver_challenge(c, authkey)
  145. return c
  146. if sys.platform != 'win32':
  147. def Pipe(duplex=True):
  148. '''
  149. Returns pair of connection objects at either end of a pipe
  150. '''
  151. if duplex:
  152. s1, s2 = socket.socketpair()
  153. s1.setblocking(True)
  154. s2.setblocking(True)
  155. c1 = _multiprocessing.Connection(os.dup(s1.fileno()))
  156. c2 = _multiprocessing.Connection(os.dup(s2.fileno()))
  157. s1.close()
  158. s2.close()
  159. else:
  160. fd1, fd2 = os.pipe()
  161. c1 = _multiprocessing.Connection(fd1, writable=False)
  162. c2 = _multiprocessing.Connection(fd2, readable=False)
  163. return c1, c2
  164. else:
  165. from _multiprocessing import win32
  166. def Pipe(duplex=True):
  167. '''
  168. Returns pair of connection objects at either end of a pipe
  169. '''
  170. address = arbitrary_address('AF_PIPE')
  171. if duplex:
  172. openmode = win32.PIPE_ACCESS_DUPLEX
  173. access = win32.GENERIC_READ | win32.GENERIC_WRITE
  174. obsize, ibsize = BUFSIZE, BUFSIZE
  175. else:
  176. openmode = win32.PIPE_ACCESS_INBOUND
  177. access = win32.GENERIC_WRITE
  178. obsize, ibsize = 0, BUFSIZE
  179. h1 = win32.CreateNamedPipe(
  180. address, openmode,
  181. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  182. win32.PIPE_WAIT,
  183. 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  184. )
  185. h2 = win32.CreateFile(
  186. address, access, 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  187. )
  188. win32.SetNamedPipeHandleState(
  189. h2, win32.PIPE_READMODE_MESSAGE, None, None
  190. )
  191. try:
  192. win32.ConnectNamedPipe(h1, win32.NULL)
  193. except WindowsError, e:
  194. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  195. raise
  196. c1 = _multiprocessing.PipeConnection(h1, writable=duplex)
  197. c2 = _multiprocessing.PipeConnection(h2, readable=duplex)
  198. return c1, c2
  199. #
  200. # Definitions for connections based on sockets
  201. #
  202. class SocketListener(object):
  203. '''
  204. Representation of a socket which is bound to an address and listening
  205. '''
  206. def __init__(self, address, family, backlog=1):
  207. self._socket = socket.socket(getattr(socket, family))
  208. try:
  209. self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  210. self._socket.setblocking(True)
  211. self._socket.bind(address)
  212. self._socket.listen(backlog)
  213. self._address = self._socket.getsockname()
  214. except socket.error:
  215. self._socket.close()
  216. raise
  217. self._family = family
  218. self._last_accepted = None
  219. if family == 'AF_UNIX':
  220. self._unlink = Finalize(
  221. self, os.unlink, args=(address,), exitpriority=0
  222. )
  223. else:
  224. self._unlink = None
  225. def accept(self):
  226. while True:
  227. try:
  228. s, self._last_accepted = self._socket.accept()
  229. except socket.error as e:
  230. if e.args[0] != errno.EINTR:
  231. raise
  232. else:
  233. break
  234. s.setblocking(True)
  235. fd = duplicate(s.fileno())
  236. conn = _multiprocessing.Connection(fd)
  237. s.close()
  238. return conn
  239. def close(self):
  240. try:
  241. self._socket.close()
  242. finally:
  243. unlink = self._unlink
  244. if unlink is not None:
  245. self._unlink = None
  246. unlink()
  247. def SocketClient(address):
  248. '''
  249. Return a connection object connected to the socket given by `address`
  250. '''
  251. family = getattr(socket, address_type(address))
  252. t = _init_timeout()
  253. while 1:
  254. s = socket.socket(family)
  255. s.setblocking(True)
  256. try:
  257. s.connect(address)
  258. except socket.error, e:
  259. s.close()
  260. if e.args[0] != errno.ECONNREFUSED or _check_timeout(t):
  261. debug('failed to connect to address %s', address)
  262. raise
  263. time.sleep(0.01)
  264. else:
  265. break
  266. else:
  267. raise
  268. fd = duplicate(s.fileno())
  269. conn = _multiprocessing.Connection(fd)
  270. s.close()
  271. return conn
  272. #
  273. # Definitions for connections based on named pipes
  274. #
  275. if sys.platform == 'win32':
  276. class PipeListener(object):
  277. '''
  278. Representation of a named pipe
  279. '''
  280. def __init__(self, address, backlog=None):
  281. self._address = address
  282. handle = win32.CreateNamedPipe(
  283. address, win32.PIPE_ACCESS_DUPLEX,
  284. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  285. win32.PIPE_WAIT,
  286. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  287. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  288. )
  289. self._handle_queue = [handle]
  290. self._last_accepted = None
  291. sub_debug('listener created with address=%r', self._address)
  292. self.close = Finalize(
  293. self, PipeListener._finalize_pipe_listener,
  294. args=(self._handle_queue, self._address), exitpriority=0
  295. )
  296. def accept(self):
  297. newhandle = win32.CreateNamedPipe(
  298. self._address, win32.PIPE_ACCESS_DUPLEX,
  299. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  300. win32.PIPE_WAIT,
  301. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  302. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  303. )
  304. self._handle_queue.append(newhandle)
  305. handle = self._handle_queue.pop(0)
  306. try:
  307. win32.ConnectNamedPipe(handle, win32.NULL)
  308. except WindowsError, e:
  309. # ERROR_NO_DATA can occur if a client has already connected,
  310. # written data and then disconnected -- see Issue 14725.
  311. if e.args[0] not in (win32.ERROR_PIPE_CONNECTED,
  312. win32.ERROR_NO_DATA):
  313. raise
  314. return _multiprocessing.PipeConnection(handle)
  315. @staticmethod
  316. def _finalize_pipe_listener(queue, address):
  317. sub_debug('closing listener with address=%r', address)
  318. for handle in queue:
  319. close(handle)
  320. def PipeClient(address):
  321. '''
  322. Return a connection object connected to the pipe given by `address`
  323. '''
  324. t = _init_timeout()
  325. while 1:
  326. try:
  327. win32.WaitNamedPipe(address, 1000)
  328. h = win32.CreateFile(
  329. address, win32.GENERIC_READ | win32.GENERIC_WRITE,
  330. 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  331. )
  332. except WindowsError, e:
  333. if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
  334. win32.ERROR_PIPE_BUSY) or _check_timeout(t):
  335. raise
  336. else:
  337. break
  338. else:
  339. raise
  340. win32.SetNamedPipeHandleState(
  341. h, win32.PIPE_READMODE_MESSAGE, None, None
  342. )
  343. return _multiprocessing.PipeConnection(h)
  344. #
  345. # Authentication stuff
  346. #
  347. MESSAGE_LENGTH = 20
  348. CHALLENGE = b'#CHALLENGE#'
  349. WELCOME = b'#WELCOME#'
  350. FAILURE = b'#FAILURE#'
  351. def deliver_challenge(connection, authkey):
  352. import hmac
  353. assert isinstance(authkey, bytes)
  354. message = os.urandom(MESSAGE_LENGTH)
  355. connection.send_bytes(CHALLENGE + message)
  356. digest = hmac.new(authkey, message).digest()
  357. response = connection.recv_bytes(256) # reject large message
  358. if response == digest:
  359. connection.send_bytes(WELCOME)
  360. else:
  361. connection.send_bytes(FAILURE)
  362. raise AuthenticationError('digest received was wrong')
  363. def answer_challenge(connection, authkey):
  364. import hmac
  365. assert isinstance(authkey, bytes)
  366. message = connection.recv_bytes(256) # reject large message
  367. assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
  368. message = message[len(CHALLENGE):]
  369. digest = hmac.new(authkey, message).digest()
  370. connection.send_bytes(digest)
  371. response = connection.recv_bytes(256) # reject large message
  372. if response != WELCOME:
  373. raise AuthenticationError('digest sent was rejected')
  374. #
  375. # Support for using xmlrpclib for serialization
  376. #
  377. class ConnectionWrapper(object):
  378. def __init__(self, conn, dumps, loads):
  379. self._conn = conn
  380. self._dumps = dumps
  381. self._loads = loads
  382. for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
  383. obj = getattr(conn, attr)
  384. setattr(self, attr, obj)
  385. def send(self, obj):
  386. s = self._dumps(obj)
  387. self._conn.send_bytes(s)
  388. def recv(self):
  389. s = self._conn.recv_bytes()
  390. return self._loads(s)
  391. def _xml_dumps(obj):
  392. return xmlrpclib.dumps((obj,), None, None, None, 1)
  393. def _xml_loads(s):
  394. (obj,), method = xmlrpclib.loads(s)
  395. return obj
  396. class XmlListener(Listener):
  397. def accept(self):
  398. global xmlrpclib
  399. import xmlrpclib
  400. obj = Listener.accept(self)
  401. return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
  402. def XmlClient(*args, **kwds):
  403. global xmlrpclib
  404. import xmlrpclib
  405. return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)