_socket.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. """
  2. """
  3. """
  4. websocket - WebSocket client library for Python
  5. Copyright (C) 2010 Hiroki Ohtani(liris)
  6. This library is free software; you can redistribute it and/or
  7. modify it under the terms of the GNU Lesser General Public
  8. License as published by the Free Software Foundation; either
  9. version 2.1 of the License, or (at your option) any later version.
  10. This library is distributed in the hope that it will be useful,
  11. but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  13. Lesser General Public License for more details.
  14. You should have received a copy of the GNU Lesser General Public
  15. License along with this library; if not, write to the Free Software
  16. Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17. """
  18. import errno
  19. import select
  20. import socket
  21. import six
  22. from ._exceptions import *
  23. from ._ssl_compat import *
  24. from ._utils import *
  25. DEFAULT_SOCKET_OPTION = [(socket.SOL_TCP, socket.TCP_NODELAY, 1)]
  26. if hasattr(socket, "SO_KEEPALIVE"):
  27. DEFAULT_SOCKET_OPTION.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1))
  28. if hasattr(socket, "TCP_KEEPIDLE"):
  29. DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPIDLE, 30))
  30. if hasattr(socket, "TCP_KEEPINTVL"):
  31. DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPINTVL, 10))
  32. if hasattr(socket, "TCP_KEEPCNT"):
  33. DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPCNT, 3))
  34. _default_timeout = None
  35. __all__ = ["DEFAULT_SOCKET_OPTION", "sock_opt", "setdefaulttimeout", "getdefaulttimeout",
  36. "recv", "recv_line", "send"]
  37. class sock_opt(object):
  38. def __init__(self, sockopt, sslopt):
  39. if sockopt is None:
  40. sockopt = []
  41. if sslopt is None:
  42. sslopt = {}
  43. self.sockopt = sockopt
  44. self.sslopt = sslopt
  45. self.timeout = None
  46. def setdefaulttimeout(timeout):
  47. """
  48. Set the global timeout setting to connect.
  49. Parameters
  50. ----------
  51. timeout: int or float
  52. default socket timeout time (in seconds)
  53. """
  54. global _default_timeout
  55. _default_timeout = timeout
  56. def getdefaulttimeout():
  57. """
  58. Get default timeout
  59. Returns
  60. ----------
  61. _default_timeout: int or float
  62. Return the global timeout setting (in seconds) to connect.
  63. """
  64. return _default_timeout
  65. def recv(sock, bufsize):
  66. if not sock:
  67. raise WebSocketConnectionClosedException("socket is already closed.")
  68. def _recv():
  69. try:
  70. return sock.recv(bufsize)
  71. except SSLWantReadError:
  72. pass
  73. except socket.error as exc:
  74. error_code = extract_error_code(exc)
  75. if error_code is None:
  76. raise
  77. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  78. raise
  79. r, w, e = select.select((sock, ), (), (), sock.gettimeout())
  80. if r:
  81. return sock.recv(bufsize)
  82. try:
  83. if sock.gettimeout() == 0:
  84. bytes_ = sock.recv(bufsize)
  85. else:
  86. bytes_ = _recv()
  87. except socket.timeout as e:
  88. message = extract_err_message(e)
  89. raise WebSocketTimeoutException(message)
  90. except SSLError as e:
  91. message = extract_err_message(e)
  92. if isinstance(message, str) and 'timed out' in message:
  93. raise WebSocketTimeoutException(message)
  94. else:
  95. raise
  96. if not bytes_:
  97. raise WebSocketConnectionClosedException(
  98. "Connection is already closed.")
  99. return bytes_
  100. def recv_line(sock):
  101. line = []
  102. while True:
  103. c = recv(sock, 1)
  104. line.append(c)
  105. if c == six.b("\n"):
  106. break
  107. return six.b("").join(line)
  108. def send(sock, data):
  109. if isinstance(data, six.text_type):
  110. data = data.encode('utf-8')
  111. if not sock:
  112. raise WebSocketConnectionClosedException("socket is already closed.")
  113. def _send():
  114. try:
  115. return sock.send(data)
  116. except SSLWantWriteError:
  117. pass
  118. except socket.error as exc:
  119. error_code = extract_error_code(exc)
  120. if error_code is None:
  121. raise
  122. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  123. raise
  124. r, w, e = select.select((), (sock, ), (), sock.gettimeout())
  125. if w:
  126. return sock.send(data)
  127. try:
  128. if sock.gettimeout() == 0:
  129. return sock.send(data)
  130. else:
  131. return _send()
  132. except socket.timeout as e:
  133. message = extract_err_message(e)
  134. raise WebSocketTimeoutException(message)
  135. except Exception as e:
  136. message = extract_err_message(e)
  137. if isinstance(message, str) and "timed out" in message:
  138. raise WebSocketTimeoutException(message)
  139. else:
  140. raise