| Alphare |
| hg-reviewers |
This will make sure both read and readline do the same processing and make
it simpler to update that processing in the future.
| Automatic diff as part of commit; lint not applicable. |
| Automatic diff as part of commit; unit tests not applicable. |
| Path | Packages | |||
|---|---|---|---|---|
| M | tests/testlib/badserverext.py (119 lines) |
| if remaining <= 0: | if remaining <= 0: | ||||
| obj._writelog(b'write limit reached; closing socket') | obj._writelog(b'write limit reached; closing socket') | ||||
| object.__getattribute__(obj, '_cond_close')() | object.__getattribute__(obj, '_cond_close')() | ||||
| raise Exception('connection closed after sending N bytes') | raise Exception('connection closed after sending N bytes') | ||||
| return result | return result | ||||
| def forward_read(self, obj, method, size=-1): | |||||
| """call an underlying read function until condition are met | |||||
| When the condition are met the socket is closed | |||||
| """ | |||||
| remaining = self.remaining_recv_bytes | |||||
| orig = object.__getattribute__(obj, '_orig') | |||||
| bmethod = method.encode('ascii') | |||||
| func = getattr(orig, method) | |||||
| # No read limit. Call original function. | |||||
| if not remaining: | |||||
| result = func(size) | |||||
| obj._writelog( | |||||
| b'%s(%d) -> (%d) %s' % (bmethod, size, len(result), result) | |||||
| ) | |||||
| return result | |||||
| origsize = size | |||||
| if size < 0: | |||||
| size = remaining | |||||
| else: | |||||
| size = min(remaining, size) | |||||
| result = func(size) | |||||
| remaining -= len(result) | |||||
| obj._writelog( | |||||
| b'%s(%d from %d) -> (%d) %s' | |||||
| % (bmethod, size, origsize, len(result), result) | |||||
| ) | |||||
| self.remaining_recv_bytes = remaining | |||||
| if remaining <= 0: | |||||
| obj._writelog(b'read limit reached; closing socket') | |||||
| obj._cond_close() | |||||
| # This is the easiest way to abort the current request. | |||||
| raise Exception('connection closed after receiving N bytes') | |||||
| return result | |||||
| # We can't adjust __class__ on a socket instance. So we define a proxy type. | # We can't adjust __class__ on a socket instance. So we define a proxy type. | ||||
| class socketproxy(object): | class socketproxy(object): | ||||
| __slots__ = ('_orig', '_logfp', '_cond') | __slots__ = ('_orig', '_logfp', '_cond') | ||||
| def __init__(self, obj, logfp, condition_tracked): | def __init__(self, obj, logfp, condition_tracked): | ||||
| object.__setattr__(self, '_orig', obj) | object.__setattr__(self, '_orig', obj) | ||||
| object.__setattr__(self, '_logfp', logfp) | object.__setattr__(self, '_logfp', logfp) | ||||
| if hasattr(orig, 'raw'): | if hasattr(orig, 'raw'): | ||||
| orig.raw._sock.shutdown(socket.SHUT_RDWR) | orig.raw._sock.shutdown(socket.SHUT_RDWR) | ||||
| else: | else: | ||||
| self.close() | self.close() | ||||
| else: | else: | ||||
| self._sock.shutdown(socket.SHUT_RDWR) | self._sock.shutdown(socket.SHUT_RDWR) | ||||
| def read(self, size=-1): | def read(self, size=-1): | ||||
| remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes | cond = object.__getattribute__(self, '_cond') | ||||
| return cond.forward_read(self, 'read', size) | |||||
| # No read limit. Call original function. | |||||
| if not remaining: | |||||
| result = object.__getattribute__(self, '_orig').read(size) | |||||
| self._writelog( | |||||
| b'read(%d) -> (%d) (%s) %s' % (size, len(result), result) | |||||
| ) | |||||
| return result | |||||
| origsize = size | |||||
| if size < 0: | |||||
| size = remaining | |||||
| else: | |||||
| size = min(remaining, size) | |||||
| result = object.__getattribute__(self, '_orig').read(size) | |||||
| remaining -= len(result) | |||||
| self._writelog( | |||||
| b'read(%d from %d) -> (%d) %s' | |||||
| % (size, origsize, len(result), result) | |||||
| ) | |||||
| object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining | |||||
| if remaining <= 0: | |||||
| self._writelog(b'read limit reached; closing socket') | |||||
| self._close() | |||||
| # This is the easiest way to abort the current request. | |||||
| raise Exception('connection closed after receiving N bytes') | |||||
| return result | |||||
| def readline(self, size=-1): | def readline(self, size=-1): | ||||
| remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes | cond = object.__getattribute__(self, '_cond') | ||||
| return cond.forward_read(self, 'readline', size) | |||||
| # No read limit. Call original function. | |||||
| if not remaining: | |||||
| result = object.__getattribute__(self, '_orig').readline(size) | |||||
| self._writelog( | |||||
| b'readline(%d) -> (%d) %s' % (size, len(result), result) | |||||
| ) | |||||
| return result | |||||
| origsize = size | |||||
| if size < 0: | |||||
| size = remaining | |||||
| else: | |||||
| size = min(remaining, size) | |||||
| result = object.__getattribute__(self, '_orig').readline(size) | |||||
| remaining -= len(result) | |||||
| self._writelog( | |||||
| b'readline(%d from %d) -> (%d) %s' | |||||
| % (size, origsize, len(result), result) | |||||
| ) | |||||
| object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining | |||||
| if remaining <= 0: | |||||
| self._writelog(b'read limit reached; closing socket') | |||||
| self._close() | |||||
| # This is the easiest way to abort the current request. | |||||
| raise Exception('connection closed after receiving N bytes') | |||||
| return result | |||||
| def write(self, data): | def write(self, data): | ||||
| cond = object.__getattribute__(self, '_cond') | cond = object.__getattribute__(self, '_cond') | ||||
| return cond.forward_write(self, 'write', data) | return cond.forward_write(self, 'write', data) | ||||
| def _cond_close(self): | def _cond_close(self): | ||||
| self._close() | self._close() | ||||