1 # badserverext.py - Extension making servers behave badly

     2 #

     3 # Copyright 2017 Gregory Szorc <gregory.szorc@gmail.com>

     4 #

     5 # This software may be used and distributed according to the terms of the

     6 # GNU General Public License version 2 or any later version.

     7 

     8 # no-check-code

     9 

    10 """Extension to make servers behave badly.

    11 

    12 This extension is useful for testing Mercurial behavior when various network

    13 events occur.

    14 

    15 Various config options in the [badserver] section influence behavior:

    16 

    17 closebeforeaccept

    18    If true, close() the server socket when a new connection arrives before

    19    accept() is called. The server will then exit.

    20 

    21 closeafteraccept

    22    If true, the server will close() the client socket immediately after

    23    accept().

    24 

    25 closeafterrecvbytes

    26    If defined, close the client socket after receiving this many bytes.

    27 

    28 closeaftersendbytes

    29    If defined, close the client socket after sending this many bytes.

    30 """

    31 

    32 from __future__ import absolute_import

    33 

    34 import socket

    35 

    36 from mercurial.hgweb import (

    37     server,

    38 )

    39 

    40 # We can't adjust __class__ on a socket instance. So we define a proxy type.

    41 class socketproxy(object):

    42     __slots__ = (

    43         '_orig',

    44         '_logfp',

    45         '_closeafterrecvbytes',

    46         '_closeaftersendbytes',

    47     )

    48 

    49     def __init__(self, obj, logfp, closeafterrecvbytes=0,

    50                  closeaftersendbytes=0):

    51         object.__setattr__(self, '_orig', obj)

    52         object.__setattr__(self, '_logfp', logfp)

    53         object.__setattr__(self, '_closeafterrecvbytes', closeafterrecvbytes)

    54         object.__setattr__(self, '_closeaftersendbytes', closeaftersendbytes)

    55 

    56     def __getattribute__(self, name):

    57         if name in ('makefile',):

    58             return object.__getattribute__(self, name)

    59 

    60         return getattr(object.__getattribute__(self, '_orig'), name)

    61 

    62     def __delattr__(self, name):

    63         delattr(object.__getattribute__(self, '_orig'), name)

    64 

    65     def __setattr__(self, name, value):

    66         setattr(object.__getattribute__(self, '_orig'), name, value)

    67 

    68     def makefile(self, mode, bufsize):

    69         f = object.__getattribute__(self, '_orig').makefile(mode, bufsize)

    70 

    71         logfp = object.__getattribute__(self, '_logfp')

    72         closeafterrecvbytes = object.__getattribute__(self,

    73                                                       '_closeafterrecvbytes')

    74         closeaftersendbytes = object.__getattribute__(self,

    75                                                       '_closeaftersendbytes')

    76 

    77         return fileobjectproxy(f, logfp,

    78                                closeafterrecvbytes=closeafterrecvbytes,

    79                                closeaftersendbytes=closeaftersendbytes)

    80 

    81 # We can't adjust __class__ on socket._fileobject, so define a proxy.

    82 class fileobjectproxy(object):

    83     __slots__ = (

    84         '_orig',

    85         '_logfp',

    86         '_closeafterrecvbytes',

    87         '_closeaftersendbytes',

    88     )

    89 

    90     def __init__(self, obj, logfp, closeafterrecvbytes=0,

    91                  closeaftersendbytes=0):

    92         object.__setattr__(self, '_orig', obj)

    93         object.__setattr__(self, '_logfp', logfp)

    94         object.__setattr__(self, '_closeafterrecvbytes', closeafterrecvbytes)

    95         object.__setattr__(self, '_closeaftersendbytes', closeaftersendbytes)

    96 

    97     def __getattribute__(self, name):

    98         if name in ('read', 'readline', 'write', '_writelog'):

    99             return object.__getattribute__(self, name)

   100 

   101         return getattr(object.__getattribute__(self, '_orig'), name)

   102 

   103     def __delattr__(self, name):

   104         delattr(object.__getattribute__(self, '_orig'), name)

   105 

   106     def __setattr__(self, name, value):

   107         setattr(object.__getattribute__(self, '_orig'), name, value)

   108 

   109     def _writelog(self, msg):

   110         msg = msg.replace('\r', '\\r').replace('\n', '\\n')

   111 

   112         object.__getattribute__(self, '_logfp').write(msg)

   113         object.__getattribute__(self, '_logfp').write('\n')

   114 

   115     def read(self, size=-1):

   116         remaining = object.__getattribute__(self, '_closeafterrecvbytes')

   117 

   118         # No read limit. Call original function.

   119         if not remaining:

   120             result = object.__getattribute__(self, '_orig').read(size)

   121             self._writelog('read(%d) -> (%d) (%s) %s' % (size,

   122                                                            len(result),

   123                                                            result))

   124             return result

   125 

   126         origsize = size

   127 

   128         if size < 0:

   129             size = remaining

   130         else:

   131             size = min(remaining, size)

   132 

   133         result = object.__getattribute__(self, '_orig').read(size)

   134         remaining -= len(result)

   135 

   136         self._writelog('read(%d from %d) -> (%d) %s' % (

   137             size, origsize, len(result), result))

   138 

   139         object.__setattr__(self, '_closeafterrecvbytes', remaining)

   140 

   141         if remaining <= 0:

   142             self._writelog('read limit reached, closing socket')

   143             self._sock.close()

   144             # This is the easiest way to abort the current request.

   145             raise Exception('connection closed after receiving N bytes')

   146 

   147         return result

   148 

   149     def readline(self, size=-1):

   150         remaining = object.__getattribute__(self, '_closeafterrecvbytes')

   151 

   152         # No read limit. Call original function.

   153         if not remaining:

   154             result = object.__getattribute__(self, '_orig').readline(size)

   155             self._writelog('readline(%d) -> (%d) %s' % (

   156                 size, len(result), result))

   157             return result

   158 

   159         origsize = size

   160 

   161         if size < 0:

   162             size = remaining

   163         else:

   164             size = min(remaining, size)

   165 

   166         result = object.__getattribute__(self, '_orig').readline(size)

   167         remaining -= len(result)

   168 

   169         self._writelog('readline(%d from %d) -> (%d) %s' % (

   170             size, origsize, len(result), result))

   171 

   172         object.__setattr__(self, '_closeafterrecvbytes', remaining)

   173 

   174         if remaining <= 0:

   175             self._writelog('read limit reached; closing socket')

   176             self._sock.close()

   177             # This is the easiest way to abort the current request.

   178             raise Exception('connection closed after receiving N bytes')

   179 

   180         return result

   181 

   182     def write(self, data):

   183         remaining = object.__getattribute__(self, '_closeaftersendbytes')

   184 

   185         # No byte limit on this operation. Call original function.

   186         if not remaining:

   187             self._writelog('write(%d) -> %s' % (len(data), data))

   188             result = object.__getattribute__(self, '_orig').write(data)

   189             return result

   190 

   191         if len(data) > remaining:

   192             newdata = data[0:remaining]

   193         else:

   194             newdata = data

   195 

   196         remaining -= len(newdata)

   197 

   198         self._writelog('write(%d from %d) -> (%d) %s' % (

   199             len(newdata), len(data), remaining, newdata))

   200 

   201         result = object.__getattribute__(self, '_orig').write(newdata)

   202 

   203         object.__setattr__(self, '_closeaftersendbytes', remaining)

   204 

   205         if remaining <= 0:

   206             self._writelog('write limit reached; closing socket')

   207             self._sock.close()

   208             raise Exception('connection closed after sending N bytes')

   209 

   210         return result

   211 

   212 def extsetup(ui):

   213     # Change the base HTTP server class so various events can be performed.

   214     # See SocketServer.BaseServer for how the specially named methods work.

   215     class badserver(server.MercurialHTTPServer):

   216         def __init__(self, ui, *args, **kwargs):

   217             self._ui = ui

   218             super(badserver, self).__init__(ui, *args, **kwargs)

   219 

   220             # Need to inherit object so super() works.

   221             class badrequesthandler(self.RequestHandlerClass, object):

   222                 def send_header(self, name, value):

   223                     # Make headers deterministic to facilitate testing.

   224                     if name.lower() == 'date':

   225                         value = 'Fri, 14 Apr 2017 00:00:00 GMT'

   226                     elif name.lower() == 'server':

   227                         value = 'badhttpserver'

   228 

   229                     return super(badrequesthandler, self).send_header(name,

   230                                                                       value)

   231 

   232             self.RequestHandlerClass = badrequesthandler

   233 

   234         # Called to accept() a pending socket.

   235         def get_request(self):

   236             if self._ui.configbool('badserver', 'closebeforeaccept'):

   237                 self.socket.close()

   238 

   239                 # Tells the server to stop processing more requests.

   240                 self.__shutdown_request = True

   241 

   242                 # Simulate failure to stop processing this request.

   243                 raise socket.error('close before accept')

   244 

   245             if self._ui.configbool('badserver', 'closeafteraccept'):

   246                 request, client_address = super(badserver, self).get_request()

   247                 request.close()

   248                 raise socket.error('close after accept')

   249 

   250             return super(badserver, self).get_request()

   251 

   252         # Does heavy lifting of processing a request. Invokes

   253         # self.finish_request() which calls self.RequestHandlerClass() which

   254         # is a hgweb.server._httprequesthandler.

   255         def process_request(self, socket, address):

   256             # Wrap socket in a proxy if we need to count bytes.

   257             closeafterrecvbytes = self._ui.configint('badserver',

   258                                                      'closeafterrecvbytes', 0)

   259             closeaftersendbytes = self._ui.configint('badserver',

   260                                                      'closeaftersendbytes', 0)

   261 

   262             if closeafterrecvbytes or closeaftersendbytes:

   263                 socket = socketproxy(socket, self.errorlog,

   264                                      closeafterrecvbytes=closeafterrecvbytes,

   265                                      closeaftersendbytes=closeaftersendbytes)

   266 

   267             return super(badserver, self).process_request(socket, address)

   268 

   269     server.MercurialHTTPServer = badserver