comparison mercurial/util.py @ 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
author Gregory Szorc <gregory.szorc@gmail.com>
date Tue, 13 Mar 2018 17:42:00 -0700
parents d683c7367989
children d3a9036d9ae9
comparison
equal deleted inserted replaced
37009:5890e5872f36 37010:8453699a1f21
687 if fn: 687 if fn:
688 fn(res) 688 fn(res)
689 689
690 return res 690 return res
691 691
692 PROXIED_SOCKET_METHODS = {
693 r'makefile',
694 r'recv',
695 r'recvfrom',
696 r'recvfrom_into',
697 r'recv_into',
698 r'send',
699 r'sendall',
700 r'sendto',
701 r'setblocking',
702 r'settimeout',
703 r'gettimeout',
704 r'setsockopt',
705 }
706
707 class socketproxy(object):
708 """A proxy around a socket that tells a watcher when events occur.
709
710 This is like ``fileobjectproxy`` except for sockets.
711
712 This type is intended to only be used for testing purposes. Think hard
713 before using it in important code.
714 """
715 __slots__ = (
716 r'_orig',
717 r'_observer',
718 )
719
720 def __init__(self, sock, observer):
721 object.__setattr__(self, r'_orig', sock)
722 object.__setattr__(self, r'_observer', observer)
723
724 def __getattribute__(self, name):
725 if name in PROXIED_SOCKET_METHODS:
726 return object.__getattribute__(self, name)
727
728 return getattr(object.__getattribute__(self, r'_orig'), name)
729
730 def __delattr__(self, name):
731 return delattr(object.__getattribute__(self, r'_orig'), name)
732
733 def __setattr__(self, name, value):
734 return setattr(object.__getattribute__(self, r'_orig'), name, value)
735
736 def __nonzero__(self):
737 return bool(object.__getattribute__(self, r'_orig'))
738
739 __bool__ = __nonzero__
740
741 def _observedcall(self, name, *args, **kwargs):
742 # Call the original object.
743 orig = object.__getattribute__(self, r'_orig')
744 res = getattr(orig, name)(*args, **kwargs)
745
746 # Call a method on the observer of the same name with arguments
747 # so it can react, log, etc.
748 observer = object.__getattribute__(self, r'_observer')
749 fn = getattr(observer, name, None)
750 if fn:
751 fn(res, *args, **kwargs)
752
753 return res
754
755 def makefile(self, *args, **kwargs):
756 res = object.__getattribute__(self, r'_observedcall')(
757 r'makefile', *args, **kwargs)
758
759 # The file object may be used for I/O. So we turn it into a
760 # proxy using our observer.
761 observer = object.__getattribute__(self, r'_observer')
762 return makeloggingfileobject(observer.fh, res, observer.name,
763 reads=observer.reads,
764 writes=observer.writes,
765 logdata=observer.logdata)
766
767 def recv(self, *args, **kwargs):
768 return object.__getattribute__(self, r'_observedcall')(
769 r'recv', *args, **kwargs)
770
771 def recvfrom(self, *args, **kwargs):
772 return object.__getattribute__(self, r'_observedcall')(
773 r'recvfrom', *args, **kwargs)
774
775 def recvfrom_into(self, *args, **kwargs):
776 return object.__getattribute__(self, r'_observedcall')(
777 r'recvfrom_into', *args, **kwargs)
778
779 def recv_into(self, *args, **kwargs):
780 return object.__getattribute__(self, r'_observedcall')(
781 r'recv_info', *args, **kwargs)
782
783 def send(self, *args, **kwargs):
784 return object.__getattribute__(self, r'_observedcall')(
785 r'send', *args, **kwargs)
786
787 def sendall(self, *args, **kwargs):
788 return object.__getattribute__(self, r'_observedcall')(
789 r'sendall', *args, **kwargs)
790
791 def sendto(self, *args, **kwargs):
792 return object.__getattribute__(self, r'_observedcall')(
793 r'sendto', *args, **kwargs)
794
795 def setblocking(self, *args, **kwargs):
796 return object.__getattribute__(self, r'_observedcall')(
797 r'setblocking', *args, **kwargs)
798
799 def settimeout(self, *args, **kwargs):
800 return object.__getattribute__(self, r'_observedcall')(
801 r'settimeout', *args, **kwargs)
802
803 def gettimeout(self, *args, **kwargs):
804 return object.__getattribute__(self, r'_observedcall')(
805 r'gettimeout', *args, **kwargs)
806
807 def setsockopt(self, *args, **kwargs):
808 return object.__getattribute__(self, r'_observedcall')(
809 r'setsockopt', *args, **kwargs)
810
692 DATA_ESCAPE_MAP = {pycompat.bytechr(i): br'\x%02x' % i for i in range(256)} 811 DATA_ESCAPE_MAP = {pycompat.bytechr(i): br'\x%02x' % i for i in range(256)}
693 DATA_ESCAPE_MAP.update({ 812 DATA_ESCAPE_MAP.update({
694 b'\\': b'\\\\', 813 b'\\': b'\\\\',
695 b'\r': br'\r', 814 b'\r': br'\r',
696 b'\n': br'\n', 815 b'\n': br'\n',
701 if isinstance(s, bytearray): 820 if isinstance(s, bytearray):
702 s = bytes(s) 821 s = bytes(s)
703 822
704 return DATA_ESCAPE_RE.sub(lambda m: DATA_ESCAPE_MAP[m.group(0)], s) 823 return DATA_ESCAPE_RE.sub(lambda m: DATA_ESCAPE_MAP[m.group(0)], s)
705 824
706 class fileobjectobserver(object): 825 class baseproxyobserver(object):
826 def _writedata(self, data):
827 if not self.logdata:
828 self.fh.write('\n')
829 self.fh.flush()
830 return
831
832 # Simple case writes all data on a single line.
833 if b'\n' not in data:
834 self.fh.write(': %s\n' % escapedata(data))
835 self.fh.flush()
836 return
837
838 # Data with newlines is written to multiple lines.
839 self.fh.write(':\n')
840 lines = data.splitlines(True)
841 for line in lines:
842 self.fh.write('%s> %s\n' % (self.name, escapedata(line)))
843 self.fh.flush()
844
845 class fileobjectobserver(baseproxyobserver):
707 """Logs file object activity.""" 846 """Logs file object activity."""
708 def __init__(self, fh, name, reads=True, writes=True, logdata=False): 847 def __init__(self, fh, name, reads=True, writes=True, logdata=False):
709 self.fh = fh 848 self.fh = fh
710 self.name = name 849 self.name = name
711 self.logdata = logdata 850 self.logdata = logdata
712 self.reads = reads 851 self.reads = reads
713 self.writes = writes 852 self.writes = writes
714 853
715 def _writedata(self, data):
716 if not self.logdata:
717 self.fh.write('\n')
718 self.fh.flush()
719 return
720
721 # Simple case writes all data on a single line.
722 if b'\n' not in data:
723 self.fh.write(': %s\n' % escapedata(data))
724 self.fh.flush()
725 return
726
727 # Data with newlines is written to multiple lines.
728 self.fh.write(':\n')
729 lines = data.splitlines(True)
730 for line in lines:
731 self.fh.write('%s> %s\n' % (self.name, escapedata(line)))
732 self.fh.flush()
733
734 def read(self, res, size=-1): 854 def read(self, res, size=-1):
735 if not self.reads: 855 if not self.reads:
736 return 856 return
737 # Python 3 can return None from reads at EOF instead of empty strings. 857 # Python 3 can return None from reads at EOF instead of empty strings.
738 if res is None: 858 if res is None:
790 """Turn a file object into a logging file object.""" 910 """Turn a file object into a logging file object."""
791 911
792 observer = fileobjectobserver(logh, name, reads=reads, writes=writes, 912 observer = fileobjectobserver(logh, name, reads=reads, writes=writes,
793 logdata=logdata) 913 logdata=logdata)
794 return fileobjectproxy(fh, observer) 914 return fileobjectproxy(fh, observer)
915
916 class socketobserver(baseproxyobserver):
917 """Logs socket activity."""
918 def __init__(self, fh, name, reads=True, writes=True, states=True,
919 logdata=False):
920 self.fh = fh
921 self.name = name
922 self.reads = reads
923 self.writes = writes
924 self.states = states
925 self.logdata = logdata
926
927 def makefile(self, res, mode=None, bufsize=None):
928 if not self.states:
929 return
930
931 self.fh.write('%s> makefile(%r, %r)\n' % (
932 self.name, mode, bufsize))
933
934 def recv(self, res, size, flags=0):
935 if not self.reads:
936 return
937
938 self.fh.write('%s> recv(%d, %d) -> %d' % (
939 self.name, size, flags, len(res)))
940 self._writedata(res)
941
942 def recvfrom(self, res, size, flags=0):
943 if not self.reads:
944 return
945
946 self.fh.write('%s> recvfrom(%d, %d) -> %d' % (
947 self.name, size, flags, len(res[0])))
948 self._writedata(res[0])
949
950 def recvfrom_into(self, res, buf, size, flags=0):
951 if not self.reads:
952 return
953
954 self.fh.write('%s> recvfrom_into(%d, %d) -> %d' % (
955 self.name, size, flags, res[0]))
956 self._writedata(buf[0:res[0]])
957
958 def recv_into(self, res, buf, size=0, flags=0):
959 if not self.reads:
960 return
961
962 self.fh.write('%s> recv_into(%d, %d) -> %d' % (
963 self.name, size, flags, res))
964 self._writedata(buf[0:res])
965
966 def send(self, res, data, flags=0):
967 if not self.writes:
968 return
969
970 self.fh.write('%s> send(%d, %d) -> %d' % (
971 self.name, len(data), flags, len(res)))
972 self._writedata(data)
973
974 def sendall(self, res, data, flags=0):
975 if not self.writes:
976 return
977
978 # Returns None on success. So don't bother reporting return value.
979 self.fh.write('%s> sendall(%d, %d)' % (
980 self.name, len(data), flags))
981 self._writedata(data)
982
983 def sendto(self, res, data, flagsoraddress, address=None):
984 if not self.writes:
985 return
986
987 if address:
988 flags = flagsoraddress
989 else:
990 flags = 0
991
992 self.fh.write('%s> sendto(%d, %d, %r) -> %d' % (
993 self.name, len(data), flags, address, res))
994 self._writedata(data)
995
996 def setblocking(self, res, flag):
997 if not self.states:
998 return
999
1000 self.fh.write('%s> setblocking(%r)\n' % (self.name, flag))
1001
1002 def settimeout(self, res, value):
1003 if not self.states:
1004 return
1005
1006 self.fh.write('%s> settimeout(%r)\n' % (self.name, value))
1007
1008 def gettimeout(self, res):
1009 if not self.states:
1010 return
1011
1012 self.fh.write('%s> gettimeout() -> %f\n' % (self.name, res))
1013
1014 def setsockopt(self, level, optname, value):
1015 if not self.states:
1016 return
1017
1018 self.fh.write('%s> setsockopt(%r, %r, %r) -> %r\n' % (
1019 self.name, level, optname, value))
1020
1021 def makeloggingsocket(logh, fh, name, reads=True, writes=True, states=True,
1022 logdata=False):
1023 """Turn a socket into a logging socket."""
1024
1025 observer = socketobserver(logh, name, reads=reads, writes=writes,
1026 states=states, logdata=logdata)
1027 return socketproxy(fh, observer)
795 1028
796 def version(): 1029 def version():
797 """Return version information if available.""" 1030 """Return version information if available."""
798 try: 1031 try:
799 from . import __version__ 1032 from . import __version__