def _recv(self):
""" Recv logic with interruptions """
# print "RECV! WAIT FOR LENGTH!"
packet = self.stream.read(self.FRAME_HEADER.size)
# If no packet - then just return
if not packet:
return None
header = packet
while len(header) != self.FRAME_HEADER.size:
packet = self.stream.read(self.FRAME_HEADER.size - len(header))
if packet:
header += packet
del packet
length, compressed = self.FRAME_HEADER.unpack(header)
# print "RECV! WAIT FOR LENGTH COMPLETE!"
required_length = length + len(self.FLUSHER)
# print "WAIT FOR", required_length
data = []
decompressor = None
if compressed:
decompressor = zlib.decompressobj()
buf = Buffer()
while required_length:
packet = self.stream.read(min(required_length, self.COMPRESSION_THRESHOLD))
if packet:
required_length -= len(packet)
# print "GET", len(packet)
if not required_length:
packet = packet[:-len(self.FLUSHER)]
if compressed:
packet = decompressor.decompress(packet)
if not packet:
continue
if packet:
buf.write(packet)
if compressed:
packet = decompressor.flush()
if packet:
buf.write(packet)
# print "COMPLETE!"
return buf
def upstream_recv(self, data):
if self.encryptor:
buf = Buffer()
ldata = len(data)
buf.write(self.encryptor.encrypt(struct.pack('<I', ldata)))
_, nw = data.write_to(buf, modificator=self.encryptor.encrypt, n=ldata)
d = self.update_encryptor()
buf.write(d)
buf.write_to(self.downstream)
else:
data.write_to(self.up_buffer)
class PupySocketStream(SocketStream):
def __init__(self, sock, transport_class, transport_kwargs):
super(PupySocketStream, self).__init__(sock)
self.MAX_IO_CHUNK = 32000
self.KEEP_ALIVE_REQUIRED = False
self.compress = True
#buffers for transport
self.upstream = Buffer(
transport_func=addGetPeer(("127.0.0.1", 443)),
shared=True
)
if sock is None:
peername = '127.0.0.1', 0
elif type(sock) is tuple:
peername = sock[0], sock[1]
else:
peername = sock.getpeername()
self.downstream = Buffer(
on_write=self._upstream_recv,
transport_func=addGetPeer(peername),
shared=True
)
self.upstream_lock = threading.Lock()
self.downstream_lock = threading.Lock()
self.transport = transport_class(self, **transport_kwargs)
#buffers for streams
self.buf_in = Buffer()
self.buf_out = Buffer()
self.on_connect()
def on_connect(self):
self.transport.on_connect()
self._upstream_recv()
def _read(self):
try:
buf = self.sock.recv(self.MAX_IO_CHUNK)
if __debug__:
logger.debug('stream: read={}'.format(len(buf) if buf else None))
except socket.timeout:
return
except socket.error:
ex = sys.exc_info()[1]
if get_exc_errno(ex) in (errno.EAGAIN, errno.EWOULDBLOCK):
# windows just has to be a b**ch
# edit: some politeness please ;)
return
self.close()
raise EOFError(ex)
if not buf:
self.close()
raise EOFError("connection closed by peer")
self.buf_in.write(buf)
# The root of evil
def poll(self, timeout):
if self.closed:
raise EOFError('polling on already closed connection')
result = ( len(self.upstream)>0 or self.sock_poll(timeout) )
return result
def sock_poll(self, timeout):
with self.downstream_lock:
to_close = None
to_read = None
while not (to_close or to_read or self.closed):
try:
to_read, _, to_close = select([self.sock], [], [self.sock], timeout)
except select_error as r:
if not r.args[0] == errno.EINTR:
to_close = True
continue
break
if to_close:
raise EOFError('sock_poll error')
if to_read:
self._read()
self.transport.downstream_recv(self.buf_in)
return True
else:
return False
def _upstream_recv(self):
""" called as a callback on the downstream.write """
if len(self.downstream)>0:
if __debug__:
logger.debug('stream: send={}'.format(len(self.downstream)))
self.downstream.write_to(super(PupySocketStream, self))
def waitfor(self, count):
if __debug__:
logger.debug('stream: waitfor={}'.format(count))
try:
while len(self.upstream)<count:
if not self.sock_poll(None) and self.closed:
return None
return self.upstream
except (EOFError, socket.error):
self.close()
raise
except Exception as e:
logger.debug(traceback.format_exc())
self.close()
raise
def read(self, count):
promise = self.waitfor(count)
if promise:
return promise.read(count)
def insert(self, data):
with self.upstream_lock:
self.buf_out.insert(data)
def flush(self):
self.buf_out.flush()
def write(self, data, notify=True):
if __debug__:
logger.debug('stream: write={} / n={}'.format(
len(data) if data else None, notify))
try:
with self.upstream_lock:
self.buf_out.write(data, notify)
del data
if notify:
self.transport.upstream_recv(self.buf_out)
#The write will be done by the _upstream_recv callback on the downstream buffer
except (EOFError, socket.error):
self.close()
raise
except Exception as e:
logger.debug(traceback.format_exc())
self.close()
raise
