changeset 37010:8453699a1f21

util: observable proxy objects for sockets

We previously introduced proxy objects and observers for file objects
to help implement low-level tests for the SSH wire protocol.

In this commit, we do the same for sockets in order to help test the
HTTP server.

We only proxy/observe some socket methods. I didn't feel like
implementing all the methods because there are so many of them and
implementing them will provide no short term value. We can always
implement them later.

# no-check-commit because we implement foo_bar methods on stdlib types

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

--- a/mercurial/util.py	Tue Mar 13 09:45:56 2018 -0700
+++ b/mercurial/util.py	Tue Mar 13 17:42:00 2018 -0700
@@ -689,6 +689,125 @@
 
         return res
 
+PROXIED_SOCKET_METHODS = {
+    r'makefile',
+    r'recv',
+    r'recvfrom',
+    r'recvfrom_into',
+    r'recv_into',
+    r'send',
+    r'sendall',
+    r'sendto',
+    r'setblocking',
+    r'settimeout',
+    r'gettimeout',
+    r'setsockopt',
+}
+
+class socketproxy(object):
+    """A proxy around a socket that tells a watcher when events occur.
+
+    This is like ``fileobjectproxy`` except for sockets.
+
+    This type is intended to only be used for testing purposes. Think hard
+    before using it in important code.
+    """
+    __slots__ = (
+        r'_orig',
+        r'_observer',
+    )
+
+    def __init__(self, sock, observer):
+        object.__setattr__(self, r'_orig', sock)
+        object.__setattr__(self, r'_observer', observer)
+
+    def __getattribute__(self, name):
+        if name in PROXIED_SOCKET_METHODS:
+            return object.__getattribute__(self, name)
+
+        return getattr(object.__getattribute__(self, r'_orig'), name)
+
+    def __delattr__(self, name):
+        return delattr(object.__getattribute__(self, r'_orig'), name)
+
+    def __setattr__(self, name, value):
+        return setattr(object.__getattribute__(self, r'_orig'), name, value)
+
+    def __nonzero__(self):
+        return bool(object.__getattribute__(self, r'_orig'))
+
+    __bool__ = __nonzero__
+
+    def _observedcall(self, name, *args, **kwargs):
+        # Call the original object.
+        orig = object.__getattribute__(self, r'_orig')
+        res = getattr(orig, name)(*args, **kwargs)
+
+        # Call a method on the observer of the same name with arguments
+        # so it can react, log, etc.
+        observer = object.__getattribute__(self, r'_observer')
+        fn = getattr(observer, name, None)
+        if fn:
+            fn(res, *args, **kwargs)
+
+        return res
+
+    def makefile(self, *args, **kwargs):
+        res = object.__getattribute__(self, r'_observedcall')(
+            r'makefile', *args, **kwargs)
+
+        # The file object may be used for I/O. So we turn it into a
+        # proxy using our observer.
+        observer = object.__getattribute__(self, r'_observer')
+        return makeloggingfileobject(observer.fh, res, observer.name,
+                                     reads=observer.reads,
+                                     writes=observer.writes,
+                                     logdata=observer.logdata)
+
+    def recv(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'recv', *args, **kwargs)
+
+    def recvfrom(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'recvfrom', *args, **kwargs)
+
+    def recvfrom_into(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'recvfrom_into', *args, **kwargs)
+
+    def recv_into(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'recv_info', *args, **kwargs)
+
+    def send(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'send', *args, **kwargs)
+
+    def sendall(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'sendall', *args, **kwargs)
+
+    def sendto(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'sendto', *args, **kwargs)
+
+    def setblocking(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'setblocking', *args, **kwargs)
+
+    def settimeout(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'settimeout', *args, **kwargs)
+
+    def gettimeout(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'gettimeout', *args, **kwargs)
+
+    def setsockopt(self, *args, **kwargs):
+        return object.__getattribute__(self, r'_observedcall')(
+            r'setsockopt', *args, **kwargs)
+
 DATA_ESCAPE_MAP = {pycompat.bytechr(i): br'\x%02x' % i for i in range(256)}
 DATA_ESCAPE_MAP.update({
     b'\\': b'\\\\',
@@ -703,15 +822,7 @@
 
     return DATA_ESCAPE_RE.sub(lambda m: DATA_ESCAPE_MAP[m.group(0)], s)
 
-class fileobjectobserver(object):
-    """Logs file object activity."""
-    def __init__(self, fh, name, reads=True, writes=True, logdata=False):
-        self.fh = fh
-        self.name = name
-        self.logdata = logdata
-        self.reads = reads
-        self.writes = writes
-
+class baseproxyobserver(object):
     def _writedata(self, data):
         if not self.logdata:
             self.fh.write('\n')
@@ -731,6 +842,15 @@
             self.fh.write('%s>     %s\n' % (self.name, escapedata(line)))
         self.fh.flush()
 
+class fileobjectobserver(baseproxyobserver):
+    """Logs file object activity."""
+    def __init__(self, fh, name, reads=True, writes=True, logdata=False):
+        self.fh = fh
+        self.name = name
+        self.logdata = logdata
+        self.reads = reads
+        self.writes = writes
+
     def read(self, res, size=-1):
         if not self.reads:
             return
@@ -793,6 +913,119 @@
                                   logdata=logdata)
     return fileobjectproxy(fh, observer)
 
+class socketobserver(baseproxyobserver):
+    """Logs socket activity."""
+    def __init__(self, fh, name, reads=True, writes=True, states=True,
+                 logdata=False):
+        self.fh = fh
+        self.name = name
+        self.reads = reads
+        self.writes = writes
+        self.states = states
+        self.logdata = logdata
+
+    def makefile(self, res, mode=None, bufsize=None):
+        if not self.states:
+            return
+
+        self.fh.write('%s> makefile(%r, %r)\n' % (
+            self.name, mode, bufsize))
+
+    def recv(self, res, size, flags=0):
+        if not self.reads:
+            return
+
+        self.fh.write('%s> recv(%d, %d) -> %d' % (
+            self.name, size, flags, len(res)))
+        self._writedata(res)
+
+    def recvfrom(self, res, size, flags=0):
+        if not self.reads:
+            return
+
+        self.fh.write('%s> recvfrom(%d, %d) -> %d' % (
+            self.name, size, flags, len(res[0])))
+        self._writedata(res[0])
+
+    def recvfrom_into(self, res, buf, size, flags=0):
+        if not self.reads:
+            return
+
+        self.fh.write('%s> recvfrom_into(%d, %d) -> %d' % (
+            self.name, size, flags, res[0]))
+        self._writedata(buf[0:res[0]])
+
+    def recv_into(self, res, buf, size=0, flags=0):
+        if not self.reads:
+            return
+
+        self.fh.write('%s> recv_into(%d, %d) -> %d' % (
+            self.name, size, flags, res))
+        self._writedata(buf[0:res])
+
+    def send(self, res, data, flags=0):
+        if not self.writes:
+            return
+
+        self.fh.write('%s> send(%d, %d) -> %d' % (
+            self.name, len(data), flags, len(res)))
+        self._writedata(data)
+
+    def sendall(self, res, data, flags=0):
+        if not self.writes:
+            return
+
+        # Returns None on success. So don't bother reporting return value.
+        self.fh.write('%s> sendall(%d, %d)' % (
+            self.name, len(data), flags))
+        self._writedata(data)
+
+    def sendto(self, res, data, flagsoraddress, address=None):
+        if not self.writes:
+            return
+
+        if address:
+            flags = flagsoraddress
+        else:
+            flags = 0
+
+        self.fh.write('%s> sendto(%d, %d, %r) -> %d' % (
+            self.name, len(data), flags, address, res))
+        self._writedata(data)
+
+    def setblocking(self, res, flag):
+        if not self.states:
+            return
+
+        self.fh.write('%s> setblocking(%r)\n' % (self.name, flag))
+
+    def settimeout(self, res, value):
+        if not self.states:
+            return
+
+        self.fh.write('%s> settimeout(%r)\n' % (self.name, value))
+
+    def gettimeout(self, res):
+        if not self.states:
+            return
+
+        self.fh.write('%s> gettimeout() -> %f\n' % (self.name, res))
+
+    def setsockopt(self, level, optname, value):
+        if not self.states:
+            return
+
+        self.fh.write('%s> setsockopt(%r, %r, %r) -> %r\n' % (
+            self.name, level, optname, value))
+
+def makeloggingsocket(logh, fh, name, reads=True, writes=True, states=True,
+                      logdata=False):
+    """Turn a socket into a logging socket."""
+
+    observer = socketobserver(logh, name, reads=reads, writes=writes,
+                              states=states, logdata=logdata)
+    return socketproxy(fh, observer)
+
 def version():
     """Return version information if available."""
     try: