class Enity(object):
def __init__(self, c, addr, opts):
self.conn = c
self.key = opts.key
self.secret = opts.secret
self.addr_proxy = addr
self.conn_proxy = Connection()
c = self.conn_proxy
if not c.connect_nonblocking(addr):
self.close()
return
c.wev.handler = lambda: self.send_connect()
add_conn(c, WRITE_EVENT)
self.timer = Timer()
self.timer.handler = lambda: self.on_timeout()
add_timer(self.timer, EXPIRE_TIME)
def on_timeout(self):
c = self.conn_proxy
log.warn(0, '*%d timeout', c.index)
self.close()
def close(self):
if self.conn:
self.conn.close()
self.conn = None
if self.conn_proxy:
self.conn_proxy.close()
self.conn_proxy = None
if self.timer:
del_timer(self.timer)
self.timer = None
def send_connect(self):
c = self.conn_proxy
log.debug(0, '*%d connect: %s', c.index, c.addr.text)
timestamp = int(time.time())
rand = '%x_%d' % (random.randint(0, 0xffffffff), get_sequence())
sign = auth.get_sign(self.secret, [timestamp, rand])
msg = Message(['connect req', self.key, timestamp, rand, sign])
log.debug(0, '*%d send message: %s', c.index, msg)
buff = msg.encode()
if not buff:
log.error(0, 'invalid message: %s', msg)
return
c.wev.buff = struct.pack('I', len(buff)) + buff
c.wev.handler = lambda: self.send()
self.send()
def send(self):
c = self.conn_proxy
buff = c.wev.buff
if len(buff) > 0:
r, size = c.send(buff)
if r != 0:
if r == 1:
log.debug(0, '*%d closed', c.index)
self.close()
return
c.wev.buff = buff[size:]
buff = c.wev.buff
if len(buff) > 0:
c.wev.handler = lambda: self.send()
del_conn(c, READ_EVENT)
add_conn(c, WRITE_EVENT)
else:
c.rev.handler = lambda: self.read_connect()
del_conn(c, WRITE_EVENT)
add_conn(c, READ_EVENT)
def read_bin(self):
c = self.conn_proxy
r, buff = c.recv(4096)
if r != 0:
if r == 1:
log.debug(0, '*%d closed', self.conn.index)
self.close()
return ''
c.rev.buff += buff
buff = c.rev.buff
size0 = struct.calcsize('I')
if len(buff) < size0:
return ''
size1, = struct.unpack('I', buff[0:size0])
if len(buff) < size1 + size0:
return ''
c.rev.buff = buff[size0 + size1:]
return buff[size0:size0 + size1]
def read_connect(self):
buff = self.read_bin()
if not buff:
return
msg = Message()
r = msg.decode(buff)
if r != 0:
return
if r < 0:
self.close()
return
c = self.conn_proxy
log.debug(0, '*%d read message: %s', c.index, msg)
cmd = msg.get(0)
if cmd != 'connect rsp':
log.error(0, 'invalid command. msg:%s', msg)
return
err = msg.get(1)
if err != 'ok':
log.error(0, 'accept fail. msg:%s', msg)
self.close()
return
Bridge(self.conn, self.conn_proxy)
self.conn = None
self.conn_proxy = None
if self.timer:
del_timer(self.timer)
self.timer = None
class Cross(object):
def __init__(self, proxy, target, ckey):
self.addr_proxy = proxy
self.addr_target = target
self.conn_proxy = None
self.conn_target = None
self.ready_proxy = False
self.ready_target = False
self.key = ckey
def init(self):
self.conn_proxy = Connection()
self.conn_target = Connection()
c = self.conn_proxy
c.connect_nonblocking(self.addr_proxy)
c.wev.handler = lambda: self.ready_connect_proxy()
add_conn(c, WRITE_EVENT)
c = self.conn_target
c.connect_nonblocking(self.addr_target)
c.wev.handler = lambda: self.ready_connect_target()
add_conn(self.conn_target, WRITE_EVENT)
def close(self):
if self.conn_proxy:
self.conn_proxy.close()
self.conn_proxy = None
if self.conn_target:
self.conn_target.close()
self.conn_target = None
def ready_connect_proxy(self):
c = self.conn_proxy
log.debug(0, '*%d connect: %s', c.index, c.addr.text)
msg = Message(['accept req', self.key])
log.debug(0, '*%d send message: %s', c.index, msg)
buff = msg.encode()
if not buff:
log.error(0, 'invalid message: %s', msg)
return
c.wev.buff = struct.pack('I', len(buff)) + buff
c.wev.handler = lambda: self.send()
self.send()
def ready_connect_target(self):
c = self.conn_target
log.debug(0, '*%d connect: %s', c.index, c.addr.text)
del_conn(c, WRITE_EVENT)
self.ready_target = True
self.check_ready()
def send(self):
c = self.conn_proxy
buff = c.wev.buff
if len(buff) > 0:
r, size = c.send(buff)
if r != 0:
if r == 1:
log.debug(0, '*%d closed', c.index)
self.close()
return
c.wev.buff = buff[size:]
buff = c.wev.buff
if len(buff) > 0:
c.wev.handler = lambda: self.send()
del_conn(c, READ_EVENT)
add_conn(c, WRITE_EVENT)
else:
c.rev.handler = lambda: self.read()
del_conn(c, WRITE_EVENT)
add_conn(c, READ_EVENT)
def read(self):
buff = self.read_bin()
if not buff:
return
msg = Message()
r = msg.decode(buff)
if r != 0:
return
if r < 0:
self.close()
return
c = self.conn_proxy
log.debug(0, '*%d read message. msg:%s', c.index, msg)
cmd = msg.get(0)
if cmd != 'accept rsp':
log.error(0, 'invalid message. msg:%s', msg)
self.close()
return
err = msg.get(1)
if err != 'ok':
log.error(0, 'accept fail. msg:%s', msg)
self.close()
return
self.ready_proxy = True
del_conn(c, WRITE_EVENT)
self.check_ready()
def read_bin(self):
c = self.conn_proxy
r, buff = c.recv(4096)
if r != 0:
if r == 1:
log.debug(0, '*%d closed', c.index)
self.close()
return ''
c.rev.buff += buff
buff = c.rev.buff
size0 = struct.calcsize('I')
if len(buff) < size0:
return ''
size1, = struct.unpack('I', buff[0:size0])
if len(buff) < size1 + size0:
return ''
c.rev.buff = buff[size0 + size1:]
return buff[size0:size0 + size1]
def check_ready(self):
if self.ready_target:
self.conn_target.wev.handler = None
del_conn(self.conn_target, WRITE_EVENT)
if self.ready_proxy:
self.conn_proxy.wev.handler = None
del_conn(self.conn_proxy, WRITE_EVENT)
if self.ready_proxy and self.ready_target:
Bridge(self.conn_proxy, self.conn_target)
self.conn_proxy = None
self.conn_target = None
