websocket.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809
  1. import base64
  2. import hashlib
  3. import os
  4. import ssl
  5. import threading
  6. import time
  7. from datetime import datetime
  8. from enum import IntEnum
  9. from typing import Any, Callable, Dict, List, Optional, Union
  10. from urllib.parse import parse_qs, urlencode, urlparse
  11. import websocket
  12. from websocket._abnf import ABNF
  13. from darabonba.request import DaraRequest
  14. from darabonba.response import DaraResponse
  15. from darabonba.runtime import RuntimeOptions
  16. class WebSocketMessageType(IntEnum):
  17. Text = 0
  18. Binary = 1
  19. Ping = 2
  20. Pong = 3
  21. Close = 4
  22. class WebSocketMessage:
  23. def __init__(
  24. self,
  25. type: WebSocketMessageType,
  26. payload: bytes,
  27. headers: Optional[Dict[str, str]] = None,
  28. timestamp: Optional[datetime] = None,
  29. ):
  30. self.type = type
  31. self.payload = payload
  32. self.headers = headers or {}
  33. self.timestamp = timestamp or datetime.now()
  34. class WebSocketSessionInfo:
  35. def __init__(
  36. self,
  37. session_id: str,
  38. request_id: str = '',
  39. connected_at: Optional[datetime] = None,
  40. remote_addr: str = '',
  41. local_addr: str = '',
  42. attributes: Optional[Dict[str, Any]] = None,
  43. ):
  44. self.session_id = session_id
  45. self.request_id = request_id
  46. self.connected_at = connected_at or datetime.now()
  47. self.remote_addr = remote_addr
  48. self.local_addr = local_addr
  49. self.attributes = attributes or {}
  50. STATE_DISCONNECTED = 0
  51. STATE_CONNECTING = 1
  52. STATE_CONNECTED = 2
  53. STATE_DISCONNECTING = 3
  54. def _get_runtime_value(runtime: Any, camel_key: str, snake_key: str = None):
  55. if runtime is None:
  56. return None
  57. if isinstance(runtime, dict):
  58. return runtime.get(camel_key)
  59. key = snake_key or camel_key
  60. return getattr(runtime, key, None)
  61. def get_web_socket_ping_interval(runtime: Any) -> Optional[int]:
  62. return _get_runtime_value(runtime, 'webSocketPingInterval', 'web_socket_ping_interval')
  63. def get_web_socket_pong_timeout(runtime: Any) -> Optional[int]:
  64. return _get_runtime_value(runtime, 'webSocketPongTimeout', 'web_socket_pong_timeout')
  65. def get_web_socket_enable_reconnect(runtime: Any) -> Optional[bool]:
  66. return _get_runtime_value(runtime, 'webSocketEnableReconnect', 'web_socket_enable_reconnect')
  67. def get_web_socket_reconnect_interval(runtime: Any) -> Optional[int]:
  68. return _get_runtime_value(runtime, 'webSocketReconnectInterval', 'web_socket_reconnect_interval')
  69. def get_web_socket_max_reconnect_times(runtime: Any) -> Optional[int]:
  70. return _get_runtime_value(runtime, 'webSocketMaxReconnectTimes', 'web_socket_max_reconnect_times')
  71. def get_web_socket_write_timeout(runtime: Any) -> Optional[int]:
  72. return _get_runtime_value(runtime, 'webSocketWriteTimeout', 'web_socket_write_timeout')
  73. def get_web_socket_handshake_timeout(runtime: Any) -> Optional[int]:
  74. return _get_runtime_value(runtime, 'webSocketHandshakeTimeout', 'web_socket_handshake_timeout')
  75. def get_web_socket_handler(runtime: Any):
  76. return _get_runtime_value(runtime, 'webSocketHandler', 'web_socket_handler')
  77. def new_websocket_response(
  78. status_code: int,
  79. status_message: str,
  80. headers: Optional[Dict[str, Any]] = None,
  81. ) -> DaraResponse:
  82. res = DaraResponse()
  83. res.status_code = status_code
  84. res.status_message = status_message
  85. normalized_headers: Dict[str, str] = {}
  86. for key, val in (headers or {}).items():
  87. if val is not None:
  88. if isinstance(val, list):
  89. normalized_headers[key.lower()] = str(val[0])
  90. else:
  91. normalized_headers[key.lower()] = str(val)
  92. res.headers = normalized_headers
  93. return res
  94. class WebSocketHandler:
  95. def after_connection_established(self, session: WebSocketSessionInfo) -> None:
  96. return None
  97. def handle_raw_message(self, session: WebSocketSessionInfo, message: WebSocketMessage) -> None:
  98. return None
  99. def handle_error(self, session: WebSocketSessionInfo, err: Exception) -> None:
  100. return None
  101. def after_connection_closed(self, session: WebSocketSessionInfo, code: int, reason: str) -> None:
  102. return None
  103. class AbstractWebSocketHandler(WebSocketHandler):
  104. pass
  105. class DefaultWebSocketClient:
  106. def __init__(self, handler: WebSocketHandler):
  107. if handler is None:
  108. raise ValueError('handler cannot be nil')
  109. self.handler = handler
  110. self.stopped = False
  111. self.pong_received = False
  112. self.state = STATE_DISCONNECTED
  113. self.ws_app: Optional[websocket.WebSocketApp] = None
  114. self.session: Optional[WebSocketSessionInfo] = None
  115. self.request: Optional[DaraRequest] = None
  116. self.runtime_object: Any = None
  117. self.reconnect_count = 0
  118. self._ping_timer: Optional[threading.Timer] = None
  119. self.ping_interval = 30000
  120. self.reconnect_interval = 5000
  121. self.write_timeout = 30000
  122. self.read_timeout = 0
  123. self.pong_timeout = 5000
  124. self.max_reconnect_times = 5
  125. self._reconnect_lock = threading.Lock()
  126. self._abort_event = threading.Event()
  127. self._ws_thread: Optional[threading.Thread] = None
  128. self._message_handlers_started = False
  129. self._connect_error: Optional[Exception] = None
  130. self._handshake_response: Dict[str, Any] = {}
  131. def connect(self, request: DaraRequest, runtime_object: Any) -> DaraResponse:
  132. if request is None:
  133. raise ValueError('request cannot be nil')
  134. if runtime_object is None:
  135. raise ValueError('runtimeObject cannot be nil')
  136. self.request = request
  137. self.runtime_object = runtime_object
  138. self._update_timeout_config(runtime_object)
  139. self.state = STATE_CONNECTING
  140. self.stopped = False
  141. self.pong_received = False
  142. self._connect_error = None
  143. self._handshake_response = {}
  144. self._abort_event.clear()
  145. request_url = build_websocket_url(request)
  146. parsed = urlparse(request_url)
  147. if not parsed.scheme or not parsed.netloc:
  148. self.state = STATE_DISCONNECTED
  149. raise ValueError(f'invalid websocket url: {request_url}')
  150. handshake_timeout = get_web_socket_handshake_timeout(runtime_object)
  151. if not handshake_timeout or handshake_timeout <= 0:
  152. handshake_timeout = 30000
  153. connect_timeout = _get_connect_timeout(runtime_object)
  154. headers: Dict[str, str] = {}
  155. if request.headers:
  156. for key, value in request.headers.items():
  157. if value and key.lower() not in ('host', 'content-length'):
  158. headers[key] = value
  159. sslopt = self._configure_tls(runtime_object, parsed.scheme)
  160. proxy_kwargs = self._configure_proxy(parsed, runtime_object, request)
  161. connected_event = threading.Event()
  162. def on_open(ws_app):
  163. self.ws_app = ws_app
  164. self.state = STATE_CONNECTED
  165. self._setup_pong_handler()
  166. self._start_message_handlers(ws_app)
  167. if self.ping_interval > 0:
  168. self._start_ping_pong()
  169. resp_headers: Dict[str, Any] = {}
  170. status_code = 101
  171. status_message = 'Switching Protocols'
  172. if ws_app.sock is not None:
  173. if hasattr(ws_app.sock, 'headers') and ws_app.sock.headers:
  174. resp_headers = dict(ws_app.sock.headers)
  175. if hasattr(ws_app.sock, 'status') and ws_app.sock.status:
  176. status_code = ws_app.sock.status
  177. if hasattr(ws_app.sock, 'status_message') and ws_app.sock.status_message:
  178. status_message = ws_app.sock.status_message
  179. session_id = (
  180. resp_headers.get('x-acs-ws-session-id')
  181. or resp_headers.get('X-Acs-Ws-Session-Id')
  182. or _generate_session_id()
  183. )
  184. request_id = (
  185. resp_headers.get('x-acs-request-id')
  186. or resp_headers.get('X-Acs-Request-Id')
  187. or ''
  188. )
  189. if isinstance(session_id, list):
  190. session_id = session_id[0]
  191. if isinstance(request_id, list):
  192. request_id = request_id[0]
  193. self.session = WebSocketSessionInfo(
  194. session_id=str(session_id),
  195. request_id=str(request_id),
  196. connected_at=datetime.now(),
  197. remote_addr=request_url,
  198. local_addr='',
  199. attributes={},
  200. )
  201. try:
  202. self.handler.after_connection_established(self.session)
  203. except Exception as err:
  204. self._connect_error = err
  205. self.state = STATE_DISCONNECTED
  206. self.stopped = True
  207. try:
  208. ws_app.close()
  209. except Exception:
  210. pass
  211. connected_event.set()
  212. return
  213. self._handshake_response = {
  214. 'status_code': status_code,
  215. 'status_message': status_message,
  216. 'headers': resp_headers,
  217. }
  218. connected_event.set()
  219. def on_error(ws_app, error):
  220. if self.state == STATE_CONNECTING:
  221. self._connect_error = error if isinstance(error, Exception) else Exception(str(error))
  222. connected_event.set()
  223. self.ws_app = websocket.WebSocketApp(
  224. request_url,
  225. header=headers,
  226. on_open=on_open,
  227. on_message=self._on_message,
  228. on_error=on_error,
  229. on_close=self._on_close,
  230. on_pong=self._on_pong,
  231. )
  232. run_kwargs = {
  233. 'sslopt': sslopt,
  234. 'ping_interval': 0,
  235. 'ping_timeout': None,
  236. }
  237. run_kwargs.update(proxy_kwargs)
  238. self._ws_thread = threading.Thread(
  239. target=self.ws_app.run_forever,
  240. kwargs=run_kwargs,
  241. daemon=True,
  242. )
  243. self._ws_thread.start()
  244. timeout_seconds = max(connect_timeout, handshake_timeout) / 1000
  245. if not connected_event.wait(timeout=timeout_seconds):
  246. self.state = STATE_DISCONNECTED
  247. self._cleanup_connection()
  248. raise TimeoutError('WebSocket connection timeout')
  249. if self._connect_error is not None:
  250. self._cleanup_connection()
  251. raise self._connect_error
  252. return new_websocket_response(
  253. self._handshake_response.get('status_code', 101),
  254. self._handshake_response.get('status_message', 'Switching Protocols'),
  255. self._handshake_response.get('headers', {}),
  256. )
  257. def disconnect(self) -> None:
  258. self._disconnect_internal(1000, 'Normal closure')
  259. def _disconnect_internal(self, code: int, reason: str) -> None:
  260. if self.state == STATE_DISCONNECTED and not self._ws_thread:
  261. return
  262. self.state = STATE_DISCONNECTING
  263. self._stop_ping_pong()
  264. self.stopped = True
  265. ws_app = self.ws_app
  266. self.ws_app = None
  267. if ws_app is not None:
  268. try:
  269. ws_app.close(status=code, reason=reason)
  270. except Exception:
  271. pass
  272. self._wait_for_close()
  273. if self.session:
  274. try:
  275. self.handler.after_connection_closed(self.session, code, reason)
  276. except Exception:
  277. pass
  278. self.state = STATE_DISCONNECTED
  279. self._abort_event.set()
  280. def reconnect(self) -> DaraResponse:
  281. return self._reconnect_internal(False)
  282. def reconnect_gracefully(self) -> DaraResponse:
  283. return self._reconnect_internal(True)
  284. def _reconnect_internal(self, graceful: bool) -> DaraResponse:
  285. if not graceful and self.is_connected():
  286. raise Exception('already connected')
  287. if self._abort_event.is_set():
  288. raise Exception('connection aborted')
  289. if not self._reconnect_lock.acquire(blocking=False):
  290. raise Exception('reconnect already in progress')
  291. try:
  292. if not get_web_socket_enable_reconnect(self.runtime_object):
  293. raise Exception('reconnect is disabled')
  294. if self.reconnect_count >= self.max_reconnect_times:
  295. raise Exception(f'max reconnect times reached: {self.max_reconnect_times}')
  296. previous_session_id = ''
  297. if graceful:
  298. if self.session and self.session.session_id:
  299. previous_session_id = self.session.session_id
  300. else:
  301. raise Exception('graceful reconnection requires existing session ID')
  302. self._cleanup_resources()
  303. self.stopped = False
  304. self.reconnect_count += 1
  305. if self.request is None:
  306. raise Exception('request is nil, cannot reconnect')
  307. if graceful and previous_session_id:
  308. if self.request.headers is None:
  309. self.request.headers = {}
  310. self.request.headers['X-Acs-Ws-Session-Id'] = previous_session_id
  311. elif self.request.headers:
  312. self.request.headers.pop('X-Acs-Ws-Session-Id', None)
  313. self._sleep_interruptible(self.reconnect_interval)
  314. if self._abort_event.is_set():
  315. raise Exception('connection aborted')
  316. result = self.connect(self.request, self.runtime_object)
  317. self.reconnect_count = 0
  318. return result
  319. finally:
  320. self._reconnect_lock.release()
  321. def _cleanup_resources(self) -> None:
  322. self.state = STATE_DISCONNECTING
  323. self._stop_ping_pong()
  324. self.stopped = True
  325. ws_app = self.ws_app
  326. self.ws_app = None
  327. if ws_app is not None:
  328. try:
  329. ws_app.close(status=1001, reason='Reconnecting')
  330. except Exception:
  331. pass
  332. self._wait_for_close()
  333. self.session = None
  334. self._message_handlers_started = False
  335. self.state = STATE_DISCONNECTED
  336. def is_connected(self) -> bool:
  337. return self.state == STATE_CONNECTED
  338. def send_text(self, text: str) -> None:
  339. if not self.is_connected():
  340. raise Exception('not connected')
  341. ws_app = self.ws_app
  342. if ws_app is None:
  343. raise Exception('connection is nil')
  344. self._send_with_timeout(lambda: ws_app.send(text))
  345. def send_binary(self, data: bytes) -> None:
  346. if not self.is_connected():
  347. raise Exception('not connected')
  348. ws_app = self.ws_app
  349. if ws_app is None:
  350. raise Exception('connection is nil')
  351. self._send_with_timeout(lambda: ws_app.send(data, opcode=ABNF.OPCODE_BINARY))
  352. def get_session_info(self) -> Optional[WebSocketSessionInfo]:
  353. return self.session
  354. def close(self) -> None:
  355. self._disconnect_internal(1000, 'Client closed')
  356. def configure_tls(self, runtime_object: Any, scheme: str = 'wss') -> Dict[str, Any]:
  357. return self._configure_tls(runtime_object, scheme)
  358. def configure_proxy(
  359. self,
  360. parsed: Any,
  361. runtime_object: Any,
  362. request: DaraRequest,
  363. ) -> Dict[str, Any]:
  364. return self._configure_proxy(parsed, runtime_object, request)
  365. def configure_socks5_proxy(self, runtime_object: Any) -> Dict[str, Any]:
  366. proxy_kwargs: Dict[str, Any] = {}
  367. socks5_proxy = _get_runtime_value(runtime_object, 'socks5Proxy', 'socks_5proxy')
  368. if socks5_proxy:
  369. proxy_kwargs['proxy_type'] = 'socks5'
  370. proxy_kwargs['http_proxy_host'] = socks5_proxy
  371. return proxy_kwargs
  372. def configure_http_proxy(
  373. self,
  374. parsed: Any,
  375. runtime_object: Any,
  376. request: DaraRequest,
  377. ) -> Dict[str, Any]:
  378. return self._configure_proxy(parsed, runtime_object, request)
  379. def _configure_tls(self, runtime_object: Any, scheme: str) -> Dict[str, Any]:
  380. if scheme not in ('wss', 'https'):
  381. return {}
  382. ignore_ssl = _get_runtime_value(runtime_object, 'ignoreSSL', 'ignore_ssl')
  383. sslopt: Dict[str, Any] = {}
  384. if ignore_ssl:
  385. sslopt['cert_reqs'] = ssl.CERT_NONE
  386. sslopt['check_hostname'] = False
  387. else:
  388. sslopt['cert_reqs'] = ssl.CERT_REQUIRED
  389. cert = _get_runtime_value(runtime_object, 'cert', 'cert')
  390. key = _get_runtime_value(runtime_object, 'key', 'key')
  391. ca = _get_runtime_value(runtime_object, 'ca', 'ca')
  392. if cert and key:
  393. sslopt['certfile'] = cert
  394. sslopt['keyfile'] = key
  395. if ca:
  396. sslopt['ca_certs'] = ca
  397. return sslopt
  398. def _configure_proxy(
  399. self,
  400. parsed: Any,
  401. runtime_object: Any,
  402. request: DaraRequest,
  403. ) -> Dict[str, Any]:
  404. proxy_kwargs: Dict[str, Any] = {}
  405. socks5_proxy = _get_runtime_value(runtime_object, 'socks5Proxy', 'socks_5proxy')
  406. if socks5_proxy:
  407. proxy_kwargs['proxy_type'] = 'socks5'
  408. proxy_kwargs['http_proxy_host'] = socks5_proxy
  409. return proxy_kwargs
  410. if parsed is None:
  411. return proxy_kwargs
  412. protocol = parsed.scheme.replace(':', '')
  413. host = parsed.hostname
  414. no_proxy_list = _get_no_proxy(protocol, runtime_object)
  415. for no_proxy_host in no_proxy_list:
  416. if no_proxy_host.strip() == host:
  417. return proxy_kwargs
  418. proxy_url = _get_http_proxy_url(protocol, host, runtime_object)
  419. if not proxy_url:
  420. return proxy_kwargs
  421. proxy_parsed = urlparse(proxy_url)
  422. proxy_kwargs['http_proxy_host'] = proxy_parsed.hostname
  423. proxy_kwargs['http_proxy_port'] = proxy_parsed.port or (
  424. 443 if protocol in ('wss', 'https') else 80
  425. )
  426. if proxy_parsed.username and proxy_parsed.password:
  427. if request.headers is None:
  428. request.headers = {}
  429. auth = f'{proxy_parsed.username}:{proxy_parsed.password}'
  430. request.headers['Proxy-Authorization'] = (
  431. 'Basic ' + base64.b64encode(auth.encode()).decode()
  432. )
  433. return proxy_kwargs
  434. def _update_timeout_config(self, runtime_object: Any) -> None:
  435. ping_interval = get_web_socket_ping_interval(runtime_object)
  436. self.ping_interval = ping_interval if ping_interval and ping_interval > 0 else 30000
  437. reconnect_interval = get_web_socket_reconnect_interval(runtime_object)
  438. self.reconnect_interval = reconnect_interval if reconnect_interval and reconnect_interval > 0 else 5000
  439. write_timeout = get_web_socket_write_timeout(runtime_object)
  440. self.write_timeout = write_timeout if write_timeout and write_timeout > 0 else 30000
  441. read_timeout = _get_runtime_value(runtime_object, 'readTimeout', 'read_timeout')
  442. self.read_timeout = read_timeout if read_timeout and read_timeout > 0 else 0
  443. pong_timeout = get_web_socket_pong_timeout(runtime_object)
  444. self.pong_timeout = pong_timeout if pong_timeout and pong_timeout > 0 else 5000
  445. max_reconnect_times = get_web_socket_max_reconnect_times(runtime_object)
  446. self.max_reconnect_times = max_reconnect_times if max_reconnect_times and max_reconnect_times > 0 else 5
  447. def _setup_pong_handler(self) -> None:
  448. self.pong_received = False
  449. def _on_pong(self, ws_app, _data) -> None:
  450. self.pong_received = True
  451. def _start_message_handlers(self, ws_app) -> None:
  452. if self._message_handlers_started:
  453. return
  454. self._message_handlers_started = True
  455. self.ws_app = ws_app
  456. def _on_message(self, ws_app, message) -> None:
  457. if self.stopped:
  458. return
  459. if isinstance(message, bytes):
  460. msg_type = WebSocketMessageType.Binary
  461. payload = message
  462. else:
  463. msg_type = WebSocketMessageType.Text
  464. payload = message.encode('utf-8') if isinstance(message, str) else bytes(message)
  465. msg = WebSocketMessage(type=msg_type, payload=payload, headers={}, timestamp=datetime.now())
  466. if not self.session:
  467. return
  468. try:
  469. self.handler.handle_raw_message(self.session, msg)
  470. except Exception as err:
  471. try:
  472. self.handler.handle_error(self.session, err)
  473. except Exception:
  474. pass
  475. def _on_error(self, ws_app, error) -> None:
  476. if self.stopped or not self.session:
  477. return
  478. err = error if isinstance(error, Exception) else Exception(str(error))
  479. try:
  480. self.handler.handle_error(self.session, err)
  481. except Exception:
  482. pass
  483. def _on_close(self, ws_app, close_status_code, close_msg) -> None:
  484. if self.stopped:
  485. return
  486. reason = close_msg or ''
  487. if self.session and close_status_code not in (1000, 1001):
  488. try:
  489. self.handler.handle_error(
  490. self.session,
  491. Exception(f'WebSocket closed: {close_status_code} {reason}'),
  492. )
  493. except Exception:
  494. pass
  495. if get_web_socket_enable_reconnect(self.runtime_object) and self.request and not self.stopped:
  496. try:
  497. self.reconnect()
  498. except Exception:
  499. pass
  500. def _start_ping_pong(self) -> None:
  501. if self.ping_interval <= 0:
  502. return
  503. self._stop_ping_pong()
  504. self._schedule_ping()
  505. def _schedule_ping(self) -> None:
  506. if self.stopped or not self.is_connected():
  507. return
  508. self.pong_received = False
  509. try:
  510. if self.ws_app and self.ws_app.sock:
  511. self.ws_app.sock.ping()
  512. except Exception as err:
  513. if self.session:
  514. try:
  515. self.handler.handle_error(self.session, err)
  516. except Exception:
  517. pass
  518. return
  519. def check_pong():
  520. if self.stopped:
  521. return
  522. if not self.pong_received and get_web_socket_enable_reconnect(self.runtime_object):
  523. try:
  524. self.reconnect()
  525. except Exception:
  526. pass
  527. if not self.stopped and self.is_connected():
  528. self._schedule_ping()
  529. self._ping_timer = threading.Timer(self.pong_timeout / 1000, check_pong)
  530. self._ping_timer.daemon = True
  531. self._ping_timer.start()
  532. def _stop_ping_pong(self) -> None:
  533. if self._ping_timer is not None:
  534. self._ping_timer.cancel()
  535. self._ping_timer = None
  536. def _cleanup_connection(self) -> None:
  537. self._stop_ping_pong()
  538. self.stopped = True
  539. ws_app = self.ws_app
  540. self.ws_app = None
  541. if ws_app is not None:
  542. try:
  543. ws_app.close()
  544. except Exception:
  545. pass
  546. self._wait_for_close()
  547. def _send_with_timeout(self, send_fn: Callable[[], None]) -> None:
  548. if self.write_timeout <= 0:
  549. send_fn()
  550. return
  551. error: List[Optional[Exception]] = [None]
  552. def do_send():
  553. try:
  554. send_fn()
  555. except Exception as err:
  556. error[0] = err
  557. thread = threading.Thread(target=do_send, daemon=True)
  558. thread.start()
  559. thread.join(timeout=self.write_timeout / 1000)
  560. if thread.is_alive():
  561. raise TimeoutError('write timeout')
  562. if error[0] is not None:
  563. raise error[0]
  564. def _wait_for_close(self, timeout: float = 5.0) -> None:
  565. thread = self._ws_thread
  566. if thread is None:
  567. return
  568. if thread is threading.current_thread():
  569. return
  570. thread.join(timeout=timeout)
  571. def _sleep_interruptible(self, ms: int) -> bool:
  572. return self._abort_event.wait(timeout=ms / 1000)
  573. def new_default_websocket_client(handler: WebSocketHandler) -> DefaultWebSocketClient:
  574. return DefaultWebSocketClient(handler)
  575. def new_websocket_client_and_connect(
  576. request: DaraRequest,
  577. runtime_object: Any,
  578. ):
  579. if runtime_object is None:
  580. raise ValueError('runtimeObject cannot be nil')
  581. handler = get_web_socket_handler(runtime_object)
  582. if handler is None:
  583. raise ValueError(
  584. 'WebSocketHandler is required: please set it in runtimeObject.webSocketHandler'
  585. )
  586. client = DefaultWebSocketClient(handler)
  587. response = client.connect(request, runtime_object)
  588. return client, response
  589. def build_websocket_url(request: DaraRequest) -> str:
  590. if request is None:
  591. raise ValueError('request cannot be nil')
  592. protocol = request.protocol or 'ws'
  593. protocol = protocol.lower()
  594. if protocol == 'http':
  595. protocol = 'ws'
  596. elif protocol == 'https':
  597. protocol = 'wss'
  598. domain = getattr(request, 'domain', None)
  599. if not domain and request.headers:
  600. domain = request.headers.get('host')
  601. if not domain:
  602. raise ValueError('domain is required (set in request.headers["host"] or request.domain)')
  603. request_url = f'{protocol}://{domain}'
  604. request_url += request.pathname or '/'
  605. if request.query:
  606. qs = urlencode(request.query, doseq=True)
  607. if qs:
  608. request_url += '&' + qs if '?' in request_url else '?' + qs
  609. return request_url
  610. def convert_to_websocket_message_type(message_type: int) -> WebSocketMessageType:
  611. mapping = {
  612. 1: WebSocketMessageType.Text,
  613. 2: WebSocketMessageType.Binary,
  614. 9: WebSocketMessageType.Ping,
  615. 10: WebSocketMessageType.Pong,
  616. 8: WebSocketMessageType.Close,
  617. }
  618. return mapping.get(message_type, WebSocketMessageType.Binary)
  619. def _generate_session_id() -> str:
  620. return f'ws-session-{int(time.time() * 1000)}{os.urandom(4).hex()}'
  621. def _get_connect_timeout(runtime_object: Any) -> int:
  622. connect_timeout = _get_runtime_value(runtime_object, 'connectTimeout', 'connect_timeout')
  623. if connect_timeout and connect_timeout > 0:
  624. return connect_timeout
  625. return 10000
  626. def _get_no_proxy(protocol: str, runtime: Any) -> List[str]:
  627. no_proxy = _get_runtime_value(runtime, 'noProxy', 'no_proxy')
  628. if no_proxy:
  629. return no_proxy.split(',')
  630. env_no_proxy = os.environ.get('NO_PROXY') or os.environ.get('no_proxy')
  631. if env_no_proxy:
  632. return env_no_proxy.split(',')
  633. return []
  634. def _get_http_proxy_url(protocol: str, host: str, runtime: Any) -> Optional[str]:
  635. no_proxy_list = _get_no_proxy(protocol, runtime)
  636. for no_proxy_host in no_proxy_list:
  637. if no_proxy_host.strip() == host:
  638. return None
  639. proxy_protocol = protocol
  640. if protocol == 'wss':
  641. proxy_protocol = 'https'
  642. elif protocol == 'ws':
  643. proxy_protocol = 'http'
  644. if proxy_protocol == 'https':
  645. proxy_url = _get_runtime_value(runtime, 'httpsProxy', 'https_proxy')
  646. if proxy_url:
  647. return proxy_url
  648. return os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
  649. proxy_url = _get_runtime_value(runtime, 'httpProxy', 'http_proxy')
  650. if proxy_url:
  651. return proxy_url
  652. return os.environ.get('HTTP_PROXY') or os.environ.get('http_proxy')