socketutils.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768
  1. # -*- coding: utf-8 -*-
  2. # Copyright (c) 2013, Mahmoud Hashemi
  3. #
  4. # Redistribution and use in source and binary forms, with or without
  5. # modification, are permitted provided that the following conditions are
  6. # met:
  7. #
  8. # * Redistributions of source code must retain the above copyright
  9. # notice, this list of conditions and the following disclaimer.
  10. #
  11. # * Redistributions in binary form must reproduce the above
  12. # copyright notice, this list of conditions and the following
  13. # disclaimer in the documentation and/or other materials provided
  14. # with the distribution.
  15. #
  16. # * The names of the contributors may not be used to endorse or
  17. # promote products derived from this software without specific
  18. # prior written permission.
  19. #
  20. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. """At its heart, Python can be viewed as an extension of the C
  32. programming language. Springing from the most popular systems
  33. programming language has made Python itself a great language for
  34. systems programming. One key to success in this domain is Python's
  35. very serviceable :mod:`socket` module and its :class:`socket.socket`
  36. type.
  37. The ``socketutils`` module provides natural next steps to the ``socket``
  38. builtin: straightforward, tested building blocks for higher-level
  39. protocols.
  40. The :class:`BufferedSocket` wraps an ordinary socket, providing a
  41. layer of intuitive buffering for both sending and receiving. This
  42. facilitates parsing messages from streams, i.e., all sockets with type
  43. ``SOCK_STREAM``. The BufferedSocket enables receiving until the next
  44. relevant token, up to a certain size, or until the connection is
  45. closed. For all of these, it provides consistent APIs to size
  46. limiting, as well as timeouts that are compatible with multiple
  47. concurrency paradigms. Use it to parse the next one-off text or binary
  48. socket protocol you encounter.
  49. This module also provides the :class:`NetstringSocket`, a pure-Python
  50. implementation of `the Netstring protocol`_, built on top of the
  51. :class:`BufferedSocket`, serving as a ready-made, production-grade example.
  52. Special thanks to `Kurt Rose`_ for his original authorship and all his
  53. contributions on this module. Also thanks to `Daniel J. Bernstein`_, the
  54. original author of `Netstring`_.
  55. .. _the Netstring protocol: https://en.wikipedia.org/wiki/Netstring
  56. .. _Kurt Rose: https://github.com/doublereedkurt
  57. .. _Daniel J. Bernstein: https://cr.yp.to/
  58. .. _Netstring: https://cr.yp.to/proto/netstrings.txt
  59. """
  60. import time
  61. import socket
  62. try:
  63. from threading import RLock
  64. except Exception:
  65. class RLock(object):
  66. 'Dummy reentrant lock for builds without threads'
  67. def __enter__(self):
  68. pass
  69. def __exit__(self, exctype, excinst, exctb):
  70. pass
  71. try:
  72. from .typeutils import make_sentinel
  73. _UNSET = make_sentinel(var_name='_UNSET')
  74. except ImportError:
  75. _UNSET = object()
  76. DEFAULT_TIMEOUT = 10 # 10 seconds
  77. DEFAULT_MAXSIZE = 32 * 1024 # 32kb
  78. _RECV_LARGE_MAXSIZE = 1024 ** 5 # 1PB
  79. class BufferedSocket(object):
  80. """Mainly provides recv_until and recv_size. recv, send, sendall, and
  81. peek all function as similarly as possible to the built-in socket
  82. API.
  83. This type has been tested against both the built-in socket type as
  84. well as those from gevent and eventlet. It also features support
  85. for sockets with timeouts set to 0 (aka nonblocking), provided the
  86. caller is prepared to handle the EWOULDBLOCK exceptions.
  87. Args:
  88. sock (socket): The connected socket to be wrapped.
  89. timeout (float): The default timeout for sends and recvs, in
  90. seconds. Set to ``None`` for no timeout, and 0 for
  91. nonblocking. Defaults to *sock*'s own timeout if already set,
  92. and 10 seconds otherwise.
  93. maxsize (int): The default maximum number of bytes to be received
  94. into the buffer before it is considered full and raises an
  95. exception. Defaults to 32 kilobytes.
  96. recvsize (int): The number of bytes to recv for every
  97. lower-level :meth:`socket.recv` call. Defaults to *maxsize*.
  98. *timeout* and *maxsize* can both be overridden on individual socket
  99. operations.
  100. All ``recv`` methods return bytestrings (:class:`bytes`) and can
  101. raise :exc:`socket.error`. :exc:`Timeout`,
  102. :exc:`ConnectionClosed`, and :exc:`MessageTooLong` all inherit
  103. from :exc:`socket.error` and exist to provide better error
  104. messages. Received bytes are always buffered, even if an exception
  105. is raised. Use :meth:`BufferedSocket.getrecvbuffer` to retrieve
  106. partial recvs.
  107. BufferedSocket does not replace the built-in socket by any
  108. means. While the overlapping parts of the API are kept parallel to
  109. the built-in :class:`socket.socket`, BufferedSocket does not
  110. inherit from socket, and most socket functionality is only
  111. available on the underlying socket. :meth:`socket.getpeername`,
  112. :meth:`socket.getsockname`, :meth:`socket.fileno`, and others are
  113. only available on the underlying socket that is wrapped. Use the
  114. ``BufferedSocket.sock`` attribute to access it. See the examples
  115. for more information on how to use BufferedSockets with built-in
  116. sockets.
  117. The BufferedSocket is threadsafe, but consider the semantics of
  118. your protocol before accessing a single socket from multiple
  119. threads. Similarly, once the BufferedSocket is constructed, avoid
  120. using the underlying socket directly. Only use it for operations
  121. unrelated to messages, e.g., :meth:`socket.getpeername`.
  122. """
  123. def __init__(self, sock, timeout=_UNSET,
  124. maxsize=DEFAULT_MAXSIZE, recvsize=_UNSET):
  125. self.sock = sock
  126. self.rbuf = b''
  127. self.sbuf = []
  128. self.maxsize = int(maxsize)
  129. if timeout is _UNSET:
  130. if self.sock.gettimeout() is None:
  131. self.timeout = DEFAULT_TIMEOUT
  132. else:
  133. self.timeout = self.sock.gettimeout()
  134. else:
  135. if timeout is None:
  136. self.timeout = timeout
  137. else:
  138. self.timeout = float(timeout)
  139. if recvsize is _UNSET:
  140. self._recvsize = self.maxsize
  141. else:
  142. self._recvsize = int(recvsize)
  143. self._send_lock = RLock()
  144. self._recv_lock = RLock()
  145. def settimeout(self, timeout):
  146. "Set the default *timeout* for future operations, in seconds."
  147. self.timeout = timeout
  148. def gettimeout(self):
  149. return self.timeout
  150. def setblocking(self, blocking):
  151. self.timeout = None if blocking else 0.0
  152. def setmaxsize(self, maxsize):
  153. """Set the default maximum buffer size *maxsize* for future
  154. operations, in bytes. Does not truncate the current buffer.
  155. """
  156. self.maxsize = maxsize
  157. def getrecvbuffer(self):
  158. "Returns the receive buffer bytestring (rbuf)."
  159. with self._recv_lock:
  160. return self.rbuf
  161. def getsendbuffer(self):
  162. "Returns a copy of the send buffer list."
  163. with self._send_lock:
  164. return b''.join(self.sbuf)
  165. def recv(self, size, flags=0, timeout=_UNSET):
  166. """Returns **up to** *size* bytes, using the internal buffer before
  167. performing a single :meth:`socket.recv` operation.
  168. Args:
  169. size (int): The maximum number of bytes to receive.
  170. flags (int): Kept for API compatibility with sockets. Only
  171. the default, ``0``, is valid.
  172. timeout (float): The timeout for this operation. Can be
  173. ``0`` for nonblocking and ``None`` for no
  174. timeout. Defaults to the value set in the constructor
  175. of BufferedSocket.
  176. If the operation does not complete in *timeout* seconds, a
  177. :exc:`Timeout` is raised. Much like the built-in
  178. :class:`socket.socket`, if this method returns an empty string,
  179. then the socket is closed and recv buffer is empty. Further
  180. calls to recv will raise :exc:`socket.error`.
  181. """
  182. with self._recv_lock:
  183. if timeout is _UNSET:
  184. timeout = self.timeout
  185. if flags:
  186. raise ValueError("non-zero flags not supported: %r" % flags)
  187. if len(self.rbuf) >= size:
  188. data, self.rbuf = self.rbuf[:size], self.rbuf[size:]
  189. return data
  190. if self.rbuf:
  191. ret, self.rbuf = self.rbuf, b''
  192. return ret
  193. self.sock.settimeout(timeout)
  194. try:
  195. data = self.sock.recv(self._recvsize)
  196. except socket.timeout:
  197. raise Timeout(timeout) # check the rbuf attr for more
  198. if len(data) > size:
  199. data, self.rbuf = data[:size], data[size:]
  200. return data
  201. def peek(self, size, timeout=_UNSET):
  202. """Returns *size* bytes from the socket and/or internal buffer. Bytes
  203. are retained in BufferedSocket's internal recv buffer. To only
  204. see bytes in the recv buffer, use :meth:`getrecvbuffer`.
  205. Args:
  206. size (int): The exact number of bytes to peek at
  207. timeout (float): The timeout for this operation. Can be 0 for
  208. nonblocking and None for no timeout. Defaults to the value
  209. set in the constructor of BufferedSocket.
  210. If the appropriate number of bytes cannot be fetched from the
  211. buffer and socket before *timeout* expires, then a
  212. :exc:`Timeout` will be raised. If the connection is closed, a
  213. :exc:`ConnectionClosed` will be raised.
  214. """
  215. with self._recv_lock:
  216. if len(self.rbuf) >= size:
  217. return self.rbuf[:size]
  218. data = self.recv_size(size, timeout=timeout)
  219. self.rbuf = data + self.rbuf
  220. return data
  221. def recv_close(self, timeout=_UNSET, maxsize=_UNSET):
  222. """Receive until the connection is closed, up to *maxsize* bytes. If
  223. more than *maxsize* bytes are received, raises :exc:`MessageTooLong`.
  224. """
  225. # recv_close works by using recv_size to request maxsize data,
  226. # and ignoring ConnectionClose, returning and clearing the
  227. # internal buffer instead. It raises an exception if
  228. # ConnectionClosed isn't raised.
  229. with self._recv_lock:
  230. if maxsize is _UNSET:
  231. maxsize = self.maxsize
  232. if maxsize is None:
  233. maxsize = _RECV_LARGE_MAXSIZE
  234. try:
  235. recvd = self.recv_size(maxsize + 1, timeout)
  236. except ConnectionClosed:
  237. ret, self.rbuf = self.rbuf, b''
  238. else:
  239. # put extra received bytes (now in rbuf) after recvd
  240. self.rbuf = recvd + self.rbuf
  241. size_read = min(maxsize, len(self.rbuf))
  242. raise MessageTooLong(size_read) # check receive buffer
  243. return ret
  244. def recv_until(self, delimiter, timeout=_UNSET, maxsize=_UNSET,
  245. with_delimiter=False):
  246. """Receive until *delimiter* is found, *maxsize* bytes have been read,
  247. or *timeout* is exceeded.
  248. Args:
  249. delimiter (bytes): One or more bytes to be searched for
  250. in the socket stream.
  251. timeout (float): The timeout for this operation. Can be 0 for
  252. nonblocking and None for no timeout. Defaults to the value
  253. set in the constructor of BufferedSocket.
  254. maxsize (int): The maximum size for the internal buffer.
  255. Defaults to the value set in the constructor.
  256. with_delimiter (bool): Whether or not to include the
  257. delimiter in the output. ``False`` by default, but
  258. ``True`` is useful in cases where one is simply
  259. forwarding the messages.
  260. ``recv_until`` will raise the following exceptions:
  261. * :exc:`Timeout` if more than *timeout* seconds expire.
  262. * :exc:`ConnectionClosed` if the underlying socket is closed
  263. by the sending end.
  264. * :exc:`MessageTooLong` if the delimiter is not found in the
  265. first *maxsize* bytes.
  266. * :exc:`socket.error` if operating in nonblocking mode
  267. (*timeout* equal to 0), or if some unexpected socket error
  268. occurs, such as operating on a closed socket.
  269. """
  270. with self._recv_lock:
  271. if maxsize is _UNSET:
  272. maxsize = self.maxsize
  273. if maxsize is None:
  274. maxsize = _RECV_LARGE_MAXSIZE
  275. if timeout is _UNSET:
  276. timeout = self.timeout
  277. len_delimiter = len(delimiter)
  278. sock = self.sock
  279. recvd = bytearray(self.rbuf)
  280. start = time.time()
  281. find_offset_start = 0 # becomes a negative index below
  282. if not timeout: # covers None (no timeout) and 0 (nonblocking)
  283. sock.settimeout(timeout)
  284. try:
  285. while 1:
  286. offset = recvd.find(delimiter, find_offset_start, maxsize)
  287. if offset != -1: # str.find returns -1 when no match found
  288. if with_delimiter: # include delimiter in return
  289. offset += len_delimiter
  290. rbuf_offset = offset
  291. else:
  292. rbuf_offset = offset + len_delimiter
  293. break
  294. elif len(recvd) > maxsize:
  295. raise MessageTooLong(maxsize, delimiter) # see rbuf
  296. if timeout:
  297. cur_timeout = timeout - (time.time() - start)
  298. if cur_timeout <= 0.0:
  299. raise socket.timeout()
  300. sock.settimeout(cur_timeout)
  301. nxt = sock.recv(self._recvsize)
  302. if not nxt:
  303. args = (len(recvd), delimiter)
  304. msg = ('connection closed after reading %s bytes'
  305. ' without finding symbol: %r' % args)
  306. raise ConnectionClosed(msg) # check the recv buffer
  307. recvd.extend(nxt)
  308. find_offset_start = -len(nxt) - len_delimiter + 1
  309. except socket.timeout:
  310. self.rbuf = bytes(recvd)
  311. msg = ('read %s bytes without finding delimiter: %r'
  312. % (len(recvd), delimiter))
  313. raise Timeout(timeout, msg) # check the recv buffer
  314. except Exception:
  315. self.rbuf = bytes(recvd)
  316. raise
  317. val, self.rbuf = bytes(recvd[:offset]), bytes(recvd[rbuf_offset:])
  318. return val
  319. def recv_size(self, size, timeout=_UNSET):
  320. """Read off of the internal buffer, then off the socket, until
  321. *size* bytes have been read.
  322. Args:
  323. size (int): number of bytes to read before returning.
  324. timeout (float): The timeout for this operation. Can be 0 for
  325. nonblocking and None for no timeout. Defaults to the value
  326. set in the constructor of BufferedSocket.
  327. If the appropriate number of bytes cannot be fetched from the
  328. buffer and socket before *timeout* expires, then a
  329. :exc:`Timeout` will be raised. If the connection is closed, a
  330. :exc:`ConnectionClosed` will be raised.
  331. """
  332. with self._recv_lock:
  333. if timeout is _UNSET:
  334. timeout = self.timeout
  335. chunks = []
  336. total_bytes = 0
  337. try:
  338. start = time.time()
  339. self.sock.settimeout(timeout)
  340. nxt = self.rbuf or self.sock.recv(self._recvsize)
  341. while nxt:
  342. total_bytes += len(nxt)
  343. if total_bytes >= size:
  344. break
  345. chunks.append(nxt)
  346. if timeout:
  347. cur_timeout = timeout - (time.time() - start)
  348. if cur_timeout <= 0.0:
  349. raise socket.timeout()
  350. self.sock.settimeout(cur_timeout)
  351. nxt = self.sock.recv(self._recvsize)
  352. else:
  353. msg = ('connection closed after reading %s of %s requested'
  354. ' bytes' % (total_bytes, size))
  355. raise ConnectionClosed(msg) # check recv buffer
  356. except socket.timeout:
  357. self.rbuf = b''.join(chunks)
  358. msg = 'read %s of %s bytes' % (total_bytes, size)
  359. raise Timeout(timeout, msg) # check recv buffer
  360. except Exception:
  361. # received data is still buffered in the case of errors
  362. self.rbuf = b''.join(chunks)
  363. raise
  364. extra_bytes = total_bytes - size
  365. if extra_bytes:
  366. last, self.rbuf = nxt[:-extra_bytes], nxt[-extra_bytes:]
  367. else:
  368. last, self.rbuf = nxt, b''
  369. chunks.append(last)
  370. return b''.join(chunks)
  371. def send(self, data, flags=0, timeout=_UNSET):
  372. """Send the contents of the internal send buffer, as well as *data*,
  373. to the receiving end of the connection. Returns the total
  374. number of bytes sent. If no exception is raised, all of *data* was
  375. sent and the internal send buffer is empty.
  376. Args:
  377. data (bytes): The bytes to send.
  378. flags (int): Kept for API compatibility with sockets. Only
  379. the default 0 is valid.
  380. timeout (float): The timeout for this operation. Can be 0 for
  381. nonblocking and None for no timeout. Defaults to the value
  382. set in the constructor of BufferedSocket.
  383. Will raise :exc:`Timeout` if the send operation fails to
  384. complete before *timeout*. In the event of an exception, use
  385. :meth:`BufferedSocket.getsendbuffer` to see which data was
  386. unsent.
  387. """
  388. with self._send_lock:
  389. if timeout is _UNSET:
  390. timeout = self.timeout
  391. if flags:
  392. raise ValueError("non-zero flags not supported")
  393. sbuf = self.sbuf
  394. sbuf.append(data)
  395. if len(sbuf) > 1:
  396. sbuf[:] = [b''.join([s for s in sbuf if s])]
  397. self.sock.settimeout(timeout)
  398. start, total_sent = time.time(), 0
  399. try:
  400. while sbuf[0]:
  401. sent = self.sock.send(sbuf[0])
  402. total_sent += sent
  403. sbuf[0] = sbuf[0][sent:]
  404. if timeout:
  405. cur_timeout = timeout - (time.time() - start)
  406. if cur_timeout <= 0.0:
  407. raise socket.timeout()
  408. self.sock.settimeout(cur_timeout)
  409. except socket.timeout:
  410. raise Timeout(timeout, '%s bytes unsent' % len(sbuf[0]))
  411. return total_sent
  412. def sendall(self, data, flags=0, timeout=_UNSET):
  413. """A passthrough to :meth:`~BufferedSocket.send`, retained for
  414. parallelism to the :class:`socket.socket` API.
  415. """
  416. return self.send(data, flags, timeout)
  417. def flush(self):
  418. "Send the contents of the internal send buffer."
  419. with self._send_lock:
  420. self.send(b'')
  421. return
  422. def buffer(self, data):
  423. "Buffer *data* bytes for the next send operation."
  424. with self._send_lock:
  425. self.sbuf.append(data)
  426. return
  427. # # #
  428. # # # Passing through some socket basics
  429. # # #
  430. def getsockname(self):
  431. """Convenience function to return the wrapped socket's own address.
  432. See :meth:`socket.getsockname` for more details.
  433. """
  434. return self.sock.getsockname()
  435. def getpeername(self):
  436. """Convenience function to return the remote address to which the
  437. wrapped socket is connected. See :meth:`socket.getpeername`
  438. for more details.
  439. """
  440. return self.sock.getpeername()
  441. def getsockopt(self, level, optname, buflen=None):
  442. """Convenience function passing through to the wrapped socket's
  443. :meth:`socket.getsockopt`.
  444. """
  445. args = (level, optname)
  446. if buflen is not None:
  447. args += (buflen,)
  448. return self.sock.getsockopt(*args)
  449. def setsockopt(self, level, optname, value):
  450. """Convenience function passing through to the wrapped socket's
  451. :meth:`socket.setsockopt`.
  452. """
  453. return self.sock.setsockopt(level, optname, value)
  454. @property
  455. def type(self):
  456. """A passthrough to the wrapped socket's type. Valid usages should
  457. only ever see :data:`socket.SOCK_STREAM`.
  458. """
  459. return self.sock.type
  460. @property
  461. def family(self):
  462. """A passthrough to the wrapped socket's family. BufferedSocket
  463. supports all widely-used families, so this read-only attribute
  464. can be one of :data:`socket.AF_INET` for IP,
  465. :data:`socket.AF_INET6` for IPv6, and :data:`socket.AF_UNIX`
  466. for UDS.
  467. """
  468. return self.sock.family
  469. @property
  470. def proto(self):
  471. """A passthrough to the wrapped socket's protocol. The ``proto``
  472. attribute is very rarely used, so it's always 0, meaning "the
  473. default" protocol. Pretty much all the practical information
  474. is in :attr:`~BufferedSocket.type` and
  475. :attr:`~BufferedSocket.family`, so you can go back to never
  476. thinking about this.
  477. """
  478. return self.sock.proto
  479. # # #
  480. # # # Now for some more advanced interpretations of the builtin socket
  481. # # #
  482. def fileno(self):
  483. """Returns the file descriptor of the wrapped socket. -1 if it has
  484. been closed on this end.
  485. Note that this makes the BufferedSocket selectable, i.e.,
  486. usable for operating system event loops without any external
  487. libraries. Keep in mind that the operating system cannot know
  488. about data in BufferedSocket's internal buffer. Exercise
  489. discipline with calling ``recv*`` functions.
  490. """
  491. return self.sock.fileno()
  492. def close(self):
  493. """Closes the wrapped socket, and empties the internal buffers. The
  494. send buffer is not flushed automatically, so if you have been
  495. calling :meth:`~BufferedSocket.buffer`, be sure to call
  496. :meth:`~BufferedSocket.flush` before calling this
  497. method. After calling this method, future socket operations
  498. will raise :exc:`socket.error`.
  499. """
  500. with self._recv_lock:
  501. with self._send_lock:
  502. self.rbuf = b''
  503. self.rbuf_unconsumed = self.rbuf
  504. self.sbuf[:] = []
  505. self.sock.close()
  506. return
  507. def shutdown(self, how):
  508. """Convenience method which passes through to the wrapped socket's
  509. :meth:`~socket.shutdown`. Semantics vary by platform, so no
  510. special internal handling is done with the buffers. This
  511. method exists to facilitate the most common usage, wherein a
  512. full ``shutdown`` is followed by a
  513. :meth:`~BufferedSocket.close`. Developers requiring more
  514. support, please open `an issue`_.
  515. .. _an issue: https://github.com/mahmoud/boltons/issues
  516. """
  517. with self._recv_lock:
  518. with self._send_lock:
  519. self.sock.shutdown(how)
  520. return
  521. # end BufferedSocket
  522. class Error(socket.error):
  523. """A subclass of :exc:`socket.error` from which all other
  524. ``socketutils`` exceptions inherit.
  525. When using :class:`BufferedSocket` and other ``socketutils``
  526. types, generally you want to catch one of the specific exception
  527. types below, or :exc:`socket.error`.
  528. """
  529. pass
  530. class ConnectionClosed(Error):
  531. """Raised when receiving and the connection is unexpectedly closed
  532. from the sending end. Raised from :class:`BufferedSocket`'s
  533. :meth:`~BufferedSocket.peek`, :meth:`~BufferedSocket.recv_until`,
  534. and :meth:`~BufferedSocket.recv_size`, and never from its
  535. :meth:`~BufferedSocket.recv` or
  536. :meth:`~BufferedSocket.recv_close`.
  537. """
  538. pass
  539. class MessageTooLong(Error):
  540. """Raised from :meth:`BufferedSocket.recv_until` and
  541. :meth:`BufferedSocket.recv_closed` when more than *maxsize* bytes are
  542. read without encountering the delimiter or a closed connection,
  543. respectively.
  544. """
  545. def __init__(self, bytes_read=None, delimiter=None):
  546. msg = 'message exceeded maximum size'
  547. if bytes_read is not None:
  548. msg += '. %s bytes read' % (bytes_read,)
  549. if delimiter is not None:
  550. msg += '. Delimiter not found: %r' % (delimiter,)
  551. super(MessageTooLong, self).__init__(msg)
  552. class Timeout(socket.timeout, Error):
  553. """Inheriting from :exc:`socket.timeout`, Timeout is used to indicate
  554. when a socket operation did not complete within the time
  555. specified. Raised from any of :class:`BufferedSocket`'s ``recv``
  556. methods.
  557. """
  558. def __init__(self, timeout, extra=""):
  559. msg = 'socket operation timed out'
  560. if timeout is not None:
  561. msg += ' after %sms.' % (timeout * 1000)
  562. if extra:
  563. msg += ' ' + extra
  564. super(Timeout, self).__init__(msg)
  565. class NetstringSocket(object):
  566. """
  567. Reads and writes using the netstring protocol.
  568. More info: https://en.wikipedia.org/wiki/Netstring
  569. Even more info: http://cr.yp.to/proto/netstrings.txt
  570. """
  571. def __init__(self, sock, timeout=DEFAULT_TIMEOUT, maxsize=DEFAULT_MAXSIZE):
  572. self.bsock = BufferedSocket(sock)
  573. self.timeout = timeout
  574. self.maxsize = maxsize
  575. self._msgsize_maxsize = len(str(maxsize)) + 1 # len(str()) == log10
  576. def fileno(self):
  577. return self.bsock.fileno()
  578. def settimeout(self, timeout):
  579. self.timeout = timeout
  580. def setmaxsize(self, maxsize):
  581. self.maxsize = maxsize
  582. self._msgsize_maxsize = self._calc_msgsize_maxsize(maxsize)
  583. def _calc_msgsize_maxsize(self, maxsize):
  584. return len(str(maxsize)) + 1 # len(str()) == log10
  585. def read_ns(self, timeout=_UNSET, maxsize=_UNSET):
  586. if timeout is _UNSET:
  587. timeout = self.timeout
  588. if maxsize is _UNSET:
  589. maxsize = self.maxsize
  590. msgsize_maxsize = self._msgsize_maxsize
  591. else:
  592. msgsize_maxsize = self._calc_msgsize_maxsize(maxsize)
  593. size_prefix = self.bsock.recv_until(b':',
  594. timeout=timeout,
  595. maxsize=msgsize_maxsize)
  596. try:
  597. size = int(size_prefix)
  598. except ValueError:
  599. raise NetstringInvalidSize('netstring message size must be valid'
  600. ' integer, not %r' % size_prefix)
  601. if size > maxsize:
  602. raise NetstringMessageTooLong(size, maxsize)
  603. payload = self.bsock.recv_size(size)
  604. if self.bsock.recv(1) != b',':
  605. raise NetstringProtocolError("expected trailing ',' after message")
  606. return payload
  607. def write_ns(self, payload):
  608. size = len(payload)
  609. if size > self.maxsize:
  610. raise NetstringMessageTooLong(size, self.maxsize)
  611. data = str(size).encode('ascii') + b':' + payload + b','
  612. self.bsock.send(data)
  613. class NetstringProtocolError(Error):
  614. "Base class for all of socketutils' Netstring exception types."
  615. pass
  616. class NetstringInvalidSize(NetstringProtocolError):
  617. """NetstringInvalidSize is raised when the ``:``-delimited size prefix
  618. of the message does not contain a valid integer.
  619. Message showing valid size::
  620. 5:hello,
  621. Here the ``5`` is the size. Anything in this prefix position that
  622. is not parsable as a Python integer (i.e., :class:`int`) will raise
  623. this exception.
  624. """
  625. def __init__(self, msg):
  626. super(NetstringInvalidSize, self).__init__(msg)
  627. class NetstringMessageTooLong(NetstringProtocolError):
  628. """NetstringMessageTooLong is raised when the size prefix contains a
  629. valid integer, but that integer is larger than the
  630. :class:`NetstringSocket`'s configured *maxsize*.
  631. When this exception is raised, it's recommended to simply close
  632. the connection instead of trying to recover.
  633. """
  634. def __init__(self, size, maxsize):
  635. msg = ('netstring message length exceeds configured maxsize: %s > %s'
  636. % (size, maxsize))
  637. super(NetstringMessageTooLong, self).__init__(msg)
  638. """
  639. attrs worth adding/passing through:
  640. properties: type, proto
  641. For its main functionality, BufferedSocket can wrap any object that
  642. has the following methods:
  643. - gettimeout()
  644. - settimeout()
  645. - recv(size)
  646. - send(data)
  647. The following methods are passed through:
  648. ...
  649. """
  650. # TODO: buffered socket check socket.type == SOCK_STREAM?
  651. # TODO: make recv_until support taking a regex
  652. # TODO: including the delimiter in the recv_until return is not
  653. # necessary, as ConnectionClosed differentiates empty messages
  654. # from socket closes.