diff --git a/lib/nntp/tiny/buffer.py b/lib/nntp/tiny/buffer.py index 020dc21..342316c 100644 --- a/lib/nntp/tiny/buffer.py +++ b/lib/nntp/tiny/buffer.py @@ -5,35 +5,35 @@ class BufferOverflow(Exception): super().__init__("Buffer overflow") class LineBuffer(): - __slots__ = 'buf', 'size', 'offset', 'count', 'eof', 'done', + __slots__ = 'buf', 'size', 'offset_i', 'offset_o', 'eof', 'done', BUFFER_SIZE = 4096 def __init__(self, size: int=BUFFER_SIZE): - self.buf: bytearray = bytearray(size) - self.size: int = size - self.offset: int = 0 - self.count: int = 0 - self.eof: bool = False - self.done: bool = False + self.buf: bytearray = bytearray(size) + self.size: int = size + self.offset_i: int = 0 + self.offset_o: int = 0 + self.eof: bool = False + self.done: bool = False def _shift(self): - count = self.size - self.offset + count = self.offset_i - self.offset_o - self.buf[0:count] = self.buf[self.offset:self.count] - self.offset = 0 - self.count = count + self.buf[0:count] = self.buf[self.offset_o:self.offset_i] + + self.offset_i = count + self.offset_o = 0 def _is_full(self): - return self.count == self.size + return self.offset_i >= self.size def _fill(self, sock: socket.socket): if self.eof: return - readlen = self.size - self.count - data = sock.recv(readlen) + data = sock.recv(self.size - self.offset_i) if data == b'': self.eof = True @@ -41,28 +41,40 @@ class LineBuffer(): readlen = len(data) - self.buf[self.offset:self.offset+readlen] = data - self.count += readlen + self.buf[self.offset_i:self.size] = data - def _flush(self, end: int): - ret = self.buf[self.offset:end+1] + self.offset_i += readlen - self.offset = end + 1 + def _drain(self, end=None): + if end is None: + end = self.offset_i + 1 + + ret = self.buf[self.offset_o:end+1] + + self.offset_o = end + 1 return str(ret, 'ascii') def readline(self, sock: socket.socket) -> str: + if self.done: + return '' + while True: - index = self.buf.find(b'\n', self.offset, self.count) + index = self.buf.find(b'\n', self.offset_o, self.offset_i) + if index < 0: if self._is_full(): - if self.offset == 0: + if self.offset_o == 0: raise BufferOverflow() - else: - self._shift() + + self._shift() + elif self.eof: + self.done = True + + return self._drain() else: self._fill(sock) else: - return self._flush(index) + return self._drain(index)