| Alphare |
| hg-reviewers |
Instead of using a fixed number a bytes, a very unstable method. We introduce a
new config that can define a regex. One this regex is detected, nothing after
it will be written and the socket will be closed.
Tests will be migrated to this new method in later changesets.
| No Linters Available |
| No Unit Test Coverage |
| Path | Packages | |||
|---|---|---|---|---|
| M | tests/testlib/badserverext.py (78 lines) |
| Commit | Parents | Author | Summary | Date |
|---|---|---|---|---|
| c096d0d7bd61 | 8d33c1a9ef1a | Pierre-Yves David | Jan 21 2022, 7:51 PM |
| If defined, close the client socket after receiving this many bytes. | If defined, close the client socket after receiving this many bytes. | ||||
| (The value is a list, multiple value can use used to close a series of | (The value is a list, multiple value can use used to close a series of | ||||
| request) | request) | ||||
| close-after-send-bytes | close-after-send-bytes | ||||
| If defined, close the client socket after sending this many bytes. | If defined, close the client socket after sending this many bytes. | ||||
| (The value is a list, multiple value can use used to close a series of | (The value is a list, multiple value can use used to close a series of | ||||
| request) | request) | ||||
| close-after-send-patterns | |||||
| If defined, close the client socket after the configured regexp is seen. | |||||
| (The value is a list, multiple value can use used to close a series of | |||||
| request) | |||||
| """ | """ | ||||
| from __future__ import absolute_import | from __future__ import absolute_import | ||||
| import re | |||||
| import socket | import socket | ||||
| from mercurial import ( | from mercurial import ( | ||||
| pycompat, | pycompat, | ||||
| registrar, | registrar, | ||||
| ) | ) | ||||
| from mercurial.hgweb import server | from mercurial.hgweb import server | ||||
| ) | ) | ||||
| configitem( | configitem( | ||||
| b'badserver', | b'badserver', | ||||
| b'close-after-send-bytes', | b'close-after-send-bytes', | ||||
| default=b'0', | default=b'0', | ||||
| ) | ) | ||||
| configitem( | configitem( | ||||
| b'badserver', | b'badserver', | ||||
| b'close-after-send-patterns', | |||||
| default=b'', | |||||
| ) | |||||
| configitem( | |||||
| b'badserver', | |||||
| b'close-before-accept', | b'close-before-accept', | ||||
| default=False, | default=False, | ||||
| ) | ) | ||||
| 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, | |||||
| close_after_send_patterns, | |||||
| ): | |||||
| 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._all_close_after_send_patterns = close_after_send_patterns | |||||
| self.target_recv_bytes = None | self.target_recv_bytes = None | ||||
| self.remaining_recv_bytes = None | self.remaining_recv_bytes = None | ||||
| self.target_send_bytes = None | self.target_send_bytes = None | ||||
| self.remaining_send_bytes = None | self.remaining_send_bytes = None | ||||
| self.send_pattern = None | |||||
| self.send_data = b'' | |||||
| 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 | ||||
| self.send_data = b'' | |||||
| if self._all_close_after_send_patterns: | |||||
| self.send_pattern = self._all_close_after_send_patterns.pop(0) | |||||
| else: | |||||
| self.send_pattern = 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 | ||||
| if self.send_pattern is not None: | |||||
| return True | |||||
| return False | return False | ||||
| def forward_write(self, obj, method, data, *args, **kwargs): | def forward_write(self, obj, method, data, *args, **kwargs): | ||||
| """call an underlying write function until condition are met | """call an underlying write function until condition are met | ||||
| When the condition are met the socket is closed | When the condition are met the socket is closed | ||||
| """ | """ | ||||
| remaining = self.remaining_send_bytes | remaining = self.remaining_send_bytes | ||||
| pattern = self.send_pattern | |||||
| orig = object.__getattribute__(obj, '_orig') | orig = object.__getattribute__(obj, '_orig') | ||||
| bmethod = method.encode('ascii') | bmethod = method.encode('ascii') | ||||
| func = getattr(orig, method) | func = getattr(orig, method) | ||||
| if pattern: | |||||
| self.send_data += data | |||||
| pieces = pattern.split(self.send_data, maxsplit=1) | |||||
| if len(pieces) > 1: | |||||
| dropped = len(pieces[-1]) | |||||
| remaining = len(data) - dropped | |||||
| if remaining: | if remaining: | ||||
| remaining = max(0, remaining) | remaining = max(0, remaining) | ||||
| if not remaining: | if not remaining: | ||||
| newdata = data | newdata = data | ||||
| else: | else: | ||||
| if remaining < len(data): | if remaining < len(data): | ||||
| newdata = data[0:remaining] | newdata = data[0:remaining] | ||||
| else: | else: | ||||
| newdata = data | newdata = data | ||||
| remaining -= len(newdata) | remaining -= len(newdata) | ||||
| self.remaining_send_bytes = remaining | self.remaining_send_bytes = remaining | ||||
| result = func(newdata, *args, **kwargs) | result = func(newdata, *args, **kwargs) | ||||
| if remaining is None: | if remaining is None: | ||||
| obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data)) | obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data)) | ||||
| else: | else: | ||||
| obj._writelog( | msg = b'%s(%d from %d) -> (%d) %s' | ||||
| b'%s(%d from %d) -> (%d) %s' | msg %= (bmethod, len(newdata), len(data), remaining, newdata) | ||||
| % ( | obj._writelog(msg) | ||||
| bmethod, | |||||
| len(newdata), | |||||
| len(data), | |||||
| remaining, | |||||
| newdata, | |||||
| ) | |||||
| ) | |||||
| if remaining is not None and remaining <= 0: | if remaining is not None and 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 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() | ||||
| def process_config(value): | def process_bytes_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 process_pattern_config(value): | |||||
| patterns = [] | |||||
| for p in value.split(b','): | |||||
| if not p: | |||||
| p = None | |||||
| else: | |||||
| p = re.compile(p, re.DOTALL | re.MULTILINE) | |||||
| patterns.append(p) | |||||
| return patterns | |||||
| def extsetup(ui): | def extsetup(ui): | ||||
| # Change the base HTTP server class so various events can be performed. | # Change the base HTTP server class so various events can be performed. | ||||
| # See SocketServer.BaseServer for how the specially named methods work. | # See SocketServer.BaseServer for how the specially named methods work. | ||||
| class badserver(server.MercurialHTTPServer): | class badserver(server.MercurialHTTPServer): | ||||
| def __init__(self, ui, *args, **kwargs): | def __init__(self, ui, *args, **kwargs): | ||||
| self._ui = ui | self._ui = ui | ||||
| super(badserver, self).__init__(ui, *args, **kwargs) | super(badserver, self).__init__(ui, *args, **kwargs) | ||||
| all_recv_bytes = self._ui.config( | all_recv_bytes = self._ui.config( | ||||
| b'badserver', b'close-after-recv-bytes' | b'badserver', b'close-after-recv-bytes' | ||||
| ) | ) | ||||
| all_recv_bytes = process_config(all_recv_bytes) | all_recv_bytes = process_bytes_config(all_recv_bytes) | ||||
| all_send_bytes = self._ui.config( | all_send_bytes = self._ui.config( | ||||
| b'badserver', b'close-after-send-bytes' | b'badserver', b'close-after-send-bytes' | ||||
| ) | ) | ||||
| all_send_bytes = process_config(all_send_bytes) | all_send_bytes = process_bytes_config(all_send_bytes) | ||||
| self._cond = ConditionTracker(all_recv_bytes, all_send_bytes) | all_send_patterns = self._ui.config( | ||||
| b'badserver', b'close-after-send-patterns' | |||||
| ) | |||||
| all_send_patterns = process_pattern_config(all_send_patterns) | |||||
| self._cond = ConditionTracker( | |||||
| all_recv_bytes, | |||||
| all_send_bytes, | |||||
| all_send_patterns, | |||||
| ) | |||||
| # Need to inherit object so super() works. | # Need to inherit object so super() works. | ||||
| class badrequesthandler(self.RequestHandlerClass, object): | class badrequesthandler(self.RequestHandlerClass, object): | ||||
| def send_header(self, name, value): | def send_header(self, name, value): | ||||
| # Make headers deterministic to facilitate testing. | # Make headers deterministic to facilitate testing. | ||||
| if name.lower() == 'date': | if name.lower() == 'date': | ||||
| value = 'Fri, 14 Apr 2017 00:00:00 GMT' | value = 'Fri, 14 Apr 2017 00:00:00 GMT' | ||||
| elif name.lower() == 'server': | elif name.lower() == 'server': | ||||