connection.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436
  1. #
  2. # A higher level module for using sockets (or Windows named pipes)
  3. #
  4. # processing/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 time
  13. import tempfile
  14. import itertools
  15. from processing import _processing, currentProcess
  16. from processing.finalize import Finalize
  17. from processing.logger import subDebug
  18. from processing.reduction import connections_are_picklable
  19. try:
  20. from processing._processing import win32
  21. except ImportError:
  22. pass
  23. #
  24. #
  25. #
  26. BUFSIZE = 8192
  27. _nextId = itertools.count().next
  28. default_family = 'AF_INET'
  29. families = ['AF_INET']
  30. if hasattr(socket, 'AF_UNIX'):
  31. default_family = 'AF_UNIX'
  32. families += ['AF_UNIX']
  33. if sys.platform == 'win32':
  34. default_family = 'AF_PIPE'
  35. families += ['AF_PIPE']
  36. #
  37. #
  38. #
  39. def arbitraryAddress(family):
  40. '''
  41. Return an arbitrary free address for the given family
  42. '''
  43. if family == 'AF_INET':
  44. return ('localhost', 0)
  45. elif family == 'AF_UNIX':
  46. return tempfile.mktemp(prefix='pyc-%d-%d-' % (os.getpid(), _nextId()))
  47. elif family == 'AF_PIPE':
  48. return tempfile.mktemp(
  49. prefix=r'\\.\pipe\pyc-%d-%d-' % (os.getpid(), _nextId())
  50. )
  51. else:
  52. raise ValueError, 'unrecognized family'
  53. def addressType(address):
  54. '''
  55. Return the types of the address
  56. This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
  57. '''
  58. if type(address) == tuple:
  59. return 'AF_INET'
  60. elif type(address) is str and address.startswith('\\\\'):
  61. return 'AF_PIPE'
  62. elif type(address) is str:
  63. return 'AF_UNIX'
  64. else:
  65. raise ValueError, 'address type of %r unrecognized' % address
  66. #
  67. # Public functions
  68. #
  69. class Listener(object):
  70. '''
  71. Returns a listener object.
  72. This is a wrapper for a bound socket which is 'listening' for
  73. connections, or for a Windows named pipe.
  74. '''
  75. def __init__(self, address=None, family=None, backlog=1,
  76. authenticate=False, authkey=None):
  77. '''
  78. `address`
  79. The address to be used by the bound socket
  80. or named pipe of `self`.
  81. `family`
  82. The type of the socket or named pipe to use.
  83. This can be one of the strings 'AF_INET' (for a TCP
  84. socket), 'AF_UNIX' (for a Unix domain socket) or 'AF_PIPE'
  85. (for a Windows named pipe). Of these only the first is
  86. guaranteed to be available.
  87. If `family` is None than the family is inferred by the
  88. format of `address`. If `address` is unspecified then
  89. a default is chosen which is dependent on the platform.
  90. This default is the family which is assumed to be the
  91. fastest available.
  92. `backlog`
  93. If the `self` uses a socket then this is passed to the
  94. `listen()` method of the socket once it has been bound.
  95. `authenticate`
  96. If this is true then digest authentication is used even if
  97. `authkey` is` None`.
  98. `authkey`
  99. If `authkey` is a string then it will be used as the
  100. authentication key; otherwise it must be `None`.
  101. If `authkey` is `None` and `authenticate` is true then
  102. `currentProcess.getAuthKey()` is used as the authentication
  103. key.
  104. If `authkey` is `None` and `authentication` is false then
  105. no authentication is done.
  106. '''
  107. family = family or (address and addressType(address)) \
  108. or default_family
  109. address = address or arbitraryAddress(family)
  110. if family == 'AF_PIPE':
  111. self._listener = PipeListener(address, backlog)
  112. else:
  113. self._listener = SocketListener(address, family, backlog)
  114. if authenticate and authkey is None:
  115. authkey = currentProcess().getAuthKey()
  116. elif authenticate:
  117. assert type(authkey) is str
  118. self._authkey = authkey
  119. def accept(self):
  120. '''
  121. Accept a connection on the bound socket or named pipe of `self`.
  122. Returns a `Connection` object.
  123. '''
  124. c = self._listener.accept()
  125. if self._authkey:
  126. deliverChallenge(c, self._authkey)
  127. answerChallenge(c, self._authkey)
  128. return c
  129. def close(self):
  130. '''
  131. Close the bound socket or named pipe of `self`.
  132. '''
  133. return self._listener.close()
  134. address = property(lambda self: self._listener._address)
  135. last_accepted = property(lambda self: self._listener._last_accepted)
  136. def Client(address, family=None, authenticate=False, authkey=None):
  137. '''
  138. Returns a connection to the address of a `Listener`
  139. '''
  140. family = family or addressType(address)
  141. if family == 'AF_PIPE':
  142. c = PipeClient(address)
  143. else:
  144. c = SocketClient(address)
  145. if authenticate and authkey is None:
  146. authkey = currentProcess().getAuthKey()
  147. elif authenticate:
  148. assert type(authkey) is str
  149. if authkey is not None:
  150. answerChallenge(c, authkey)
  151. deliverChallenge(c, authkey)
  152. return c
  153. if sys.platform != 'win32':
  154. def Pipe(duplex=True):
  155. '''
  156. Returns pair of connection objects at either end of a pipe
  157. '''
  158. if duplex:
  159. s1, s2 = socket.socketpair()
  160. c1 = _processing.Connection(s1.fileno())
  161. c2 = _processing.Connection(s2.fileno())
  162. s1.close()
  163. s2.close()
  164. else:
  165. fd1, fd2 = os.pipe()
  166. c1 = _processing.Connection(fd1, duplicate=False)
  167. c2 = _processing.Connection(fd2, duplicate=False)
  168. return c1, c2
  169. else:
  170. def Pipe(duplex=True):
  171. '''
  172. Returns pair of connection objects at either end of a pipe
  173. '''
  174. address = arbitraryAddress('AF_PIPE')
  175. if duplex:
  176. openmode = win32.PIPE_ACCESS_DUPLEX
  177. access = win32.GENERIC_READ | win32.GENERIC_WRITE
  178. obsize, ibsize = BUFSIZE, BUFSIZE
  179. else:
  180. openmode = win32.PIPE_ACCESS_INBOUND
  181. access = win32.GENERIC_WRITE
  182. obsize, ibsize = 0, BUFSIZE
  183. h1 = win32.CreateNamedPipe(
  184. address, openmode,
  185. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  186. win32.PIPE_WAIT,
  187. 1, obsize, ibsize, win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  188. )
  189. h2 = win32.CreateFile(
  190. address, access, 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  191. )
  192. win32.SetNamedPipeHandleState(
  193. h2, win32.PIPE_READMODE_MESSAGE, None, None
  194. )
  195. try:
  196. win32.ConnectNamedPipe(h1, win32.NULL)
  197. except WindowsError, e:
  198. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  199. raise
  200. c1 = _processing.PipeConnection(h1, duplicate=False)
  201. c2 = _processing.PipeConnection(h2, duplicate=False)
  202. return c1, c2
  203. #
  204. # Definitions for connections based on sockets
  205. #
  206. class SocketListener(object):
  207. '''
  208. Represtation of a socket which is bound to an address and listening
  209. '''
  210. def __init__(self, address, family, backlog=1):
  211. self._socket = socket.socket(getattr(socket, family))
  212. self._socket.bind(address)
  213. if family == 'AF_UNIX':
  214. os.chmod(address, int('0600', 8)) # readable/writable only by user
  215. self._socket.listen(backlog)
  216. address = self._socket.getsockname()
  217. if type(address) is tuple:
  218. address = (socket.getfqdn(address[0]),) + address[1:]
  219. self._address = address
  220. self._family = family
  221. self._last_accepted = None
  222. subDebug('listener bound to address %r', self._address)
  223. if family == 'AF_UNIX':
  224. self._unlink = Finalize(
  225. self, os.unlink, args=(self._address,), exitpriority=0
  226. )
  227. else:
  228. self._unlink = None
  229. def accept(self):
  230. s, self._last_accepted = self._socket.accept()
  231. conn = _processing.Connection(s.fileno())
  232. s.close()
  233. return conn
  234. def close(self):
  235. self._socket.close()
  236. if self._unlink is not None:
  237. self._unlink()
  238. def SocketClient(address):
  239. '''
  240. Return a connection object connected to the socket given by `address`
  241. '''
  242. family = addressType(address)
  243. s = socket.socket( getattr(socket, family) )
  244. endtime = time.time() + 10
  245. while endtime > time.time():
  246. try:
  247. s.connect(address)
  248. except socket.error, e:
  249. if e.args[0] != 10061: # 10061 => connection refused
  250. raise
  251. time.sleep(0.01)
  252. else:
  253. break
  254. else:
  255. raise
  256. conn = _processing.Connection(s.fileno())
  257. s.close()
  258. return conn
  259. #
  260. # Definitions for connections based on named pipes
  261. #
  262. if sys.platform == 'win32':
  263. class PipeListener(object):
  264. '''
  265. Representation of a named pipe
  266. '''
  267. def __init__(self, address, backlog=None):
  268. self._address = address
  269. handle = win32.CreateNamedPipe(
  270. address, win32.PIPE_ACCESS_DUPLEX,
  271. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  272. win32.PIPE_WAIT,
  273. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  274. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  275. )
  276. self._handle_queue = [handle]
  277. self._last_accepted = None
  278. subDebug('listener created with address=%r', self._address)
  279. self.close = Finalize(
  280. self, PipeListener._finalize_pipelistener,
  281. args=(self._handle_queue, self._address), exitpriority=0
  282. )
  283. def accept(self):
  284. newhandle = win32.CreateNamedPipe(
  285. self._address, win32.PIPE_ACCESS_DUPLEX,
  286. win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE |
  287. win32.PIPE_WAIT,
  288. win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
  289. win32.NMPWAIT_WAIT_FOREVER, win32.NULL
  290. )
  291. self._handle_queue.append(newhandle)
  292. handle = self._handle_queue.pop(0)
  293. try:
  294. win32.ConnectNamedPipe(handle, win32.NULL)
  295. except WindowsError, e:
  296. if e.args[0] != win32.ERROR_PIPE_CONNECTED:
  297. raise
  298. return _processing.PipeConnection(handle, duplicate=False)
  299. @staticmethod
  300. def _finalize_pipelistener(queue, address):
  301. subDebug('closing listener with address=%r', address)
  302. for handle in queue:
  303. win32.CloseHandle(handle)
  304. def PipeClient(address):
  305. '''
  306. Return a connection object connected to the pipe given by `address`
  307. '''
  308. endtime = time.time() + 10
  309. while endtime > time.time():
  310. try:
  311. win32.WaitNamedPipe(address, 1000)
  312. h = win32.CreateFile(
  313. address, win32.GENERIC_READ | win32.GENERIC_WRITE,
  314. 0, win32.NULL, win32.OPEN_EXISTING, 0, win32.NULL
  315. )
  316. except WindowsError, e:
  317. if e.args[0] not in (win32.ERROR_SEM_TIMEOUT,
  318. win32.ERROR_PIPE_BUSY):
  319. raise
  320. else:
  321. break
  322. else:
  323. raise
  324. win32.SetNamedPipeHandleState(
  325. h, win32.PIPE_READMODE_MESSAGE, None, None
  326. )
  327. return _processing.PipeConnection(h, duplicate=False)
  328. #
  329. # Authentication stuff
  330. #
  331. class AuthenticationError(Exception):
  332. pass
  333. def deliverChallenge(connection, authkey):
  334. import hmac, sha
  335. assert type(authkey) is str, '%r is not a string' % authkey
  336. try:
  337. message = os.urandom(20)
  338. except AttributeError:
  339. import random
  340. message = ''.join(str(random.randrange(256)) for i in range(20))
  341. connection.sendBytes('#CHALLENGE:' + message)
  342. digest = hmac.new(authkey, message, sha).digest()
  343. response = connection.recvBytes()
  344. if response == digest:
  345. connection.sendBytes('#WELCOME')
  346. else:
  347. connection.sendBytes('#AUTHENTICATION_FAILED')
  348. raise AuthenticationError, 'digest received was wrong'
  349. def answerChallenge(connection, authkey):
  350. import hmac, sha
  351. assert type(authkey) is str, '%r is not a string' % authkey
  352. message = connection.recvBytes()
  353. assert message[:11] == '#CHALLENGE:', 'message = %r' % message
  354. message = message[11:]
  355. digest = hmac.new(authkey, message, sha).digest()
  356. connection.sendBytes(digest)
  357. response = connection.recvBytes()
  358. if response != '#WELCOME':
  359. raise AuthenticationError, 'digest sent was rejected'
  360. # deprecated
  361. deliver_Challenge = deliverChallenge
  362. answer_challenge = answerChallenge