diff --git a/lib/nntp/tiny/buffer.py b/lib/nntp/tiny/buffer.py index 499f8d9..020dc21 100644 --- a/lib/nntp/tiny/buffer.py +++ b/lib/nntp/tiny/buffer.py @@ -5,66 +5,64 @@ class BufferOverflow(Exception): super().__init__("Buffer overflow") class LineBuffer(): - __slots__ = 'buf', 'offset', + __slots__ = 'buf', 'size', 'offset', 'count', '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 - def readline(self, sock: socket.socket) -> bytes: + def _shift(self): + count = self.size - self.offset + + self.buf[0:count] = self.buf[self.offset:self.count] + + self.offset = 0 + self.count = count + + def _is_full(self): + return self.count == self.size + + def _fill(self, sock: socket.socket): + if self.eof: + return + + readlen = self.size - self.count + data = sock.recv(readlen) + + if data == b'': + self.eof = True + return + + readlen = len(data) + + self.buf[self.offset:self.offset+readlen] = data + + self.count += readlen + + def _flush(self, end: int): + ret = self.buf[self.offset:end+1] + + self.offset = end + 1 + + return str(ret, 'ascii') + + def readline(self, sock: socket.socket) -> str: while True: - # - # Check and see if there is already a line in the - # internal buffer. - # - offset = self.buf.find(b'\n', self.offset, len(self.buf)) + index = self.buf.find(b'\n', self.offset, self.count) - if offset < 0: - # - # There is no line in the buffer. Check and see if - # there is more room in the buffer to continue - # reading. - # - if offset == self.BUFFER_SIZE - 1: - # - # There is no room left in the buffer; the line - # must be too long. Panic. - # - raise BufferOverflow() - - # - # Read from the socket to attempt to fill the buffer - # more. - # - data = sock.recv(self.BUFFER_SIZE - self.offset) - - # - # Check to see if recv() returned any data. If not, - # the socket is closed. Return an empty string. - # - if data == b'': - return '' - - count = len(data) - - # - # Now, fill the buffer. The next iteration of the - # loop will be used to check for a newline. - # - self.buf[self.offset:self.offset+count] = data + if index < 0: + if self._is_full(): + if self.offset == 0: + raise BufferOverflow() + else: + self._shift() + else: + self._fill(sock) else: - # - # There is a line in the buffer. Prepare to return - # it. - # - ret = self.buf[self.offset:offset] - - # - # Advance the offset to one character beyond the - # newline. - # - self.offset = offset + 1 - - return str(ret, 'ascii') + return self._flush(index)