_socket.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  1. """
  2. websocket - WebSocket client library for Python
  3. Copyright (C) 2010 Hiroki Ohtani(liris)
  4. This library is free software; you can redistribute it and/or
  5. modify it under the terms of the GNU Lesser General Public
  6. License as published by the Free Software Foundation; either
  7. version 2.1 of the License, or (at your option) any later version.
  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. You should have received a copy of the GNU Lesser General Public
  13. License along with this library; if not, write to the Free Software
  14. Foundation, Inc., 51 Franklin Street, Fifth Floor,
  15. Boston, MA 02110-1335 USA
  16. """
  17. import errno
  18. import select
  19. import socket
  20. import six
  21. import sys
  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. timeout: default socket timeout time. This value is second.
  50. """
  51. global _default_timeout
  52. _default_timeout = timeout
  53. def getdefaulttimeout():
  54. """
  55. Return the global timeout setting(second) to connect.
  56. """
  57. return _default_timeout
  58. def recv(sock, bufsize):
  59. if not sock:
  60. raise WebSocketConnectionClosedException("socket is already closed.")
  61. def _recv():
  62. try:
  63. return sock.recv(bufsize)
  64. except SSLWantReadError:
  65. pass
  66. except socket.error as exc:
  67. error_code = extract_error_code(exc)
  68. if error_code is None:
  69. raise
  70. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  71. raise
  72. r, w, e = select.select((sock, ), (), (), sock.gettimeout())
  73. if r:
  74. return sock.recv(bufsize)
  75. try:
  76. if sock.gettimeout() == 0:
  77. bytes_ = sock.recv(bufsize)
  78. else:
  79. bytes_ = _recv()
  80. except socket.timeout as e:
  81. message = extract_err_message(e)
  82. raise WebSocketTimeoutException(message)
  83. except SSLError as e:
  84. message = extract_err_message(e)
  85. if isinstance(message, str) and 'timed out' in message:
  86. raise WebSocketTimeoutException(message)
  87. else:
  88. raise
  89. if not bytes_:
  90. raise WebSocketConnectionClosedException(
  91. "Connection is already closed.")
  92. return bytes_
  93. def recv_line(sock):
  94. line = []
  95. while True:
  96. c = recv(sock, 1)
  97. line.append(c)
  98. if c == six.b("\n"):
  99. break
  100. return six.b("").join(line)
  101. def send(sock, data):
  102. if isinstance(data, six.text_type):
  103. data = data.encode('utf-8')
  104. if not sock:
  105. raise WebSocketConnectionClosedException("socket is already closed.")
  106. def _send():
  107. try:
  108. return sock.send(data)
  109. except SSLWantWriteError:
  110. pass
  111. except socket.error as exc:
  112. error_code = extract_error_code(exc)
  113. if error_code is None:
  114. raise
  115. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  116. raise
  117. r, w, e = select.select((), (sock, ), (), sock.gettimeout())
  118. if w:
  119. return sock.send(data)
  120. try:
  121. if sock.gettimeout() == 0:
  122. return sock.send(data)
  123. else:
  124. return _send()
  125. except socket.timeout as e:
  126. message = extract_err_message(e)
  127. raise WebSocketTimeoutException(message)
  128. except Exception as e:
  129. message = extract_err_message(e)
  130. if isinstance(message, str) and "timed out" in message:
  131. raise WebSocketTimeoutException(message)
  132. else:
  133. raise