/Lib/multiprocessing/connection.py

http://unladen-swallow.googlecode.com/ · Python · 417 lines · 284 code · 67 blank · 66 comment · 59 complexity · 9fe5f28ba266ef9a12094378d3726dc3 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.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  175. self._socket.bind(address)
  176. self._socket.listen(backlog)
  177. self._address = self._socket.getsockname()
  178. self._family = family
  179. self._last_accepted = None
  180. if family == 'AF_UNIX':
  181. self._unlink = Finalize(
  182. self, os.unlink, args=(address,), exitpriority=0
  183. )
  184. else:
  185. self._unlink = None
  186. def accept(self):
  187. s, self._last_accepted = self._socket.accept()
  188. fd = duplicate(s.fileno())
  189. conn = _multiprocessing.Connection(fd)
  190. s.close()
  191. return conn
  192. def close(self):
  193. self._socket.close()
  194. if self._unlink is not None:
  195. self._unlink()
  196. def SocketClient(address):
  197. '''
  198. Return a connection object connected to the socket given by `address`
  199. '''
  200. family = address_type(address)
  201. s = socket.socket( getattr(socket, family) )
  202. while 1:
  203. try:
  204. s.connect(address)
  205. except socket.error, e:
  206. if e.args[0] != errno.ECONNREFUSED: # connection refused
  207. debug('failed to connect to address %s', address)
  208. raise
  209. time.sleep(0.01)
  210. else:
  211. break
  212. else:
  213. raise
  214. fd = duplicate(s.fileno())
  215. conn = _multiprocessing.Connection(fd)
  216. s.close()
  217. return conn
  218. #
  219. # Definitions for connections based on named pipes
  220. #
  221. if sys.platform == 'win32':
  222. class PipeListener(object):
  223. '''
  224. Representation of a named pipe
  225. '''
  226. def __init__(self, address, backlog=None):
  227. self._address = address
  228. handle = win32.CreateNamedPipe(
  229. address, win32.PIPE_ACCESS_DUPLEX,
  230. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  231. win32.PIPE_WAIT,
  232. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  233. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  234. )
  235. self._handle_queue = [handle]
  236. self._last_accepted = None
  237. sub_debug('listener created with address=%r', self._address)
  238. self.close = Finalize(
  239. self, PipeListener._finalize_pipe_listener,
  240. args=(self._handle_queue, self._address), exitpriority=0
  241. )
  242. def accept(self):
  243. newhandle = win32.CreateNamedPipe(
  244. self._address, win32.PIPE_ACCESS_DUPLEX,
  245. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  246. win32.PIPE_WAIT,
  247. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  248. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  249. )
  250. self._handle_queue.append(newhandle)
  251. handle = self._handle_queue.pop(0)
  252. try:
  253. win32.ConnectNamedPipe(handle, win32.NULL)
  254. except WindowsError, e:
  255. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  256. raise
  257. return _multiprocessing.PipeConnection(handle)
  258. @staticmethod
  259. def _finalize_pipe_listener(queue, address):
  260. sub_debug('closing listener with address=%r', address)
  261. for handle in queue:
  262. close(handle)
  263. def PipeClient(address):
  264. '''
  265. Return a connection object connected to the pipe given by `address`
  266. '''
  267. while 1:
  268. try:
  269. win32.WaitNamedPipe(address, 1000)
  270. h = win32.CreateFile(
  271. address, win32.GENERIC_READ | win32.GENERIC_WRITE,
  272. 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  273. )
  274. except WindowsError, e:
  275. if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
  276. win32.ERROR_PIPE_BUSY):
  277. raise
  278. else:
  279. break
  280. else:
  281. raise
  282. win32.SetNamedPipeHandleState(
  283. h, win32.PIPE_READMODE_MESSAGE, None, None
  284. )
  285. return _multiprocessing.PipeConnection(h)
  286. #
  287. # Authentication stuff
  288. #
  289. MESSAGE_LENGTH = 20
  290. CHALLENGE = b'#CHALLENGE#'
  291. WELCOME = b'#WELCOME#'
  292. FAILURE = b'#FAILURE#'
  293. def deliver_challenge(connection, authkey):
  294. import hmac
  295. assert isinstance(authkey, bytes)
  296. message = os.urandom(MESSAGE_LENGTH)
  297. connection.send_bytes(CHALLENGE + message)
  298. digest = hmac.new(authkey, message).digest()
  299. response = connection.recv_bytes(256) # reject large message
  300. if response == digest:
  301. connection.send_bytes(WELCOME)
  302. else:
  303. connection.send_bytes(FAILURE)
  304. raise AuthenticationError('digest received was wrong')
  305. def answer_challenge(connection, authkey):
  306. import hmac
  307. assert isinstance(authkey, bytes)
  308. message = connection.recv_bytes(256) # reject large message
  309. assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
  310. message = message[len(CHALLENGE):]
  311. digest = hmac.new(authkey, message).digest()
  312. connection.send_bytes(digest)
  313. response = connection.recv_bytes(256) # reject large message
  314. if response != WELCOME:
  315. raise AuthenticationError('digest sent was rejected')
  316. #
  317. # Support for using xmlrpclib for serialization
  318. #
  319. class ConnectionWrapper(object):
  320. def __init__(self, conn, dumps, loads):
  321. self._conn = conn
  322. self._dumps = dumps
  323. self._loads = loads
  324. for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
  325. obj = getattr(conn, attr)
  326. setattr(self, attr, obj)
  327. def send(self, obj):
  328. s = self._dumps(obj)
  329. self._conn.send_bytes(s)
  330. def recv(self):
  331. s = self._conn.recv_bytes()
  332. return self._loads(s)
  333. def _xml_dumps(obj):
  334. return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf8')
  335. def _xml_loads(s):
  336. (obj,), method = xmlrpclib.loads(s.decode('utf8'))
  337. return obj
  338. class XmlListener(Listener):
  339. def accept(self):
  340. global xmlrpclib
  341. import xmlrpclib
  342. obj = Listener.accept(self)
  343. return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
  344. def XmlClient(*args, **kwds):
  345. global xmlrpclib
  346. import xmlrpclib
  347. return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)