PageRenderTime 54ms CodeModel.GetById 20ms RepoModel.GetById 1ms app.codeStats 0ms

/python/Lib/multiprocessing/connection.py

https://github.com/vesnican/play
Python | 416 lines | 364 code | 24 blank | 28 comment | 14 complexity | 59ef3d301715159b5a72bb7a7462ee92 MD5 | raw file
  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 --- see COPYING.txt
  7. #
  8. __all__ = [ 'Client', 'Listener', 'Pipe' ]
  9. import os
  10. import sys
  11. import socket
  12. import errno
  13. import time
  14. import tempfile
  15. import itertools
  16. import _multiprocessing
  17. from multiprocessing import current_process, AuthenticationError
  18. from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
  19. from multiprocessing.forking import duplicate, close
  20. #
  21. #
  22. #
  23. BUFSIZE = 8192
  24. _mmap_counter = itertools.count()
  25. default_family = 'AF_INET'
  26. families = ['AF_INET']
  27. if hasattr(socket, 'AF_UNIX'):
  28. default_family = 'AF_UNIX'
  29. families += ['AF_UNIX']
  30. if sys.platform == 'win32':
  31. default_family = 'AF_PIPE'
  32. families += ['AF_PIPE']
  33. #
  34. #
  35. #
  36. def arbitrary_address(family):
  37. '''
  38. Return an arbitrary free address for the given family
  39. '''
  40. if family == 'AF_INET':
  41. return ('localhost', 0)
  42. elif family == 'AF_UNIX':
  43. return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
  44. elif family == 'AF_PIPE':
  45. return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
  46. (os.getpid(), _mmap_counter.next()))
  47. else:
  48. raise ValueError('unrecognized family')
  49. def address_type(address):
  50. '''
  51. Return the types of the address
  52. This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
  53. '''
  54. if type(address) == tuple:
  55. return 'AF_INET'
  56. elif type(address) is str and address.startswith('\\\\'):
  57. return 'AF_PIPE'
  58. elif type(address) is str:
  59. return 'AF_UNIX'
  60. else:
  61. raise ValueError('address type of %r unrecognized' % address)
  62. #
  63. # Public functions
  64. #
  65. class Listener(object):
  66. '''
  67. Returns a listener object.
  68. This is a wrapper for a bound socket which is 'listening' for
  69. connections, or for a Windows named pipe.
  70. '''
  71. def __init__(self, address=None, family=None, backlog=1, authkey=None):
  72. family = family or (address and address_type(address)) \
  73. or default_family
  74. address = address or arbitrary_address(family)
  75. if family == 'AF_PIPE':
  76. self._listener = PipeListener(address, backlog)
  77. else:
  78. self._listener = SocketListener(address, family, backlog)
  79. if authkey is not None and not isinstance(authkey, bytes):
  80. raise TypeError, 'authkey should be a byte string'
  81. self._authkey = authkey
  82. def accept(self):
  83. '''
  84. Accept a connection on the bound socket or named pipe of `self`.
  85. Returns a `Connection` object.
  86. '''
  87. c = self._listener.accept()
  88. if self._authkey:
  89. deliver_challenge(c, self._authkey)
  90. answer_challenge(c, self._authkey)
  91. return c
  92. def close(self):
  93. '''
  94. Close the bound socket or named pipe of `self`.
  95. '''
  96. return self._listener.close()
  97. address = property(lambda self: self._listener._address)
  98. last_accepted = property(lambda self: self._listener._last_accepted)
  99. def Client(address, family=None, authkey=None):
  100. '''
  101. Returns a connection to the address of a `Listener`
  102. '''
  103. family = family or address_type(address)
  104. if family == 'AF_PIPE':
  105. c = PipeClient(address)
  106. else:
  107. c = SocketClient(address)
  108. if authkey is not None and not isinstance(authkey, bytes):
  109. raise TypeError, 'authkey should be a byte string'
  110. if authkey is not None:
  111. answer_challenge(c, authkey)
  112. deliver_challenge(c, authkey)
  113. return c
  114. if sys.platform != 'win32':
  115. def Pipe(duplex=True):
  116. '''
  117. Returns pair of connection objects at either end of a pipe
  118. '''
  119. if duplex:
  120. s1, s2 = socket.socketpair()
  121. c1 = _multiprocessing.Connection(os.dup(s1.fileno()))
  122. c2 = _multiprocessing.Connection(os.dup(s2.fileno()))
  123. s1.close()
  124. s2.close()
  125. else:
  126. fd1, fd2 = os.pipe()
  127. c1 = _multiprocessing.Connection(fd1, writable=False)
  128. c2 = _multiprocessing.Connection(fd2, readable=False)
  129. return c1, c2
  130. else:
  131. from ._multiprocessing import win32
  132. def Pipe(duplex=True):
  133. '''
  134. Returns pair of connection objects at either end of a pipe
  135. '''
  136. address = arbitrary_address('AF_PIPE')
  137. if duplex:
  138. openmode = win32.PIPE_ACCESS_DUPLEX
  139. access = win32.GENERIC_READ | win32.GENERIC_WRITE
  140. obsize, ibsize = BUFSIZE, BUFSIZE
  141. else:
  142. openmode = win32.PIPE_ACCESS_INBOUND
  143. access = win32.GENERIC_WRITE
  144. obsize, ibsize = 0, BUFSIZE
  145. h1 = win32.CreateNamedPipe(
  146. address, openmode,
  147. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  148. win32.PIPE_WAIT,
  149. 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  150. )
  151. h2 = win32.CreateFile(
  152. address, access, 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  153. )
  154. win32.SetNamedPipeHandleState(
  155. h2, win32.PIPE_READMODE_MESSAGE, None, None
  156. )
  157. try:
  158. win32.ConnectNamedPipe(h1, win32.NULL)
  159. except WindowsError, e:
  160. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  161. raise
  162. c1 = _multiprocessing.PipeConnection(h1, writable=duplex)
  163. c2 = _multiprocessing.PipeConnection(h2, readable=duplex)
  164. return c1, c2
  165. #
  166. # Definitions for connections based on sockets
  167. #
  168. class SocketListener(object):
  169. '''
  170. Representation of a socket which is bound to an address and listening
  171. '''
  172. def __init__(self, address, family, backlog=1):
  173. self._socket = socket.socket(getattr(socket, family))
  174. self._socket.bind(address)
  175. self._socket.listen(backlog)
  176. self._address = self._socket.getsockname()
  177. self._family = family
  178. self._last_accepted = None
  179. if family == 'AF_UNIX':
  180. self._unlink = Finalize(
  181. self, os.unlink, args=(address,), exitpriority=0
  182. )
  183. else:
  184. self._unlink = None
  185. def accept(self):
  186. s, self._last_accepted = self._socket.accept()
  187. fd = duplicate(s.fileno())
  188. conn = _multiprocessing.Connection(fd)
  189. s.close()
  190. return conn
  191. def close(self):
  192. self._socket.close()
  193. if self._unlink is not None:
  194. self._unlink()
  195. def SocketClient(address):
  196. '''
  197. Return a connection object connected to the socket given by `address`
  198. '''
  199. family = address_type(address)
  200. s = socket.socket( getattr(socket, family) )
  201. while 1:
  202. try:
  203. s.connect(address)
  204. except socket.error, e:
  205. if e.args[0] != errno.ECONNREFUSED: # connection refused
  206. debug('failed to connect to address %s', address)
  207. raise
  208. time.sleep(0.01)
  209. else:
  210. break
  211. else:
  212. raise
  213. fd = duplicate(s.fileno())
  214. conn = _multiprocessing.Connection(fd)
  215. s.close()
  216. return conn
  217. #
  218. # Definitions for connections based on named pipes
  219. #
  220. if sys.platform == 'win32':
  221. class PipeListener(object):
  222. '''
  223. Representation of a named pipe
  224. '''
  225. def __init__(self, address, backlog=None):
  226. self._address = address
  227. handle = win32.CreateNamedPipe(
  228. address, win32.PIPE_ACCESS_DUPLEX,
  229. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  230. win32.PIPE_WAIT,
  231. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  232. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  233. )
  234. self._handle_queue = [handle]
  235. self._last_accepted = None
  236. sub_debug('listener created with address=%r', self._address)
  237. self.close = Finalize(
  238. self, PipeListener._finalize_pipe_listener,
  239. args=(self._handle_queue, self._address), exitpriority=0
  240. )
  241. def accept(self):
  242. newhandle = win32.CreateNamedPipe(
  243. self._address, win32.PIPE_ACCESS_DUPLEX,
  244. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  245. win32.PIPE_WAIT,
  246. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  247. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  248. )
  249. self._handle_queue.append(newhandle)
  250. handle = self._handle_queue.pop(0)
  251. try:
  252. win32.ConnectNamedPipe(handle, win32.NULL)
  253. except WindowsError, e:
  254. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  255. raise
  256. return _multiprocessing.PipeConnection(handle)
  257. @staticmethod
  258. def _finalize_pipe_listener(queue, address):
  259. sub_debug('closing listener with address=%r', address)
  260. for handle in queue:
  261. close(handle)
  262. def PipeClient(address):
  263. '''
  264. Return a connection object connected to the pipe given by `address`
  265. '''
  266. while 1:
  267. try:
  268. win32.WaitNamedPipe(address, 1000)
  269. h = win32.CreateFile(
  270. address, win32.GENERIC_READ | win32.GENERIC_WRITE,
  271. 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  272. )
  273. except WindowsError, e:
  274. if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
  275. win32.ERROR_PIPE_BUSY):
  276. raise
  277. else:
  278. break
  279. else:
  280. raise
  281. win32.SetNamedPipeHandleState(
  282. h, win32.PIPE_READMODE_MESSAGE, None, None
  283. )
  284. return _multiprocessing.PipeConnection(h)
  285. #
  286. # Authentication stuff
  287. #
  288. MESSAGE_LENGTH = 20
  289. CHALLENGE = b'#CHALLENGE#'
  290. WELCOME = b'#WELCOME#'
  291. FAILURE = b'#FAILURE#'
  292. def deliver_challenge(connection, authkey):
  293. import hmac
  294. assert isinstance(authkey, bytes)
  295. message = os.urandom(MESSAGE_LENGTH)
  296. connection.send_bytes(CHALLENGE + message)
  297. digest = hmac.new(authkey, message).digest()
  298. response = connection.recv_bytes(256) # reject large message
  299. if response == digest:
  300. connection.send_bytes(WELCOME)
  301. else:
  302. connection.send_bytes(FAILURE)
  303. raise AuthenticationError('digest received was wrong')
  304. def answer_challenge(connection, authkey):
  305. import hmac
  306. assert isinstance(authkey, bytes)
  307. message = connection.recv_bytes(256) # reject large message
  308. assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
  309. message = message[len(CHALLENGE):]
  310. digest = hmac.new(authkey, message).digest()
  311. connection.send_bytes(digest)
  312. response = connection.recv_bytes(256) # reject large message
  313. if response != WELCOME:
  314. raise AuthenticationError('digest sent was rejected')
  315. #
  316. # Support for using xmlrpclib for serialization
  317. #
  318. class ConnectionWrapper(object):
  319. def __init__(self, conn, dumps, loads):
  320. self._conn = conn
  321. self._dumps = dumps
  322. self._loads = loads
  323. for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
  324. obj = getattr(conn, attr)
  325. setattr(self, attr, obj)
  326. def send(self, obj):
  327. s = self._dumps(obj)
  328. self._conn.send_bytes(s)
  329. def recv(self):
  330. s = self._conn.recv_bytes()
  331. return self._loads(s)
  332. def _xml_dumps(obj):
  333. return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf8')
  334. def _xml_loads(s):
  335. (obj,), method = xmlrpclib.loads(s.decode('utf8'))
  336. return obj
  337. class XmlListener(Listener):
  338. def accept(self):
  339. global xmlrpclib
  340. import xmlrpclib
  341. obj = Listener.accept(self)
  342. return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
  343. def XmlClient(*args, **kwds):
  344. global xmlrpclib
  345. import xmlrpclib
  346. return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)