tests/badserverext.py
changeset 41463 ba7298160357
parent 39428 cbfab495dbcf
child 41464 d343d9ac173e
equal deleted inserted replaced
41462:9b2b8794f801 41463:ba7298160357
    46 
    46 
    47 configitem(b'badserver', b'closeafteraccept',
    47 configitem(b'badserver', b'closeafteraccept',
    48     default=False,
    48     default=False,
    49 )
    49 )
    50 configitem(b'badserver', b'closeafterrecvbytes',
    50 configitem(b'badserver', b'closeafterrecvbytes',
    51     default='0',
    51     default=b'0',
    52 )
    52 )
    53 configitem(b'badserver', b'closeaftersendbytes',
    53 configitem(b'badserver', b'closeaftersendbytes',
    54     default='0',
    54     default=b'0',
    55 )
    55 )
    56 configitem(b'badserver', b'closebeforeaccept',
    56 configitem(b'badserver', b'closebeforeaccept',
    57     default=False,
    57     default=False,
    58 )
    58 )
    59 
    59 
   125 
   125 
   126     def __setattr__(self, name, value):
   126     def __setattr__(self, name, value):
   127         setattr(object.__getattribute__(self, '_orig'), name, value)
   127         setattr(object.__getattribute__(self, '_orig'), name, value)
   128 
   128 
   129     def _writelog(self, msg):
   129     def _writelog(self, msg):
   130         msg = msg.replace('\r', '\\r').replace('\n', '\\n')
   130         msg = msg.replace(b'\r', b'\\r').replace(b'\n', b'\\n')
   131 
   131 
   132         object.__getattribute__(self, '_logfp').write(msg)
   132         object.__getattribute__(self, '_logfp').write(msg)
   133         object.__getattribute__(self, '_logfp').write('\n')
   133         object.__getattribute__(self, '_logfp').write(b'\n')
   134         object.__getattribute__(self, '_logfp').flush()
   134         object.__getattribute__(self, '_logfp').flush()
   135 
   135 
   136     def read(self, size=-1):
   136     def read(self, size=-1):
   137         remaining = object.__getattribute__(self, '_closeafterrecvbytes')
   137         remaining = object.__getattribute__(self, '_closeafterrecvbytes')
   138 
   138 
   139         # No read limit. Call original function.
   139         # No read limit. Call original function.
   140         if not remaining:
   140         if not remaining:
   141             result = object.__getattribute__(self, '_orig').read(size)
   141             result = object.__getattribute__(self, '_orig').read(size)
   142             self._writelog('read(%d) -> (%d) (%s) %s' % (size,
   142             self._writelog(b'read(%d) -> (%d) (%s) %s' % (size,
   143                                                            len(result),
   143                                                           len(result),
   144                                                            result))
   144                                                           result))
   145             return result
   145             return result
   146 
   146 
   147         origsize = size
   147         origsize = size
   148 
   148 
   149         if size < 0:
   149         if size < 0:
   152             size = min(remaining, size)
   152             size = min(remaining, size)
   153 
   153 
   154         result = object.__getattribute__(self, '_orig').read(size)
   154         result = object.__getattribute__(self, '_orig').read(size)
   155         remaining -= len(result)
   155         remaining -= len(result)
   156 
   156 
   157         self._writelog('read(%d from %d) -> (%d) %s' % (
   157         self._writelog(b'read(%d from %d) -> (%d) %s' % (
   158             size, origsize, len(result), result))
   158             size, origsize, len(result), result))
   159 
   159 
   160         object.__setattr__(self, '_closeafterrecvbytes', remaining)
   160         object.__setattr__(self, '_closeafterrecvbytes', remaining)
   161 
   161 
   162         if remaining <= 0:
   162         if remaining <= 0:
   163             self._writelog('read limit reached, closing socket')
   163             self._writelog(b'read limit reached, closing socket')
   164             self._sock.close()
   164             self._sock.close()
   165             # This is the easiest way to abort the current request.
   165             # This is the easiest way to abort the current request.
   166             raise Exception('connection closed after receiving N bytes')
   166             raise Exception('connection closed after receiving N bytes')
   167 
   167 
   168         return result
   168         return result
   171         remaining = object.__getattribute__(self, '_closeafterrecvbytes')
   171         remaining = object.__getattribute__(self, '_closeafterrecvbytes')
   172 
   172 
   173         # No read limit. Call original function.
   173         # No read limit. Call original function.
   174         if not remaining:
   174         if not remaining:
   175             result = object.__getattribute__(self, '_orig').readline(size)
   175             result = object.__getattribute__(self, '_orig').readline(size)
   176             self._writelog('readline(%d) -> (%d) %s' % (
   176             self._writelog(b'readline(%d) -> (%d) %s' % (
   177                 size, len(result), result))
   177                 size, len(result), result))
   178             return result
   178             return result
   179 
   179 
   180         origsize = size
   180         origsize = size
   181 
   181 
   185             size = min(remaining, size)
   185             size = min(remaining, size)
   186 
   186 
   187         result = object.__getattribute__(self, '_orig').readline(size)
   187         result = object.__getattribute__(self, '_orig').readline(size)
   188         remaining -= len(result)
   188         remaining -= len(result)
   189 
   189 
   190         self._writelog('readline(%d from %d) -> (%d) %s' % (
   190         self._writelog(b'readline(%d from %d) -> (%d) %s' % (
   191             size, origsize, len(result), result))
   191             size, origsize, len(result), result))
   192 
   192 
   193         object.__setattr__(self, '_closeafterrecvbytes', remaining)
   193         object.__setattr__(self, '_closeafterrecvbytes', remaining)
   194 
   194 
   195         if remaining <= 0:
   195         if remaining <= 0:
   196             self._writelog('read limit reached; closing socket')
   196             self._writelog(b'read limit reached; closing socket')
   197             self._sock.close()
   197             self._sock.close()
   198             # This is the easiest way to abort the current request.
   198             # This is the easiest way to abort the current request.
   199             raise Exception('connection closed after receiving N bytes')
   199             raise Exception('connection closed after receiving N bytes')
   200 
   200 
   201         return result
   201         return result
   203     def write(self, data):
   203     def write(self, data):
   204         remaining = object.__getattribute__(self, '_closeaftersendbytes')
   204         remaining = object.__getattribute__(self, '_closeaftersendbytes')
   205 
   205 
   206         # No byte limit on this operation. Call original function.
   206         # No byte limit on this operation. Call original function.
   207         if not remaining:
   207         if not remaining:
   208             self._writelog('write(%d) -> %s' % (len(data), data))
   208             self._writelog(b'write(%d) -> %s' % (len(data), data))
   209             result = object.__getattribute__(self, '_orig').write(data)
   209             result = object.__getattribute__(self, '_orig').write(data)
   210             return result
   210             return result
   211 
   211 
   212         if len(data) > remaining:
   212         if len(data) > remaining:
   213             newdata = data[0:remaining]
   213             newdata = data[0:remaining]
   214         else:
   214         else:
   215             newdata = data
   215             newdata = data
   216 
   216 
   217         remaining -= len(newdata)
   217         remaining -= len(newdata)
   218 
   218 
   219         self._writelog('write(%d from %d) -> (%d) %s' % (
   219         self._writelog(b'write(%d from %d) -> (%d) %s' % (
   220             len(newdata), len(data), remaining, newdata))
   220             len(newdata), len(data), remaining, newdata))
   221 
   221 
   222         result = object.__getattribute__(self, '_orig').write(newdata)
   222         result = object.__getattribute__(self, '_orig').write(newdata)
   223 
   223 
   224         object.__setattr__(self, '_closeaftersendbytes', remaining)
   224         object.__setattr__(self, '_closeaftersendbytes', remaining)
   225 
   225 
   226         if remaining <= 0:
   226         if remaining <= 0:
   227             self._writelog('write limit reached; closing socket')
   227             self._writelog(b'write limit reached; closing socket')
   228             self._sock.close()
   228             self._sock.close()
   229             raise Exception('connection closed after sending N bytes')
   229             raise Exception('connection closed after sending N bytes')
   230 
   230 
   231         return result
   231         return result
   232 
   232 
   237         def __init__(self, ui, *args, **kwargs):
   237         def __init__(self, ui, *args, **kwargs):
   238             self._ui = ui
   238             self._ui = ui
   239             super(badserver, self).__init__(ui, *args, **kwargs)
   239             super(badserver, self).__init__(ui, *args, **kwargs)
   240 
   240 
   241             recvbytes = self._ui.config(b'badserver', b'closeafterrecvbytes')
   241             recvbytes = self._ui.config(b'badserver', b'closeafterrecvbytes')
   242             recvbytes = recvbytes.split(',')
   242             recvbytes = recvbytes.split(b',')
   243             self.closeafterrecvbytes = [int(v) for v in recvbytes if v]
   243             self.closeafterrecvbytes = [int(v) for v in recvbytes if v]
   244             sendbytes = self._ui.config(b'badserver', b'closeaftersendbytes')
   244             sendbytes = self._ui.config(b'badserver', b'closeaftersendbytes')
   245             sendbytes = sendbytes.split(',')
   245             sendbytes = sendbytes.split(b',')
   246             self.closeaftersendbytes = [int(v) for v in sendbytes if v]
   246             self.closeaftersendbytes = [int(v) for v in sendbytes if v]
   247 
   247 
   248             # Need to inherit object so super() works.
   248             # Need to inherit object so super() works.
   249             class badrequesthandler(self.RequestHandlerClass, object):
   249             class badrequesthandler(self.RequestHandlerClass, object):
   250                 def send_header(self, name, value):
   250                 def send_header(self, name, value):
   268                 self.__shutdown_request = True
   268                 self.__shutdown_request = True
   269 
   269 
   270                 # Simulate failure to stop processing this request.
   270                 # Simulate failure to stop processing this request.
   271                 raise socket.error('close before accept')
   271                 raise socket.error('close before accept')
   272 
   272 
   273             if self._ui.configbool('badserver', 'closeafteraccept'):
   273             if self._ui.configbool(b'badserver', b'closeafteraccept'):
   274                 request, client_address = super(badserver, self).get_request()
   274                 request, client_address = super(badserver, self).get_request()
   275                 request.close()
   275                 request.close()
   276                 raise socket.error('close after accept')
   276                 raise socket.error('close after accept')
   277 
   277 
   278             return super(badserver, self).get_request()
   278             return super(badserver, self).get_request()