From f2bc58cdd2f2b9b0025a88c0faccf55e10b29353 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 5 May 2015 10:47:02 +1200 Subject: Add tcp.Reader.safe_read, use it in socks and websockets safe_read is guaranteed to raise or return a byte string of the requested length. It's particularly useful for implementing binary protocols. --- netlib/socks.py | 32 +++++++++----------------------- 1 file changed, 9 insertions(+), 23 deletions(-) (limited to 'netlib/socks.py') diff --git a/netlib/socks.py b/netlib/socks.py index 497b8eef..6f9f57bd 100644 --- a/netlib/socks.py +++ b/netlib/socks.py @@ -52,20 +52,6 @@ METHOD = utils.BiDi( ) -def _read(f, n): - try: - d = f.read(n) - if len(d) == n: - return d - else: - raise SocksError( - REP.GENERAL_SOCKS_SERVER_FAILURE, - "Incomplete Read" - ) - except socket.error as e: - raise SocksError(REP.GENERAL_SOCKS_SERVER_FAILURE, str(e)) - - class ClientGreeting(object): __slots__ = ("ver", "methods") @@ -75,9 +61,9 @@ class ClientGreeting(object): @classmethod def from_file(cls, f): - ver, nmethods = struct.unpack("!BB", _read(f, 2)) + ver, nmethods = struct.unpack("!BB", f.safe_read(2)) methods = array.array("B") - methods.fromstring(_read(f, nmethods)) + methods.fromstring(f.safe_read(nmethods)) return cls(ver, methods) def to_file(self, f): @@ -94,7 +80,7 @@ class ServerGreeting(object): @classmethod def from_file(cls, f): - ver, method = struct.unpack("!BB", _read(f, 2)) + ver, method = struct.unpack("!BB", f.safe_read(2)) return cls(ver, method) def to_file(self, f): @@ -112,27 +98,27 @@ class Message(object): @classmethod def from_file(cls, f): - ver, msg, rsv, atyp = struct.unpack("!BBBB", _read(f, 4)) + ver, msg, rsv, atyp = struct.unpack("!BBBB", f.safe_read(4)) if rsv != 0x00: raise SocksError(REP.GENERAL_SOCKS_SERVER_FAILURE, "Socks Request: Invalid reserved byte: %s" % rsv) if atyp == ATYP.IPV4_ADDRESS: # We use tnoa here as ntop is not commonly available on Windows. - host = socket.inet_ntoa(_read(f, 4)) + host = socket.inet_ntoa(f.safe_read(4)) use_ipv6 = False elif atyp == ATYP.IPV6_ADDRESS: - host = socket.inet_ntop(socket.AF_INET6, _read(f, 16)) + host = socket.inet_ntop(socket.AF_INET6, f.safe_read(16)) use_ipv6 = True elif atyp == ATYP.DOMAINNAME: - length, = struct.unpack("!B", _read(f, 1)) - host = _read(f, length) + length, = struct.unpack("!B", f.safe_read(1)) + host = f.safe_read(length) use_ipv6 = False else: raise SocksError(REP.ADDRESS_TYPE_NOT_SUPPORTED, "Socks Request: Unknown ATYP: %s" % atyp) - port, = struct.unpack("!H", _read(f, 2)) + port, = struct.unpack("!H", f.safe_read(2)) addr = tcp.Address((host, port), use_ipv6=use_ipv6) return cls(ver, msg, atyp, addr) -- cgit v1.2.3