| Alphare |
| hg-reviewers |
This will make sure both sendall and write 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 (130 lines) |
| ) | ) | ||||
| class ConditionTracker(object): | class ConditionTracker(object): | ||||
| def __init__(self, close_after_recv_bytes, close_after_send_bytes): | def __init__(self, close_after_recv_bytes, close_after_send_bytes): | ||||
| self._all_close_after_recv_bytes = close_after_recv_bytes | self._all_close_after_recv_bytes = close_after_recv_bytes | ||||
| self._all_close_after_send_bytes = close_after_send_bytes | self._all_close_after_send_bytes = close_after_send_bytes | ||||
| self.target_recv_bytes = None | |||||
| self.remaining_recv_bytes = None | |||||
| self.target_send_bytes = None | |||||
| self.remaining_send_bytes = None | |||||
| def start_next_request(self): | def start_next_request(self): | ||||
| """move to the next set of close condition""" | """move to the next set of close condition""" | ||||
| if self._all_close_after_recv_bytes: | if self._all_close_after_recv_bytes: | ||||
| self.target_recv_bytes = self._all_close_after_recv_bytes.pop(0) | self.target_recv_bytes = self._all_close_after_recv_bytes.pop(0) | ||||
| self.remaining_recv_bytes = self.target_recv_bytes | self.remaining_recv_bytes = self.target_recv_bytes | ||||
| else: | else: | ||||
| self.target_recv_bytes = None | self.target_recv_bytes = None | ||||
| self.remaining_recv_bytes = None | self.remaining_recv_bytes = None | ||||
| if self._all_close_after_send_bytes: | if self._all_close_after_send_bytes: | ||||
| self.target_send_bytes = self._all_close_after_send_bytes.pop(0) | self.target_send_bytes = self._all_close_after_send_bytes.pop(0) | ||||
| self.remaining_send_bytes = self.target_send_bytes | self.remaining_send_bytes = self.target_send_bytes | ||||
| else: | else: | ||||
| self.target_send_bytes = None | self.target_send_bytes = None | ||||
| self.remaining_send_bytes = None | self.remaining_send_bytes = None | ||||
| def might_close(self): | def might_close(self): | ||||
| """True, if any processing will be needed""" | """True, if any processing will be needed""" | ||||
| if self.remaining_recv_bytes is not None: | if self.remaining_recv_bytes is not None: | ||||
| return True | return True | ||||
| if self.remaining_send_bytes is not None: | if self.remaining_send_bytes is not None: | ||||
| return True | return True | ||||
| return False | return False | ||||
| def forward_write(self, obj, method, data, *args, **kwargs): | |||||
| """call an underlying write function until condition are met | |||||
| When the condition are met the socket is closed | |||||
| """ | |||||
| remaining = self.remaining_send_bytes | |||||
| orig = object.__getattribute__(obj, '_orig') | |||||
| bmethod = method.encode('ascii') | |||||
| func = getattr(orig, method) | |||||
| # No byte limit on this operation. Call original function. | |||||
| if not remaining: | |||||
| result = func(data, *args, **kwargs) | |||||
| obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data)) | |||||
| return result | |||||
| remaining = max(0, remaining) | |||||
| if remaining > 0: | |||||
| if remaining < len(data): | |||||
| newdata = data[0:remaining] | |||||
| else: | |||||
| newdata = data | |||||
| remaining -= len(newdata) | |||||
| obj._writelog( | |||||
| b'%s(%d from %d) -> (%d) %s' | |||||
| % ( | |||||
| bmethod, | |||||
| len(newdata), | |||||
| len(data), | |||||
| remaining, | |||||
| newdata, | |||||
| ) | |||||
| ) | |||||
| result = func(newdata, *args, **kwargs) | |||||
| self.remaining_send_bytes = remaining | |||||
| if remaining <= 0: | |||||
| obj._writelog(b'write limit reached; closing socket') | |||||
| object.__getattribute__(obj, '_cond_close')() | |||||
| raise Exception('connection closed after sending 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) | ||||
| f = object.__getattribute__(self, '_orig').makefile(mode, bufsize) | f = object.__getattribute__(self, '_orig').makefile(mode, bufsize) | ||||
| logfp = object.__getattribute__(self, '_logfp') | logfp = object.__getattribute__(self, '_logfp') | ||||
| cond = object.__getattribute__(self, '_cond') | cond = object.__getattribute__(self, '_cond') | ||||
| return fileobjectproxy(f, logfp, cond) | return fileobjectproxy(f, logfp, cond) | ||||
| def sendall(self, data, flags=0): | def sendall(self, data, flags=0): | ||||
| remaining = object.__getattribute__(self, '_cond').remaining_send_bytes | cond = object.__getattribute__(self, '_cond') | ||||
| return cond.forward_write(self, 'sendall', data, flags) | |||||
| # No read limit. Call original function. | |||||
| if not remaining: | |||||
| result = object.__getattribute__(self, '_orig').sendall(data, flags) | |||||
| self._writelog(b'sendall(%d) -> %s' % (len(data), data)) | |||||
| return result | |||||
| if len(data) > remaining: | |||||
| newdata = data[0:remaining] | |||||
| else: | |||||
| newdata = data | |||||
| remaining -= len(newdata) | |||||
| result = object.__getattribute__(self, '_orig').sendall(newdata, flags) | |||||
| self._writelog( | |||||
| b'sendall(%d from %d) -> (%d) %s' | |||||
| % (len(newdata), len(data), remaining, newdata) | |||||
| ) | |||||
| object.__getattribute__(self, '_cond').remaining_send_bytes = remaining | |||||
| if remaining <= 0: | def _cond_close(self): | ||||
| self._writelog(b'write limit reached; closing socket') | |||||
| object.__getattribute__(self, '_orig').shutdown(socket.SHUT_RDWR) | object.__getattribute__(self, '_orig').shutdown(socket.SHUT_RDWR) | ||||
| raise Exception('connection closed after sending N bytes') | |||||
| return result | |||||
| # We can't adjust __class__ on socket._fileobject, so define a proxy. | # We can't adjust __class__ on socket._fileobject, so define a proxy. | ||||
| class fileobjectproxy(object): | class fileobjectproxy(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) | ||||
| object.__setattr__(self, '_cond', condition_tracked) | object.__setattr__(self, '_cond', condition_tracked) | ||||
| def __getattribute__(self, name): | def __getattribute__(self, name): | ||||
| if name in ('_close', 'read', 'readline', 'write', '_writelog'): | if name in ( | ||||
| '_close', | |||||
| 'read', | |||||
| 'readline', | |||||
| 'write', | |||||
| '_writelog', | |||||
| '_cond_close', | |||||
| ): | |||||
| return object.__getattribute__(self, name) | return object.__getattribute__(self, name) | ||||
| return getattr(object.__getattribute__(self, '_orig'), name) | return getattr(object.__getattribute__(self, '_orig'), name) | ||||
| def __delattr__(self, name): | def __delattr__(self, name): | ||||
| delattr(object.__getattribute__(self, '_orig'), name) | delattr(object.__getattribute__(self, '_orig'), name) | ||||
| def __setattr__(self, name, value): | def __setattr__(self, name, value): | ||||
| self._close() | self._close() | ||||
| # This is the easiest way to abort the current request. | # This is the easiest way to abort the current request. | ||||
| raise Exception('connection closed after receiving N bytes') | raise Exception('connection closed after receiving N bytes') | ||||
| return result | return result | ||||
| def write(self, data): | def write(self, data): | ||||
| remaining = object.__getattribute__(self, '_cond').remaining_send_bytes | cond = object.__getattribute__(self, '_cond') | ||||
| return cond.forward_write(self, 'write', data) | |||||
| # No byte limit on this operation. Call original function. | |||||
| if not remaining: | |||||
| result = object.__getattribute__(self, '_orig').write(data) | |||||
| self._writelog(b'write(%d) -> %s' % (len(data), data)) | |||||
| return result | |||||
| if len(data) > remaining: | |||||
| newdata = data[0:remaining] | |||||
| else: | |||||
| newdata = data | |||||
| remaining -= len(newdata) | |||||
| result = object.__getattribute__(self, '_orig').write(newdata) | |||||
| self._writelog( | |||||
| b'write(%d from %d) -> (%d) %s' | |||||
| % (len(newdata), len(data), remaining, newdata) | |||||
| ) | |||||
| object.__getattribute__(self, '_cond').remaining_send_bytes = remaining | |||||
| if remaining <= 0: | def _cond_close(self): | ||||
| self._writelog(b'write limit reached; closing socket') | |||||
| self._close() | self._close() | ||||
| raise Exception('connection closed after sending N bytes') | |||||
| return result | |||||
| def process_config(value): | def process_config(value): | ||||
| parts = value.split(b',') | parts = value.split(b',') | ||||
| integers = [int(v) for v in parts if v] | integers = [int(v) for v in parts if v] | ||||
| return [v if v else None for v in integers] | return [v if v else None for v in integers] | ||||
| def extsetup(ui): | def extsetup(ui): | ||||