Adds zlib StreamReader wrapper for zlib encoding (6)

asyncvnc.py

@@ -1,4 +1,4 @@
-from asyncio import StreamReader, StreamWriter, open_connection
+from asyncio import StreamReader, StreamWriter, open_connection, IncompleteReadError
 from contextlib import asynccontextmanager, contextmanager, ExitStack
 from dataclasses import dataclass, field
 from enum import Enum
@@ -259,6 +259,58 @@ def score(self) -> float:
         return value
 
 
+@dataclass
+class StreamZReader:
+    """
+    aio StreamReader wrapper for zlib
+    """
+    reader: StreamReader = field(repr=False)
+    decompress: object = field(repr=False)
+    length: int
+    _head: int = 0
+    _buffer: bytes = b''
+    _zbuffer: bytes = b''
+
+    async def read(self, n: int =-1) -> bytes:
+
+        if n == -1:
+            try:
+                self._zbuffer += await self.reader.readexactly(self.length)
+            except IncompleteReadError as e:
+                self._zbuffer += e.partial
+            self.length = 0
+            rdata = self._buffer[self._head:] + self.decompress.decompress(self._zbuffer)
+            self._buffer = b''
+            self._head = 0
+            self._zbuffer = self.decompress.unconsumed_tail
+            return rdata
+
+        if (n <= 0) or ((len(self._buffer) <= self._head) and (self.length <= 0)):
+            return b''
+
+        while (len(self._buffer) == self._head) and (self.length > 0):
+            ndata = await self.reader.readexactly(self.length)
+            self.length -= len(ndata)
+            self._zbuffer += ndata
+            self._buffer = self.decompress.decompress(self._zbuffer)
+            self._head = 0
+            self._zbuffer = self.decompress.unconsumed_tail
+        rdata = self._buffer[self._head:self._head + n]
+        self._head += len(rdata)
+        return rdata
+
+    async def readexactly(self, n: int) -> bytes:
+        data = b''
+        _n = n
+        while _n > 0:
+            ndata = await self.read(_n)
+            if len(ndata) == 0:
+                raise IncompleteReadError(data, n)
+            data += ndata
+            _n = n - len(data)
+        return data
+
+
 @dataclass
 class Video:
     """
@@ -267,7 +319,7 @@ class Video:
 
     reader: StreamReader = field(repr=False)
     writer: StreamWriter = field(repr=False)
-    decompress: Callable[[bytes], bytes] = field(repr=False)
+    decompress: object = field(repr=False)
 
     #: Desktop name.
     name: str
@@ -301,7 +353,7 @@ async def create(cls, reader: StreamReader, writer: StreamWriter) -> 'Video':
             writer.write(b'\x00\x00\x00\x00\x20\x18\x00\x01\x00\xff'
                          b'\x00\xff\x00\xff\x00\x08\x10\x00\x00\x00')
         writer.write(b'\x02\x00\x00\x01\x00\x00\x00\x06')
-        decompress = decompressobj().decompress
+        decompress = decompressobj()
         return cls(reader, writer, decompress, name, width, height, mode)
 
     def refresh(self, x: int = 0, y: int = 0, width: Optional[int] = None, height: Optional[int] = None):
@@ -329,12 +381,14 @@ async def read(self):
         height = await read_int(self.reader, 2)
         encoding = await read_int(self.reader, 4)
 
-        if encoding == 0:  # Raw
-            data = await self.reader.readexactly(height * width * 4)
-        elif encoding == 6:  # ZLib
+        if (encoding == 6) or (encoding == 16):
             length = await read_int(self.reader, 4)
-            data = await self.reader.readexactly(length)
-            data = self.decompress(data)
+            _reader = StreamZReader(self.reader, self.decompress, length)
+        else:
+            _reader = self.reader
+
+        if (encoding == 0) or (encoding == 6):  # Raw
+            data = await _reader.readexactly(height * width * 4)
         else:
             raise ValueError(encoding)
 

tests/test_zlib.py

@@ -0,0 +1,57 @@
+from io import BytesIO
+from asyncio import StreamReader, run, IncompleteReadError
+from asyncvnc import StreamZReader
+from zlib import compress, decompressobj
+import pytest
+
+
+@pytest.fixture
+def textandreader():
+    text = b"""
+Nv5InqATV5iLkFt5CVGEDS0XqhmStqZL
+oLYBnj5CklOUoFZGzR6OgyP7swr57MWK
+Ffm8ZyrMhf52NdYde8Y6P8tOk64hpGoA
+o15WdpBg75YiSgWNy6f3EKWqnPzgSvFU
+bVyfs2xoT5yUD1GLnX4JW3r8cM9O4JCM
+tUZnva9xjLJwrxCvzWN87V78EQVHY8qW
+vzXGdXb92w0dJZ7E0f6jQ2HzcBGfzzkt
+haBB7pKB4IOSH5V9AQ7D6wRA3zK4Q7BE
+da8lQnGPniWz3ROWqvbc4eHMI9K4ivmC
+OzN5edG3iOW57xS6c9mJhQYzpWtT8SvO
+"""
+    ztext = compress(text)
+    reader = StreamReader()
+    reader._buffer.extend(bytearray(ztext))
+    zreader = StreamZReader(reader, decompressobj(), len(ztext))
+    return (text, zreader)
+
+@pytest.mark.asyncio
+async def test_readall(textandreader):
+    text, zreader = textandreader
+    assert text  == await zreader.read()
+
+@pytest.mark.asyncio
+async def test_readexactly5(textandreader):
+    text, zreader = textandreader
+    result = b''
+    if len(text) % 5:
+        result += await zreader.readexactly(len(text) % 5)
+    try:
+        while True:
+            result += await zreader.readexactly(5)
+    except IncompleteReadError as e:
+        if len(e.partial):
+            result += b"ErrorErrorError"
+    assert text == result
+
+@pytest.mark.asyncio
+async def test_readexactlyexpect(textandreader):
+    text, zreader = textandreader
+    result = b''
+    try:
+        while True:
+            result += await zreader.readexactly(len(text)+10)
+            result += b"ErrorErrorError"
+    except IncompleteReadError as e:
+        result += e.partial
+    assert text == result