diff --git a/curl_cffi/curl.py b/curl_cffi/curl.py index 2ff7889..9ebca8c 100644 --- a/curl_cffi/curl.py +++ b/curl_cffi/curl.py @@ -234,11 +234,13 @@ def getinfo(self, option: CurlInfo) -> Union[bytes, int, float, List]: 0x200000: "long*", 0x300000: "double*", 0x400000: "struct curl_slist **", + 0x500000: "long*" } ret_cast_option = { 0x100000: ffi.string, 0x200000: int, 0x300000: float, + 0x500000: int, } c_value = ffi.new(ret_option[option & 0xF00000]) ret = lib.curl_easy_getinfo(self._curl, option, c_value) diff --git a/curl_cffi/requests/websockets.py b/curl_cffi/requests/websockets.py index 8a00b3e..1ddbf2e 100644 --- a/curl_cffi/requests/websockets.py +++ b/curl_cffi/requests/websockets.py @@ -1,9 +1,12 @@ import asyncio +import select import struct from enum import IntEnum from typing import Callable, Optional, Tuple -from ..const import CurlECode, CurlWsFlag +from .._wrapper import ffi, lib +from ..aio import CURL_SOCKET_BAD +from ..const import CurlECode, CurlWsFlag, CurlInfo from ..curl import CurlError ON_MESSAGE_T = Callable[["WebSocket", bytes], None] @@ -64,14 +67,20 @@ def recv(self) -> Tuple[bytes, int]: """ chunks = [] flags = 0 - # TODO use select here + sock_fd = self.curl.getinfo(CurlInfo.ACTIVESOCKET) + if sock_fd == CURL_SOCKET_BAD: + raise WebSocketError( + "Invalid active socket", CurlECode.NO_CONNECTION_AVAILABLE + ) while True: try: - chunk, frame = self.curl.ws_recv() - flags = frame.flags - chunks.append(chunk) - if frame.bytesleft == 0 and flags & CurlWsFlag.CONT == 0: - break + rlist, _, _ = select.select([sock_fd], [], [], 5.0) + if rlist: + chunk, frame = self.curl.ws_recv() + flags = frame.flags + chunks.append(chunk) + if frame.bytesleft == 0 and flags & CurlWsFlag.CONT == 0: + break except CurlError as e: if e.code == CurlECode.AGAIN: pass