diff options
author | Aldo Cortesi <aldo@nullcube.com> | 2011-01-27 10:52:42 +1300 |
---|---|---|
committer | Aldo Cortesi <aldo@nullcube.com> | 2011-01-27 10:52:42 +1300 |
commit | 077272ec9703ff1b3f20f69c971adf63a6dd05c1 (patch) | |
tree | c45e28fc09abaecd8f968f73f94ac73f8deec99d /libmproxy/contrib/bson/network.py | |
parent | 460107589c2ef425ef124c74908be1fa86ebab9d (diff) | |
download | mitmproxy-077272ec9703ff1b3f20f69c971adf63a6dd05c1.tar.gz mitmproxy-077272ec9703ff1b3f20f69c971adf63a6dd05c1.tar.bz2 mitmproxy-077272ec9703ff1b3f20f69c971adf63a6dd05c1.zip |
Switch to BSON for data serialization.
Diffstat (limited to 'libmproxy/contrib/bson/network.py')
-rw-r--r-- | libmproxy/contrib/bson/network.py | 64 |
1 files changed, 64 insertions, 0 deletions
diff --git a/libmproxy/contrib/bson/network.py b/libmproxy/contrib/bson/network.py new file mode 100644 index 00000000..f9a90d48 --- /dev/null +++ b/libmproxy/contrib/bson/network.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python + +import socket +try: + from cStringIO import StringIO +except ImportError, e: + from StringIO import StringIO +from struct import unpack +from __init__ import dumps, loads + +def _bintoint(data): + return unpack("<i", data)[0] + +def _sendobj(self, obj): + """ + Atomically send a BSON message. + """ + data = dumps(obj) + self.sendall(data) + +def _recvobj(self): + """ + Atomic read of a BSON message. + + This function either returns a dict, None, or raises a socket error. + + If the return value is None, it means the socket is closed by the other side. + """ + sock_buf = self.recvbytes(4) + if sock_buf is None: + return None + + message_length = _bintoint(sock_buf.getvalue()) + sock_buf = self.recvbytes(message_length - 4, sock_buf) + if sock_buf is None: + return None + + retval = loads(sock_buf.getvalue()) + return retval + + +def _recvbytes(self, bytes_needed, sock_buf = None): + """ + Atomic read of bytes_needed bytes. + + This function either returns exactly the nmber of bytes requested in a + StringIO buffer, None, or raises a socket error. + + If the return value is None, it means the socket is closed by the other side. + """ + if sock_buf is None: + sock_buf = StringIO() + bytes_count = 0 + while bytes_count < bytes_needed: + chunk = self.recv(min(bytes_needed - bytes_count, 32768)) + part_count = len(chunk) + + if part_count < 1: + return None + + bytes_count += part_count + sock_buf.write(chunk) + + return sock_buf |