comparison tests/testlib/badserverext.py @ 48613:b060e305d79f

test-http-bad-server: factor code dealing with "read" in the new object

This will make sure both `read` and `readline` do the same processing and make
it simpler to update that processing in the future.

Differential Revision: https://phab.mercurial-scm.org/D12044

48612:11e5cb170d36

            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.
class socketproxy(object):
    __slots__ = ('_orig', '_logfp', '_cond')

                self.close()
        else:
            self._sock.shutdown(socket.SHUT_RDWR)

    def read(self, size=-1):
        remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes

        # No read limit. Call original function.
        if not remaining:
            result = object.__getattribute__(self, '_orig').read(size)
            self._writelog(
                b'read(%d) -> (%d) (%s) %s' % (size, len(result), result)
            )
            return result

        origsize = size

        if size < 0:
            size = remaining
        else:
            size = min(remaining, size)

        result = object.__getattribute__(self, '_orig').read(size)
        remaining -= len(result)

        self._writelog(
            b'read(%d from %d) -> (%d) %s'
            % (size, origsize, len(result), result)
        )

        object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining

        if remaining <= 0:
            self._writelog(b'read limit reached; closing socket')
            self._close()

            # This is the easiest way to abort the current request.
            raise Exception('connection closed after receiving N bytes')

        return result

    def readline(self, size=-1):
        remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes

        # No read limit. Call original function.
        if not remaining:
            result = object.__getattribute__(self, '_orig').readline(size)
            self._writelog(
                b'readline(%d) -> (%d) %s' % (size, len(result), result)
            )
            return result

        origsize = size

        if size < 0:
            size = remaining
        else:
            size = min(remaining, size)

        result = object.__getattribute__(self, '_orig').readline(size)
        remaining -= len(result)

        self._writelog(
            b'readline(%d from %d) -> (%d) %s'
            % (size, origsize, len(result), result)
        )

        object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining

        if remaining <= 0:
            self._writelog(b'read limit reached; closing socket')
            self._close()

            # This is the easiest way to abort the current request.
            raise Exception('connection closed after receiving N bytes')

        return result

    def write(self, data):
        cond = object.__getattribute__(self, '_cond')
        return cond.forward_write(self, 'write', data)

48613:b060e305d79f

            object.__getattribute__(obj, '_cond_close')()
            raise Exception('connection closed after sending N bytes')

        return result

    def forward_read(self, obj, method, size=-1):
        """call an underlying read function until condition are met

        When the condition are met the socket is closed
        """
        remaining = self.remaining_recv_bytes

        orig = object.__getattribute__(obj, '_orig')
        bmethod = method.encode('ascii')
        func = getattr(orig, method)

        # No read limit. Call original function.
        if not remaining:
            result = func(size)
            obj._writelog(
                b'%s(%d) -> (%d) %s' % (bmethod, size, len(result), result)
            )
            return result

        origsize = size

        if size < 0:
            size = remaining
        else:
            size = min(remaining, size)

        result = func(size)
        remaining -= len(result)

        obj._writelog(
            b'%s(%d from %d) -> (%d) %s'
            % (bmethod, size, origsize, len(result), result)
        )

        self.remaining_recv_bytes = remaining

        if remaining <= 0:
            obj._writelog(b'read limit reached; closing socket')
            obj._cond_close()

            # This is the easiest way to abort the current request.
            raise Exception('connection closed after receiving N bytes')

        return result


# We can't adjust __class__ on a socket instance. So we define a proxy type.
class socketproxy(object):
    __slots__ = ('_orig', '_logfp', '_cond')

                self.close()
        else:
            self._sock.shutdown(socket.SHUT_RDWR)

    def read(self, size=-1):
        cond = object.__getattribute__(self, '_cond')
        return cond.forward_read(self, 'read', size)

    def readline(self, size=-1):
        cond = object.__getattribute__(self, '_cond')
        return cond.forward_read(self, 'readline', size)

    def write(self, data):
        cond = object.__getattribute__(self, '_cond')
        return cond.forward_write(self, 'write', data)