| 10 | |
| 11 | |
| 12 | class ChunkedReader: |
| 13 | def __init__(self, req, unreader): |
| 14 | self.req = req |
| 15 | self.parser = self.parse_chunked(unreader) |
| 16 | self.buf = io.BytesIO() |
| 17 | |
| 18 | def read(self, size): |
| 19 | if not isinstance(size, int): |
| 20 | raise TypeError("size must be an integer type") |
| 21 | if size < 0: |
| 22 | raise ValueError("Size must be positive.") |
| 23 | if size == 0: |
| 24 | return b"" |
| 25 | |
| 26 | if self.parser: |
| 27 | while self.buf.tell() < size: |
| 28 | try: |
| 29 | self.buf.write(next(self.parser)) |
| 30 | except StopIteration: |
| 31 | self.parser = None |
| 32 | break |
| 33 | |
| 34 | data = self.buf.getvalue() |
| 35 | ret, rest = data[:size], data[size:] |
| 36 | self.buf = io.BytesIO() |
| 37 | self.buf.write(rest) |
| 38 | return ret |
| 39 | |
| 40 | def parse_trailers(self, unreader, data): |
| 41 | buf = io.BytesIO() |
| 42 | buf.write(data) |
| 43 | |
| 44 | idx = buf.getvalue().find(b"\r\n\r\n") |
| 45 | done = buf.getvalue()[:2] == b"\r\n" |
| 46 | while idx < 0 and not done: |
| 47 | self.get_data(unreader, buf) |
| 48 | idx = buf.getvalue().find(b"\r\n\r\n") |
| 49 | done = buf.getvalue()[:2] == b"\r\n" |
| 50 | if done: |
| 51 | unreader.unread(buf.getvalue()[2:]) |
| 52 | return b"" |
| 53 | self.req.trailers = self.req.parse_headers(buf.getvalue()[:idx], from_trailer=True) |
| 54 | unreader.unread(buf.getvalue()[idx + 4:]) |
| 55 | |
| 56 | def parse_chunked(self, unreader): |
| 57 | (size, rest) = self.parse_chunk_size(unreader) |
| 58 | while size > 0: |
| 59 | while size > len(rest): |
| 60 | size -= len(rest) |
| 61 | yield rest |
| 62 | rest = unreader.read() |
| 63 | if not rest: |
| 64 | raise NoMoreData() |
| 65 | yield rest[:size] |
| 66 | # Remove \r\n after chunk |
| 67 | rest = rest[size:] |
| 68 | while len(rest) < 2: |
| 69 | new_data = unreader.read() |