_socket.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. """
  2. """
  3. """
  4. _socket.py
  5. websocket - WebSocket client library for Python
  6. Copyright 2021 engn33r
  7. Licensed under the Apache License, Version 2.0 (the "License");
  8. you may not use this file except in compliance with the License.
  9. You may obtain a copy of the License at
  10. http://www.apache.org/licenses/LICENSE-2.0
  11. Unless required by applicable law or agreed to in writing, software
  12. distributed under the License is distributed on an "AS IS" BASIS,
  13. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. See the License for the specific language governing permissions and
  15. limitations under the License.
  16. """
  17. import errno
  18. import selectors
  19. import socket
  20. from ._exceptions import *
  21. from ._ssl_compat import *
  22. from ._utils import *
  23. DEFAULT_SOCKET_OPTION = [(socket.SOL_TCP, socket.TCP_NODELAY, 1)]
  24. #if hasattr(socket, "SO_KEEPALIVE"):
  25. # DEFAULT_SOCKET_OPTION.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1))
  26. #if hasattr(socket, "TCP_KEEPIDLE"):
  27. # DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPIDLE, 30))
  28. #if hasattr(socket, "TCP_KEEPINTVL"):
  29. # DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPINTVL, 10))
  30. #if hasattr(socket, "TCP_KEEPCNT"):
  31. # DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPCNT, 3))
  32. _default_timeout = None
  33. __all__ = ["DEFAULT_SOCKET_OPTION", "sock_opt", "setdefaulttimeout", "getdefaulttimeout",
  34. "recv", "recv_line", "send"]
  35. class sock_opt:
  36. def __init__(self, sockopt, sslopt):
  37. if sockopt is None:
  38. sockopt = []
  39. if sslopt is None:
  40. sslopt = {}
  41. self.sockopt = sockopt
  42. self.sslopt = sslopt
  43. self.timeout = None
  44. def setdefaulttimeout(timeout):
  45. """
  46. Set the global timeout setting to connect.
  47. Parameters
  48. ----------
  49. timeout: int or float
  50. default socket timeout time (in seconds)
  51. """
  52. global _default_timeout
  53. _default_timeout = timeout
  54. def getdefaulttimeout():
  55. """
  56. Get default timeout
  57. Returns
  58. ----------
  59. _default_timeout: int or float
  60. Return the global timeout setting (in seconds) to connect.
  61. """
  62. return _default_timeout
  63. def recv(sock, bufsize):
  64. if not sock:
  65. raise WebSocketConnectionClosedException("socket is already closed.")
  66. def _recv():
  67. try:
  68. return sock.recv(bufsize)
  69. except SSLWantReadError:
  70. pass
  71. except socket.error as exc:
  72. error_code = extract_error_code(exc)
  73. if error_code is None:
  74. raise
  75. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  76. raise
  77. sel = selectors.DefaultSelector()
  78. sel.register(sock, selectors.EVENT_READ)
  79. r = sel.select(sock.gettimeout())
  80. sel.close()
  81. if r:
  82. return sock.recv(bufsize)
  83. try:
  84. if sock.gettimeout() == 0:
  85. bytes_ = sock.recv(bufsize)
  86. else:
  87. bytes_ = _recv()
  88. except socket.timeout as e:
  89. message = extract_err_message(e)
  90. raise WebSocketTimeoutException(message)
  91. except SSLError as e:
  92. message = extract_err_message(e)
  93. if isinstance(message, str) and 'timed out' in message:
  94. raise WebSocketTimeoutException(message)
  95. else:
  96. raise
  97. if not bytes_:
  98. raise WebSocketConnectionClosedException(
  99. "Connection to remote host was lost.")
  100. return bytes_
  101. def recv_line(sock):
  102. line = []
  103. while True:
  104. c = recv(sock, 1)
  105. line.append(c)
  106. if c == b'\n':
  107. break
  108. return b''.join(line)
  109. def send(sock, data):
  110. if isinstance(data, str):
  111. data = data.encode('utf-8')
  112. if not sock:
  113. raise WebSocketConnectionClosedException("socket is already closed.")
  114. def _send():
  115. try:
  116. return sock.send(data)
  117. except SSLWantWriteError:
  118. pass
  119. except socket.error as exc:
  120. error_code = extract_error_code(exc)
  121. if error_code is None:
  122. raise
  123. if error_code != errno.EAGAIN or error_code != errno.EWOULDBLOCK:
  124. raise
  125. sel = selectors.DefaultSelector()
  126. sel.register(sock, selectors.EVENT_WRITE)
  127. w = sel.select(sock.gettimeout())
  128. sel.close()
  129. if w:
  130. return sock.send(data)
  131. try:
  132. if sock.gettimeout() == 0:
  133. return sock.send(data)
  134. else:
  135. return _send()
  136. except socket.timeout as e:
  137. message = extract_err_message(e)
  138. raise WebSocketTimeoutException(message)
  139. except Exception as e:
  140. message = extract_err_message(e)
  141. if isinstance(message, str) and "timed out" in message:
  142. raise WebSocketTimeoutException(message)
  143. else:
  144. raise