test_socketutils.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429
  1. # -*- coding: utf-8 -*-
  2. import sys
  3. import time
  4. import errno
  5. import socket
  6. import threading
  7. from boltons.socketutils import (BufferedSocket,
  8. NetstringSocket,
  9. ConnectionClosed,
  10. NetstringMessageTooLong,
  11. MessageTooLong,
  12. Timeout)
  13. import pytest
  14. # skip if there's no socketpair
  15. pytestmark = pytest.mark.skipif(getattr(socket, 'socketpair', None) is None,
  16. reason='no socketpair (likely Py2 on Windows)')
  17. def test_short_lines():
  18. for ms in (2, 4, 6, 1024, None):
  19. x, y = socket.socketpair()
  20. bs = BufferedSocket(x)
  21. y.sendall(b'1\n2\n3\n')
  22. assert bs.recv_until(b'\n', maxsize=ms) == b'1'
  23. assert bs.recv_until(b'\n', maxsize=ms) == b'2'
  24. y.close()
  25. assert bs.recv_close(maxsize=ms) == b'3\n'
  26. try:
  27. bs.recv_size(1)
  28. except ConnectionClosed:
  29. pass
  30. else:
  31. assert False, 'expected ConnectionClosed'
  32. bs.close()
  33. return
  34. def test_multibyte_delim():
  35. """Primarily tests recv_until with various maxsizes and True/False
  36. for with_delimiter.
  37. """
  38. delim = b'\r\n'
  39. for with_delim in (True, False):
  40. if with_delim:
  41. cond_delim = b'\r\n'
  42. else:
  43. cond_delim = b''
  44. empty = b''
  45. small_one = b'1'
  46. big_two = b'2' * 2048
  47. for ms in (3, 5, 1024, None):
  48. x, y = socket.socketpair()
  49. bs = BufferedSocket(x)
  50. y.sendall(empty + delim)
  51. y.sendall(small_one + delim)
  52. y.sendall(big_two + delim)
  53. kwargs = {'maxsize': ms, 'with_delimiter': with_delim}
  54. assert bs.recv_until(delim, **kwargs) == empty + cond_delim
  55. assert bs.recv_until(delim, **kwargs) == small_one + cond_delim
  56. try:
  57. assert bs.recv_until(delim, **kwargs) == big_two + cond_delim
  58. except MessageTooLong:
  59. if ms is None:
  60. assert False, 'unexpected MessageTooLong'
  61. else:
  62. if ms is not None:
  63. assert False, 'expected MessageTooLong'
  64. return
  65. def test_props():
  66. x, y = socket.socketpair()
  67. bs = BufferedSocket(x)
  68. assert bs.type == x.type
  69. assert bs.proto == x.proto
  70. assert bs.family == x.family
  71. return
  72. def test_buffers():
  73. x, y = socket.socketpair()
  74. bx, by = BufferedSocket(x), BufferedSocket(y)
  75. assert by.getrecvbuffer() == b''
  76. assert by.getsendbuffer() == b''
  77. assert bx.getrecvbuffer() == b''
  78. by.buffer(b'12')
  79. by.sendall(b'3')
  80. assert bx.recv_size(1) == b'1'
  81. assert bx.getrecvbuffer() == b'23'
  82. return
  83. IS_PYPY_2 = ('__pypy__' in sys.builtin_module_names
  84. and sys.version_info[0] == 2)
  85. @pytest.mark.xfail(IS_PYPY_2, reason="pypy2 bug, fixed in 7.2. unmark when this test stops failing on travis (when they upgrade from 7.1)")
  86. def test_client_disconnecting():
  87. def get_bs_pair():
  88. x, y = socket.socketpair()
  89. bx, by = BufferedSocket(x), BufferedSocket(y)
  90. # sanity check
  91. by.sendall(b'123')
  92. bx.recv_size(3) == b'123'
  93. return bx, by
  94. bx, by = get_bs_pair()
  95. assert bx.fileno() > 0
  96. bx.close()
  97. assert bx.getrecvbuffer() == b''
  98. try:
  99. bx.recv(1)
  100. except socket.error:
  101. pass
  102. else:
  103. assert False, 'expected socket.error on closed recv'
  104. assert bx.fileno() == -1
  105. by.buffer(b'123')
  106. assert by.getsendbuffer()
  107. try:
  108. by.flush()
  109. except socket.error:
  110. assert by.getsendbuffer() == b'123'
  111. else:
  112. if sys.platform != 'win32': # Windows socketpairs are kind of bad
  113. assert False, 'expected socket.error broken pipe'
  114. try:
  115. by.shutdown(socket.SHUT_RDWR)
  116. except socket.error:
  117. # Mac sockets are already shut down at this point. See #71.
  118. if sys.platform != 'darwin':
  119. raise
  120. by.close()
  121. assert not by.getsendbuffer()
  122. try:
  123. by.send(b'123')
  124. except socket.error:
  125. pass
  126. else:
  127. assert False, 'expected socket.error on closed send'
  128. return
  129. def test_split_delim():
  130. delim = b'\r\n'
  131. first = b'1234\r'
  132. second = b'\n5'
  133. x, y = socket.socketpair()
  134. bs = BufferedSocket(x)
  135. y.sendall(first)
  136. try:
  137. bs.recv_until(delim, timeout=0.0001)
  138. except Timeout:
  139. pass
  140. y.sendall(second)
  141. assert bs.recv_until(delim, with_delimiter=True) == b'1234\r\n'
  142. assert bs.recv_size(1) == b'5'
  143. return
  144. def test_basic_nonblocking():
  145. delim = b'\n'
  146. # test with per-call timeout
  147. x, y = socket.socketpair()
  148. bs = BufferedSocket(x)
  149. try:
  150. bs.recv_until(delim, timeout=0)
  151. except socket.error as se:
  152. assert se.errno == errno.EWOULDBLOCK
  153. y.sendall(delim) # sending an empty message, effectively
  154. assert bs.recv_until(delim) == b''
  155. # test with instance-level default timeout
  156. x, y = socket.socketpair()
  157. bs = BufferedSocket(x, timeout=0)
  158. try:
  159. bs.recv_until(delim)
  160. except socket.error as se:
  161. assert se.errno == errno.EWOULDBLOCK
  162. y.sendall(delim)
  163. assert bs.recv_until(delim) == b''
  164. # test with setblocking(0) on the underlying socket
  165. x, y = socket.socketpair()
  166. x.setblocking(0)
  167. bs = BufferedSocket(x)
  168. try:
  169. bs.recv_until(delim)
  170. except socket.error as se:
  171. assert se.errno == errno.EWOULDBLOCK
  172. y.sendall(delim)
  173. assert bs.recv_until(delim) == b''
  174. return
  175. def test_simple_buffered_socket_passthroughs():
  176. x, y = socket.socketpair()
  177. bs = BufferedSocket(x)
  178. assert bs.getsockname() == x.getsockname()
  179. assert bs.getpeername() == x.getpeername()
  180. def test_timeout_setters_getters():
  181. x, y = socket.socketpair()
  182. bs = BufferedSocket(x)
  183. assert bs.settimeout(1.0) is None
  184. assert bs.gettimeout() == 1.0
  185. assert bs.setblocking(False) is None
  186. assert bs.gettimeout() == 0.0
  187. assert bs.setblocking(True) is None
  188. assert bs.gettimeout() is None
  189. def netstring_server(server_socket):
  190. "A basic netstring server loop, supporting a few operations"
  191. try:
  192. while True:
  193. clientsock, addr = server_socket.accept()
  194. client = NetstringSocket(clientsock)
  195. while 1:
  196. request = client.read_ns()
  197. if request == b'close':
  198. clientsock.close()
  199. break
  200. elif request == b'shutdown':
  201. return
  202. elif request == b'reply4k':
  203. client.write_ns(b'a' * 4096)
  204. elif request == b'ping':
  205. client.write_ns(b'pong')
  206. elif request == b'reply128k':
  207. client.setmaxsize(128 * 1024)
  208. client.write_ns(b'huge' * 32 * 1024) # 128kb
  209. client.setmaxsize(32768) # back to default
  210. except Exception as e:
  211. print(u'netstring_server exiting with error: %r' % e)
  212. raise
  213. def test_socketutils_netstring():
  214. """A holistic feature test of BufferedSocket via the NetstringSocket
  215. wrapper. Runs
  216. """
  217. print("running self tests")
  218. # Set up server
  219. server_socket = socket.socket()
  220. server_socket.bind(('127.0.0.1', 0)) # localhost with ephemeral port
  221. server_socket.listen(100)
  222. ip, port = server_socket.getsockname()
  223. start_server = lambda: netstring_server(server_socket)
  224. threading.Thread(target=start_server).start()
  225. # set up client
  226. def client_connect():
  227. clientsock = socket.create_connection((ip, port))
  228. client = NetstringSocket(clientsock)
  229. return client
  230. # connect, ping-pong
  231. client = client_connect()
  232. client.write_ns(b'ping')
  233. assert client.read_ns() == b'pong'
  234. s = time.time()
  235. for i in range(1000):
  236. client.write_ns(b'ping')
  237. assert client.read_ns() == b'pong'
  238. dur = time.time() - s
  239. print("netstring ping-pong latency", dur, "ms")
  240. s = time.time()
  241. for i in range(1000):
  242. client.write_ns(b'ping')
  243. resps = []
  244. for i in range(1000):
  245. resps.append(client.read_ns())
  246. e = time.time()
  247. assert all([r == b'pong' for r in resps])
  248. assert client.bsock.getrecvbuffer() == b''
  249. dur = e - s
  250. print("netstring pipelined ping-pong latency", dur, "ms")
  251. # tell the server to close the socket and then try a failure case
  252. client.write_ns(b'close')
  253. try:
  254. client.read_ns()
  255. raise Exception('read from closed socket')
  256. except ConnectionClosed:
  257. print("raised ConnectionClosed correctly")
  258. # test big messages
  259. client = client_connect()
  260. client.setmaxsize(128 * 1024)
  261. client.write_ns(b'reply128k')
  262. res = client.read_ns()
  263. assert len(res) == (128 * 1024)
  264. client.write_ns(b'close')
  265. # test that read timeouts work
  266. client = client_connect()
  267. client.settimeout(0.1)
  268. try:
  269. client.read_ns()
  270. raise Exception('did not timeout')
  271. except Timeout:
  272. print("read_ns raised timeout correctly")
  273. client.write_ns(b'close')
  274. # test that netstring max sizes work
  275. client = client_connect()
  276. client.setmaxsize(2048)
  277. client.write_ns(b'reply4k')
  278. try:
  279. client.read_ns()
  280. raise Exception('read more than maxsize')
  281. except NetstringMessageTooLong:
  282. print("raised MessageTooLong correctly")
  283. try:
  284. client.bsock.recv_until(b'b', maxsize=4096)
  285. raise Exception('recv_until did not raise MessageTooLong')
  286. except MessageTooLong:
  287. print("raised MessageTooLong correctly")
  288. assert client.bsock.recv_size(4097) == b'a' * 4096 + b','
  289. print('correctly maintained buffer after exception raised')
  290. # test BufferedSocket read timeouts with recv_until and recv_size
  291. client.bsock.settimeout(0.01)
  292. try:
  293. client.bsock.recv_until(b'a')
  294. raise Exception('recv_until did not raise Timeout')
  295. except Timeout:
  296. print('recv_until correctly raised Timeout')
  297. try:
  298. client.bsock.recv_size(1)
  299. raise Exception('recv_size did not raise Timeout')
  300. except Timeout:
  301. print('recv_size correctly raised Timeout')
  302. client.write_ns(b'shutdown')
  303. print("all passed")
  304. def netstring_server_timeout_override(server_socket):
  305. """Netstring socket has an unreasonably low timeout,
  306. however it should be overridden by the `read_ns` argument."""
  307. try:
  308. while True:
  309. clientsock, addr = server_socket.accept()
  310. client = NetstringSocket(clientsock, timeout=0.01)
  311. while 1:
  312. request = client.read_ns(1)
  313. if request == b'close':
  314. clientsock.close()
  315. break
  316. elif request == b'shutdown':
  317. return
  318. elif request == b'ping':
  319. client.write_ns(b'pong')
  320. except Exception as e:
  321. print(u'netstring_server exiting with error: %r' % e)
  322. raise
  323. def test_socketutils_netstring_timeout():
  324. """Tests that server socket timeout is overridden by the argument to read call.
  325. Server has timeout of 10 ms, and we will sleep for 20 ms. If timeout is not overridden correctly,
  326. a timeout exception will be raised."""
  327. print("running timeout test")
  328. # Set up server
  329. server_socket = socket.socket()
  330. server_socket.bind(('127.0.0.1', 0)) # localhost with ephemeral port
  331. server_socket.listen(100)
  332. ip, port = server_socket.getsockname()
  333. start_server = lambda: netstring_server_timeout_override(server_socket)
  334. threading.Thread(target=start_server).start()
  335. # set up client
  336. def client_connect():
  337. clientsock = socket.create_connection((ip, port))
  338. client = NetstringSocket(clientsock)
  339. return client
  340. # connect, ping-pong
  341. client = client_connect()
  342. time.sleep(0.02)
  343. client.write_ns(b'ping')
  344. assert client.read_ns() == b'pong'
  345. client.write_ns(b'shutdown')
  346. print("no timeout occurred - all good.")