chore: generic fixes (#1187)

proxy/bridge.py

@@ -202,6 +202,7 @@ async def _cfproxy_worker_fallback(reader, writer, relay_init, label,
             try:
                 ws = await RawWebSocket.connect(worker_domain, worker_domain,
                                                 timeout=10.0, path=path)
+                break
             except Exception as exc:
                 cf_worker_pool.report_failure(worker_domain, exc)
                 log.warning("[%s] DC%d%s CF worker %s failed: %s",

proxy/config.py

@@ -209,11 +209,11 @@ def parse_dc_ip_list(dc_ip_list: List[str]) -> Dict[int, str]:
         dc_s, ip_s = entry.split(':', 1)
         try:
             dc_n = int(dc_s)
-            _socket.inet_aton(ip_s)
+            _socket.inet_pton(_socket.AF_INET, ip_s)
         except (ValueError, OSError):
             err = ValueError(f"Invalid --dc-ip {entry!r}")
             err.entry = entry
             err.kind = "invalid"
-            raise err
+            raise err from None
         dc_redirects[dc_n] = ip_s
     return dc_redirects

proxy/pool.py

@@ -296,7 +296,7 @@ class _CfWorkerPool:
 
     def available_domains(self, worker_domains: List[str]) -> List[str]:
         now = time.time()
-        domains = list()
+        domains = []
         for domain in worker_domains:
             if domain in domains:
                 continue

proxy/raw_websocket.py

@@ -67,18 +67,22 @@ def set_sock_opts(transport, buffer_size):
 
 
 class RawWebSocket:
-    __slots__ = ('reader', 'writer', '_closed')
+    __slots__ = ('reader', 'writer', '_closed', '_frag')
 
+    OP_CONT = 0x0
     OP_BINARY = 0x2
     OP_CLOSE = 0x8
     OP_PING = 0x9
     OP_PONG = 0xA
 
+    MAX_MESSAGE_LEN = 16 * 1024 * 1024
+
     def __init__(self, reader: asyncio.StreamReader,
                  writer: asyncio.StreamWriter):
         self.reader = reader
         self.writer = writer
         self._closed = False
+        self._frag = bytearray()
 
     @staticmethod
     async def connect(host: str, domain: str, timeout: float = 10.0,
@@ -164,7 +168,7 @@ class RawWebSocket:
 
     async def recv(self) -> Optional[bytes]:
         while not self._closed:
-            opcode, payload = await self._read_frame()
+            opcode, payload, fin = await self._read_frame()
 
             if opcode == self.OP_CLOSE:
                 self._closed = True
@@ -192,8 +196,18 @@ class RawWebSocket:
             if opcode == self.OP_PONG:
                 continue
 
-            if opcode in (0x1, 0x2):
-                return payload
+            if opcode in (self.OP_CONT, 0x1, self.OP_BINARY):
+                if fin and not self._frag:
+                    return payload
+                self._frag.extend(payload)
+                if len(self._frag) > self.MAX_MESSAGE_LEN:
+                    raise ConnectionError(
+                        f"WS message too large: {len(self._frag)} bytes")
+                if not fin:
+                    continue
+                message = bytes(self._frag)
+                self._frag.clear()
+                return message
             continue
         return None
 
@@ -251,17 +265,20 @@ class RawWebSocket:
             return _st_BBH4s.pack(fb, 0x80 | 126, length, mask_key) + masked
         return _st_BBQ4s.pack(fb, 0x80 | 127, length, mask_key) + masked
 
-    async def _read_frame(self) -> Tuple[int, bytes]:
+    async def _read_frame(self) -> Tuple[int, bytes, bool]:
         hdr = await self.reader.readexactly(2)
+        fin = bool(hdr[0] & 0x80)
         opcode = hdr[0] & 0x0F
         length = hdr[1] & 0x7F
         if length == 126:
             length = _st_H.unpack(await self.reader.readexactly(2))[0]
         elif length == 127:
             length = _st_Q.unpack(await self.reader.readexactly(8))[0]
+        if length > self.MAX_MESSAGE_LEN:
+            raise ConnectionError(f"WS frame too large: {length} bytes")
         if hdr[1] & 0x80:
             mask_key = await self.reader.readexactly(4)
             payload = await self.reader.readexactly(length)
-            return opcode, _xor_mask(payload, mask_key)
+            return opcode, _xor_mask(payload, mask_key), fin
         payload = await self.reader.readexactly(length)
-        return opcode, payload
+        return opcode, payload, fin