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.

transport.py 9.8KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  1. # Copyright (C) 2009 Barry Pederson <bp@barryp.org>
  2. #
  3. # This library is free software; you can redistribute it and/or
  4. # modify it under the terms of the GNU Lesser General Public
  5. # License as published by the Free Software Foundation; either
  6. # version 2.1 of the License, or (at your option) any later version.
  7. #
  8. # This library is distributed in the hope that it will be useful,
  9. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  10. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  11. # Lesser General Public License for more details.
  12. #
  13. # You should have received a copy of the GNU Lesser General Public
  14. # License along with this library; if not, write to the Free Software
  15. # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301
  16. from __future__ import absolute_import
  17. import errno
  18. import re
  19. import socket
  20. import ssl
  21. # Jython does not have this attribute
  22. try:
  23. from socket import SOL_TCP
  24. except ImportError: # pragma: no cover
  25. from socket import IPPROTO_TCP as SOL_TCP # noqa
  26. try:
  27. from ssl import SSLError
  28. except ImportError:
  29. class SSLError(Exception): # noqa
  30. pass
  31. from struct import pack, unpack
  32. from .exceptions import UnexpectedFrame
  33. from .utils import get_errno, set_cloexec
  34. _UNAVAIL = errno.EAGAIN, errno.EINTR, errno.ENOENT
  35. AMQP_PORT = 5672
  36. EMPTY_BUFFER = bytes()
  37. # Yes, Advanced Message Queuing Protocol Protocol is redundant
  38. AMQP_PROTOCOL_HEADER = 'AMQP\x01\x01\x00\x09'.encode('latin_1')
  39. # Match things like: [fe80::1]:5432, from RFC 2732
  40. IPV6_LITERAL = re.compile(r'\[([\.0-9a-f:]+)\](?::(\d+))?')
  41. class _AbstractTransport(object):
  42. """Common superclass for TCP and SSL transports"""
  43. connected = False
  44. def __init__(self, host, connect_timeout):
  45. self.connected = True
  46. msg = None
  47. port = AMQP_PORT
  48. m = IPV6_LITERAL.match(host)
  49. if m:
  50. host = m.group(1)
  51. if m.group(2):
  52. port = int(m.group(2))
  53. else:
  54. if ':' in host:
  55. host, port = host.rsplit(':', 1)
  56. port = int(port)
  57. self.sock = None
  58. last_err = None
  59. for res in socket.getaddrinfo(host, port, 0,
  60. socket.SOCK_STREAM, SOL_TCP):
  61. af, socktype, proto, canonname, sa = res
  62. try:
  63. self.sock = socket.socket(af, socktype, proto)
  64. try:
  65. set_cloexec(self.sock, True)
  66. except NotImplementedError:
  67. pass
  68. self.sock.settimeout(connect_timeout)
  69. self.sock.connect(sa)
  70. except socket.error as exc:
  71. msg = exc
  72. self.sock.close()
  73. self.sock = None
  74. last_err = msg
  75. continue
  76. break
  77. if not self.sock:
  78. # Didn't connect, return the most recent error message
  79. raise socket.error(last_err)
  80. try:
  81. self.sock.settimeout(None)
  82. self.sock.setsockopt(SOL_TCP, socket.TCP_NODELAY, 1)
  83. self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
  84. self._setup_transport()
  85. self._write(AMQP_PROTOCOL_HEADER)
  86. except (OSError, IOError, socket.error) as exc:
  87. if get_errno(exc) not in _UNAVAIL:
  88. self.connected = False
  89. raise
  90. def __del__(self):
  91. try:
  92. # socket module may have been collected by gc
  93. # if this is called by a thread at shutdown.
  94. if socket is not None:
  95. try:
  96. self.close()
  97. except socket.error:
  98. pass
  99. finally:
  100. self.sock = None
  101. def _read(self, n, initial=False):
  102. """Read exactly n bytes from the peer"""
  103. raise NotImplementedError('Must be overriden in subclass')
  104. def _setup_transport(self):
  105. """Do any additional initialization of the class (used
  106. by the subclasses)."""
  107. pass
  108. def _shutdown_transport(self):
  109. """Do any preliminary work in shutting down the connection."""
  110. pass
  111. def _write(self, s):
  112. """Completely write a string to the peer."""
  113. raise NotImplementedError('Must be overriden in subclass')
  114. def close(self):
  115. if self.sock is not None:
  116. self._shutdown_transport()
  117. # Call shutdown first to make sure that pending messages
  118. # reach the AMQP broker if the program exits after
  119. # calling this method.
  120. self.sock.shutdown(socket.SHUT_RDWR)
  121. self.sock.close()
  122. self.sock = None
  123. self.connected = False
  124. def read_frame(self, unpack=unpack):
  125. read = self._read
  126. read_frame_buffer = EMPTY_BUFFER
  127. try:
  128. frame_header = read(7, True)
  129. read_frame_buffer += frame_header
  130. frame_type, channel, size = unpack('>BHI', frame_header)
  131. payload = read(size)
  132. read_frame_buffer += payload
  133. ch = ord(read(1))
  134. except socket.timeout:
  135. self._read_buffer = read_frame_buffer + self._read_buffer
  136. raise
  137. except (OSError, IOError, socket.error) as exc:
  138. # Don't disconnect for ssl read time outs
  139. # http://bugs.python.org/issue10272
  140. if isinstance(exc, SSLError) and 'timed out' in str(exc):
  141. raise socket.timeout()
  142. if get_errno(exc) not in _UNAVAIL:
  143. self.connected = False
  144. raise
  145. if ch == 206: # '\xce'
  146. return frame_type, channel, payload
  147. else:
  148. raise UnexpectedFrame(
  149. 'Received 0x{0:02x} while expecting 0xce'.format(ch))
  150. def write_frame(self, frame_type, channel, payload):
  151. size = len(payload)
  152. try:
  153. self._write(pack(
  154. '>BHI%dsB' % size,
  155. frame_type, channel, size, payload, 0xce,
  156. ))
  157. except socket.timeout:
  158. raise
  159. except (OSError, IOError, socket.error) as exc:
  160. if get_errno(exc) not in _UNAVAIL:
  161. self.connected = False
  162. raise
  163. class SSLTransport(_AbstractTransport):
  164. """Transport that works over SSL"""
  165. def __init__(self, host, connect_timeout, ssl):
  166. if isinstance(ssl, dict):
  167. self.sslopts = ssl
  168. self._read_buffer = EMPTY_BUFFER
  169. super(SSLTransport, self).__init__(host, connect_timeout)
  170. def _setup_transport(self):
  171. """Wrap the socket in an SSL object."""
  172. if hasattr(self, 'sslopts'):
  173. self.sock = ssl.wrap_socket(self.sock, **self.sslopts)
  174. else:
  175. self.sock = ssl.wrap_socket(self.sock)
  176. self.sock.do_handshake()
  177. self._quick_recv = self.sock.read
  178. def _shutdown_transport(self):
  179. """Unwrap a Python 2.6 SSL socket, so we can call shutdown()"""
  180. if self.sock is not None:
  181. try:
  182. unwrap = self.sock.unwrap
  183. except AttributeError:
  184. return
  185. self.sock = unwrap()
  186. def _read(self, n, initial=False,
  187. _errnos=(errno.ENOENT, errno.EAGAIN, errno.EINTR)):
  188. # According to SSL_read(3), it can at most return 16kb of data.
  189. # Thus, we use an internal read buffer like TCPTransport._read
  190. # to get the exact number of bytes wanted.
  191. recv = self._quick_recv
  192. rbuf = self._read_buffer
  193. try:
  194. while len(rbuf) < n:
  195. try:
  196. s = recv(n - len(rbuf)) # see note above
  197. except socket.error as exc:
  198. # ssl.sock.read may cause ENOENT if the
  199. # operation couldn't be performed (Issue celery#1414).
  200. if not initial and exc.errno in _errnos:
  201. continue
  202. raise
  203. if not s:
  204. raise IOError('Socket closed')
  205. rbuf += s
  206. except:
  207. self._read_buffer = rbuf
  208. raise
  209. result, self._read_buffer = rbuf[:n], rbuf[n:]
  210. return result
  211. def _write(self, s):
  212. """Write a string out to the SSL socket fully."""
  213. try:
  214. write = self.sock.write
  215. except AttributeError:
  216. # Works around a bug in python socket library
  217. raise IOError('Socket closed')
  218. else:
  219. while s:
  220. n = write(s)
  221. if not n:
  222. raise IOError('Socket closed')
  223. s = s[n:]
  224. class TCPTransport(_AbstractTransport):
  225. """Transport that deals directly with TCP socket."""
  226. def _setup_transport(self):
  227. """Setup to _write() directly to the socket, and
  228. do our own buffered reads."""
  229. self._write = self.sock.sendall
  230. self._read_buffer = EMPTY_BUFFER
  231. self._quick_recv = self.sock.recv
  232. def _read(self, n, initial=False, _errnos=(errno.EAGAIN, errno.EINTR)):
  233. """Read exactly n bytes from the socket"""
  234. recv = self._quick_recv
  235. rbuf = self._read_buffer
  236. try:
  237. while len(rbuf) < n:
  238. try:
  239. s = recv(n - len(rbuf))
  240. except socket.error as exc:
  241. if not initial and exc.errno in _errnos:
  242. continue
  243. raise
  244. if not s:
  245. raise IOError('Socket closed')
  246. rbuf += s
  247. except:
  248. self._read_buffer = rbuf
  249. raise
  250. result, self._read_buffer = rbuf[:n], rbuf[n:]
  251. return result
  252. def create_transport(host, connect_timeout, ssl=False):
  253. """Given a few parameters from the Connection constructor,
  254. select and create a subclass of _AbstractTransport."""
  255. if ssl:
  256. return SSLTransport(host, connect_timeout, ssl)
  257. else:
  258. return TCPTransport(host, connect_timeout)