Skip to content

Commit

Permalink
Do non-blocking connect().
Browse files Browse the repository at this point in the history
This way we don't freeze the entire proxy when someone tries to connect to a
nonexistent IP address (oops).
  • Loading branch information
apenwarr committed May 2, 2010
1 parent 81c89ce commit 9b23fd2
Showing 1 changed file with 36 additions and 14 deletions.
50 changes: 36 additions & 14 deletions ssnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,13 +46,15 @@ def _try_peername(sock):


class SockWrapper:
def __init__(self, rsock, wsock, peername=None):
def __init__(self, rsock, wsock, connect_to=None, peername=None):
self.exc = None
self.rsock = rsock
self.wsock = wsock
self.shut_read = self.shut_write = False
self.buf = []
self.connect_to = connect_to
self.peername = peername or _try_peername(self.rsock)
self.try_connect()

def __del__(self):
debug1('%r: deleting\n' % self)
Expand All @@ -66,6 +68,23 @@ def seterr(self, e):
if not self.exc:
self.exc = e

def try_connect(self):
if not self.connect_to:
return # already connected
self.rsock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42)
self.rsock.setblocking(False)
try:
self.rsock.connect(self.connect_to)
self.connect_to = None
except socket.error, e:
if e.args[0] in [errno.EINPROGRESS, errno.EALREADY]:
pass # not connected yet
elif e.args[0] in [errno.ECONNREFUSED, errno.ETIMEDOUT]:
# a "normal" kind of error
self.seterr(e)
else:
raise # error we've never heard of?! barf completely.

def noread(self):
if not self.shut_read:
debug2('%r: done reading\n' % self)
Expand All @@ -82,6 +101,8 @@ def nowrite(self):
self.seterr(e)

def uwrite(self, buf):
if self.connect_to:
return 0 # still connecting
self.wsock.setblocking(False)
try:
return _nb_clean(os.write, self.wsock.fileno(), buf)
Expand All @@ -97,6 +118,8 @@ def write(self, buf):
return self.uwrite(buf)

def uread(self):
if self.connect_to:
return None # still connecting
if self.shut_read:
return
self.rsock.setblocking(False)
Expand Down Expand Up @@ -154,16 +177,23 @@ def __init__(self, wrap1, wrap2):
self.wrap2 = wrap2

def pre_select(self, r, w, x):
if self.wrap1.buf:
if self.wrap1.connect_to:
w.add(self.wrap1.rsock)
elif self.wrap1.buf:
w.add(self.wrap2.wsock)
elif not self.wrap1.shut_read:
r.add(self.wrap1.rsock)
if self.wrap2.buf:

if self.wrap2.connect_to:
w.add(self.wrap2.rsock)
elif self.wrap2.buf:
w.add(self.wrap1.wsock)
elif not self.wrap2.shut_read:
r.add(self.wrap2.rsock)

def callback(self):
self.wrap1.try_connect()
self.wrap2.try_connect()
self.wrap1.fill()
self.wrap2.fill()
self.wrap1.copy_to(self.wrap2)
Expand Down Expand Up @@ -324,14 +354,6 @@ def got_packet(self, cmd, data):
def connect_dst(ip, port):
debug2('Connecting to %s:%d\n' % (ip, port))
outsock = socket.socket()
outsock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42)
e = None
try:
outsock.connect((ip,port))
except socket.error, e:
if e.args[0] not in [errno.ECONNREFUSED]:
raise
sw = SockWrapper(outsock, outsock, peername = '%s:%d' % (ip,port))
if e:
sw.seterr(e)
return sw
return SockWrapper(outsock, outsock,
connect_to = (ip,port),
peername = '%s:%d' % (ip,port))

0 comments on commit 9b23fd2

Please sign in to comment.