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 import (

    37     pycompat,

    38     registrar,

    39 )

    40 

    41 from mercurial.hgweb import server

    42 

    43 configtable = {}

    44 configitem = registrar.configitem(configtable)

    45 

    46 configitem(

    47     b'badserver',

    48     b'closeafteraccept',

    49     default=False,

    50 )

    51 configitem(

    52     b'badserver',

    53     b'closeafterrecvbytes',

    54     default=b'0',

    55 )

    56 configitem(

    57     b'badserver',

    58     b'closeaftersendbytes',

    59     default=b'0',

    60 )

    61 configitem(

    62     b'badserver',

    63     b'closebeforeaccept',

    64     default=False,

    65 )

    66 

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

    68 class socketproxy(object):

    69     __slots__ = (

    70         '_orig',

    71         '_logfp',

    72         '_closeafterrecvbytes',

    73         '_closeaftersendbytes',

    74     )

    75 

    76     def __init__(

    77         self, obj, logfp, closeafterrecvbytes=0, closeaftersendbytes=0

    78     ):

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

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

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

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

    83 

    84     def __getattribute__(self, name):

    85         if name in ('makefile', 'sendall', '_writelog'):

    86             return object.__getattribute__(self, name)

    87 

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

    89 

    90     def __delattr__(self, name):

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

    92 

    93     def __setattr__(self, name, value):

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

    95 

    96     def _writelog(self, msg):

    97         msg = msg.replace(b'\r', b'\\r').replace(b'\n', b'\\n')

    98 

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

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

   101         object.__getattribute__(self, '_logfp').flush()

   102 

   103     def makefile(self, mode, bufsize):

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

   105 

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

   107         closeafterrecvbytes = object.__getattribute__(

   108             self, '_closeafterrecvbytes'

   109         )

   110         closeaftersendbytes = object.__getattribute__(

   111             self, '_closeaftersendbytes'

   112         )

   113 

   114         return fileobjectproxy(

   115             f,

   116             logfp,

   117             closeafterrecvbytes=closeafterrecvbytes,

   118             closeaftersendbytes=closeaftersendbytes,

   119         )

   120 

   121     def sendall(self, data, flags=0):

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

   123 

   124         # No read limit. Call original function.

   125         if not remaining:

   126             result = object.__getattribute__(self, '_orig').sendall(data, flags)

   127             self._writelog(b'sendall(%d) -> %s' % (len(data), data))

   128             return result

   129 

   130         if len(data) > remaining:

   131             newdata = data[0:remaining]

   132         else:

   133             newdata = data

   134 

   135         remaining -= len(newdata)

   136 

   137         result = object.__getattribute__(self, '_orig').sendall(newdata, flags)

   138 

   139         self._writelog(

   140             b'sendall(%d from %d) -> (%d) %s'

   141             % (len(newdata), len(data), remaining, newdata)

   142         )

   143 

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

   145 

   146         if remaining <= 0:

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

   148             object.__getattribute__(self, '_orig').shutdown(socket.SHUT_RDWR)

   149 

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

   151 

   152         return result

   153 

   154 

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

   156 class fileobjectproxy(object):

   157     __slots__ = (

   158         '_orig',

   159         '_logfp',

   160         '_closeafterrecvbytes',

   161         '_closeaftersendbytes',

   162     )

   163 

   164     def __init__(

   165         self, obj, logfp, closeafterrecvbytes=0, closeaftersendbytes=0

   166     ):

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

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

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

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

   171 

   172     def __getattribute__(self, name):

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

   174             return object.__getattribute__(self, name)

   175 

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

   177 

   178     def __delattr__(self, name):

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

   180 

   181     def __setattr__(self, name, value):

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

   183 

   184     def _writelog(self, msg):

   185         msg = msg.replace(b'\r', b'\\r').replace(b'\n', b'\\n')

   186 

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

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

   189         object.__getattribute__(self, '_logfp').flush()

   190 

   191     def _close(self):

   192         # Python 3 uses an io.BufferedIO instance. Python 2 uses some file

   193         # object wrapper.

   194         if pycompat.ispy3:

   195             orig = object.__getattribute__(self, '_orig')

   196 

   197             if hasattr(orig, 'raw'):

   198                 orig.raw._sock.shutdown(socket.SHUT_RDWR)

   199             else:

   200                 self.close()

   201         else:

   202             self._sock.shutdown(socket.SHUT_RDWR)

   203 

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

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

   206 

   207         # No read limit. Call original function.

   208         if not remaining:

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

   210             self._writelog(

   211                 b'read(%d) -> (%d) (%s) %s' % (size, len(result), result)

   212             )

   213             return result

   214 

   215         origsize = size

   216 

   217         if size < 0:

   218             size = remaining

   219         else:

   220             size = min(remaining, size)

   221 

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

   223         remaining -= len(result)

   224 

   225         self._writelog(

   226             b'read(%d from %d) -> (%d) %s'

   227             % (size, origsize, len(result), result)

   228         )

   229 

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

   231 

   232         if remaining <= 0:

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

   234             self._close()

   235 

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

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

   238 

   239         return result

   240 

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

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

   243 

   244         # No read limit. Call original function.

   245         if not remaining:

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

   247             self._writelog(

   248                 b'readline(%d) -> (%d) %s' % (size, len(result), result)

   249             )

   250             return result

   251 

   252         origsize = size

   253 

   254         if size < 0:

   255             size = remaining

   256         else:

   257             size = min(remaining, size)

   258 

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

   260         remaining -= len(result)

   261 

   262         self._writelog(

   263             b'readline(%d from %d) -> (%d) %s'

   264             % (size, origsize, len(result), result)

   265         )

   266 

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

   268 

   269         if remaining <= 0:

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

   271             self._close()

   272 

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

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

   275 

   276         return result

   277 

   278     def write(self, data):

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

   280 

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

   282         if not remaining:

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

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

   285             return result

   286 

   287         if len(data) > remaining:

   288             newdata = data[0:remaining]

   289         else:

   290             newdata = data

   291 

   292         remaining -= len(newdata)

   293 

   294         self._writelog(

   295             b'write(%d from %d) -> (%d) %s'

   296             % (len(newdata), len(data), remaining, newdata)

   297         )

   298 

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

   300 

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

   302 

   303         if remaining <= 0:

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

   305             self._close()

   306 

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

   308 

   309         return result

   310 

   311 

   312 def extsetup(ui):

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

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

   315     class badserver(server.MercurialHTTPServer):

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

   317             self._ui = ui

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

   319 

   320             recvbytes = self._ui.config(b'badserver', b'closeafterrecvbytes')

   321             recvbytes = recvbytes.split(b',')

   322             self.closeafterrecvbytes = [int(v) for v in recvbytes if v]

   323             sendbytes = self._ui.config(b'badserver', b'closeaftersendbytes')

   324             sendbytes = sendbytes.split(b',')

   325             self.closeaftersendbytes = [int(v) for v in sendbytes if v]

   326 

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

   328             class badrequesthandler(self.RequestHandlerClass, object):

   329                 def send_header(self, name, value):

   330                     # Make headers deterministic to facilitate testing.

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

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

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

   334                         value = 'badhttpserver'

   335 

   336                     return super(badrequesthandler, self).send_header(

   337                         name, value

   338                     )

   339 

   340             self.RequestHandlerClass = badrequesthandler

   341 

   342         # Called to accept() a pending socket.

   343         def get_request(self):

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

   345                 self.socket.close()

   346 

   347                 # Tells the server to stop processing more requests.

   348                 self.__shutdown_request = True

   349 

   350                 # Simulate failure to stop processing this request.

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

   352 

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

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

   355                 request.close()

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

   357 

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

   359 

   360         # Does heavy lifting of processing a request. Invokes

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

   362         # is a hgweb.server._httprequesthandler.

   363         def process_request(self, socket, address):

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

   365             if self.closeafterrecvbytes:

   366                 closeafterrecvbytes = self.closeafterrecvbytes.pop(0)

   367             else:

   368                 closeafterrecvbytes = 0

   369             if self.closeaftersendbytes:

   370                 closeaftersendbytes = self.closeaftersendbytes.pop(0)

   371             else:

   372                 closeaftersendbytes = 0

   373 

   374             if closeafterrecvbytes or closeaftersendbytes:

   375                 socket = socketproxy(

   376                     socket,

   377                     self.errorlog,

   378                     closeafterrecvbytes=closeafterrecvbytes,

   379                     closeaftersendbytes=closeaftersendbytes,

   380                 )

   381 

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

   383 

   384     server.MercurialHTTPServer = badserver