Disable NAGLE algorithm to decrease TCP tunnel delay
[nepi.git] / src / nepi / util / tunchannel.py
index d622170..03adfe5 100644 (file)
@@ -4,8 +4,10 @@ import os
 import struct
 import socket
 import threading
+import traceback
 import errno
 import fcntl
+import random
 import traceback
 import functools
 import collections
@@ -83,7 +85,7 @@ def formatPacket(packet, ether_mode):
             packet[48:] if (int(packet[1],16) > 5) else packet[40:], # payload
         ) ) )
 
-def _packetReady(buf, ether_mode=False, len=len):
+def _packetReady(buf, ether_mode=False, len=len, str=str):
     if not buf:
         return False
         
@@ -98,35 +100,36 @@ def _packetReady(buf, ether_mode=False, len=len):
             totallen = socket.htons(totallen)
             rv = len(buf[0]) >= totallen
         if not rv and len(buf) > 1:
-            nbuf = ''.join(buf)
-            buf.clear()
-            buf.append(nbuf)
+            # collapse only first two buffers
+            # as needed, to mantain len(buf) meaningful
+            p1 = buf.popleft()
+            buf[0] = p1+str(buf[0])
         else:
             return rv
     return rv
 
-def _pullPacket(buf, ether_mode=False, len=len):
+def _pullPacket(buf, ether_mode=False, len=len, buffer=buffer):
     if ether_mode:
         return buf.popleft()
     else:
         _,totallen = struct.unpack('HH',buf[0][:4])
         totallen = socket.htons(totallen)
-        if len(buf[0]) < totallen:
-            rv = buf[0][:totallen]
-            buf[0] = buf[0][totallen:]
+        if len(buf[0]) > totallen:
+            rv = buffer(buf[0],0,totallen)
+            buf[0] = buffer(buf[0],totallen)
         else:
             rv = buf.popleft()
         return rv
 
-def etherStrip(buf):
+def etherStrip(buf, buffer=buffer, len=len):
     if len(buf) < 14:
         return ""
     if buf[12:14] == '\x08\x10' and buf[16:18] == '\x08\x00':
         # tagged ethernet frame
-        return buf[18:]
+        return buffer(buf, 18)
     elif buf[12:14] == '\x08\x00':
         # untagged ethernet frame
-        return buf[14:]
+        return buffer(buf, 14)
     else:
         return ""
 
@@ -192,12 +195,14 @@ def nonblock(fd):
         return False
 
 def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr=sys.stderr, reconnect=None, rwrite=None, rread=None, tunqueue=1000, tunkqueue=1000,
-        cipher='AES', accept_local=None, accept_remote=None, slowlocal=True,
-        len=len, max=max, OSError=OSError, select=select.select, selecterror=select.error, os=os, socket=socket,
+        cipher='AES', accept_local=None, accept_remote=None, slowlocal=True, queueclass=None, bwlimit=None,
+        len=len, max=max, min=min, buffer=buffer, OSError=OSError, select=select.select, selecterror=select.error, os=os, socket=socket,
         retrycodes=(os.errno.EWOULDBLOCK, os.errno.EAGAIN, os.errno.EINTR) ):
     crypto_mode = False
+    crypter = None
+
     try:
-        if cipher_key:
+        if cipher_key and cipher:
             import Crypto.Cipher
             import hashlib
             __import__('Crypto.Cipher.'+cipher)
@@ -232,7 +237,7 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
         if rread is None:
             def rread(remote, maxlen, os_read=os.read):
                 return os_read(remote_fd, maxlen)
-    
     rnonblock = nonblock(remote)
     tnonblock = nonblock(tun)
     
@@ -290,33 +295,58 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
                 else:
                     return None
     
-    # Limited frame parsing, to preserve packet boundaries.
-    # Which is needed, since /dev/net/tun is unbuffered
     maxbkbuf = maxfwbuf = max(10,tunqueue-tunkqueue)
     tunhurry = max(0,maxbkbuf/2)
-    fwbuf = collections.deque()
-    bkbuf = collections.deque()
+    
+    if queueclass is None:
+        queueclass = collections.deque
+        maxbatch = 2000
+        maxtbatch = 50
+    else:
+        maxfwbuf = maxbkbuf = 2000000000
+        maxbatch = 50
+        maxtbatch = 30
+        tunhurry = 30
+    
+    fwbuf = queueclass()
+    bkbuf = queueclass()
     nfwbuf = 0
     nbkbuf = 0
-    if ether_mode:
+    
+    # backwards queue functions
+    # they may need packet inspection to 
+    # reconstruct packet boundaries
+    if ether_mode or udp:
         packetReady = bool
-        pullPacket = collections.deque.popleft
-        reschedule = collections.deque.appendleft
+        pullPacket = queueclass.popleft
+        reschedule = queueclass.appendleft
     else:
         packetReady = _packetReady
         pullPacket = _pullPacket
-        reschedule = collections.deque.appendleft
+        reschedule = queueclass.appendleft
+    
+    # forward queue functions
+    # no packet inspection needed
+    fpacketReady = bool
+    fpullPacket = queueclass.popleft
+    freschedule = queueclass.appendleft
+    
     tunfd = tun.fileno()
     os_read = os.read
     os_write = os.write
     
+    tget = time.time
+    maxbwfree = bwfree = 1500 * tunqueue
+    lastbwtime = tget()
+    
     remoteok = True
     
+    
     while not TERMINATE:
         wset = []
         if packetReady(bkbuf):
             wset.append(tun)
-        if remoteok and packetReady(fwbuf):
+        if remoteok and fpacketReady(fwbuf) and (not bwlimit or bwfree > 0):
             wset.append(remote)
         
         rset = []
@@ -336,6 +366,9 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
             if e.args[0] == errno.EINTR:
                 # just retry
                 continue
+            else:
+                traceback.print_exc(file=sys.stderr)
+                raise
 
         # check for errors
         if errs:
@@ -359,18 +392,26 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
         # check to see if we can write
         #rr = wr = rt = wt = 0
         if remote in wrdy:
+            sent = 0
             try:
                 try:
-                    for x in xrange(2000):
+                    for x in xrange(maxbatch):
                         packet = pullPacket(fwbuf)
 
                         if crypto_mode:
                             packet = encrypt_(packet, crypter)
                         
-                        rwrite(remote, packet)
+                        sentnow = rwrite(remote, packet)
+                        sent += sentnow
                         #wr += 1
                         
-                        if not rnonblock or not packetReady(fwbuf):
+                        if not udp and 0 <= sentnow < len(packet):
+                            # packet partially sent
+                            # reschedule the remaining part
+                            # this doesn't happen ever in udp mode
+                            freschedule(fwbuf, buffer(packet,sentnow))
+                        
+                        if not rnonblock or not fpacketReady(fwbuf):
                             break
                 except OSError,e:
                     # This except handles the entire While block on PURPOSE
@@ -378,7 +419,7 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
                     # The only operation that can raise this exception is rwrite
                     if e.errno in retrycodes:
                         # re-schedule packet
-                        reschedule(fwbuf, packet)
+                        freschedule(fwbuf, packet)
                     else:
                         raise
             except:
@@ -392,9 +433,12 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
                     # in UDP mode, we ignore errors - packet loss man...
                     raise
                 #traceback.print_exc(file=sys.stderr)
+            
+            if bwlimit:
+                bwfree -= sent
         if tun in wrdy:
             try:
-                for x in xrange(50):
+                for x in xrange(maxtbatch):
                     packet = pullPacket(bkbuf)
                     twrite(tunfd, packet)
                     #wt += 1
@@ -422,7 +466,7 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
         # check incoming data packets
         if tun in rdrdy:
             try:
-                for x in xrange(2000):
+                for x in xrange(maxbatch):
                     packet = tread(tunfd,2000) # tun.read blocks until it gets 2k!
                     if not packet:
                         continue
@@ -440,8 +484,9 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
         if remote in rdrdy:
             try:
                 try:
-                    for x in xrange(2000):
+                    for x in xrange(maxbatch):
                         packet = rread(remote,2000)
+                        
                         #rr += 1
                         
                         if crypto_mode:
@@ -476,8 +521,222 @@ def tun_fwd(tun, remote, with_pi, ether_mode, cipher_key, udp, TERMINATE, stderr
                     # in UDP mode, we ignore errors - packet loss man...
                     raise
                 traceback.print_exc(file=sys.stderr)
+
+        if bwlimit:
+            tnow = tget()
+            delta = tnow - lastbwtime
+            if delta > 0.001:
+                delta = int(bwlimit * delta)
+                if delta > 0:
+                    bwfree = min(bwfree+delta, maxbwfree)
+                    lastbwtime = tnow
         
         #print >>sys.stderr, "rr:%d\twr:%d\trt:%d\twt:%d" % (rr,wr,rt,wt)
 
+def udp_connect(TERMINATE, local_addr, local_port, peer_addr, peer_port):
+    rsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 0)
+    retrydelay = 1.0
+    for i in xrange(30):
+        # TERMINATE is a array. An item can be added to TERMINATE, from
+        # outside this function to force termination of the loop
+        if TERMINATE:
+            raise OSError, "Killed"
+        try:
+            rsock.bind((local_addr, local_port))
+            break
+        except socket.error:
+            # wait a while, retry
+            print >>sys.stderr, "%s: Could not bind. Retrying in a sec..." % (time.strftime('%c'),)
+            time.sleep(min(30.0,retrydelay))
+            retrydelay *= 1.1
+    else:
+        rsock.bind((local_addr, local_port))
+    print >>sys.stderr, "Listening UDP at: %s:%d" % (local_addr, local_port)
+    print >>sys.stderr, "Connecting UDP to: %s:%d" % (peer_addr, peer_port)
+    rsock.connect((peer_addr, peer_port))
+    return rsock
+
+def udp_handshake(TERMINATE, rsock):
+    endme = False
+    def keepalive():
+        while not endme and not TERMINATE:
+            try:
+                rsock.send('')
+            except:
+                pass
+            time.sleep(1)
+        try:
+            rsock.send('')
+        except:
+            pass
+    keepalive_thread = threading.Thread(target=keepalive)
+    keepalive_thread.start()
+    for i in xrange(900):
+        if TERMINATE:
+            raise OSError, "Killed"
+        try:
+            heartbeat = rsock.recv(10)
+            break
+        except:
+            time.sleep(1)
+    else:
+        heartbeat = rsock.recv(10)
+    endme = True
+    keepalive_thread.join()
+
+def udp_establish(TERMINATE, local_addr, local_port, peer_addr, peer_port):
+    rsock = udp_connect(TERMINATE, local_addr, local_port, peer_addr,
+            peer_port)
+    udp_handshake(TERMINATE, rsock)
+    return rsock 
+
+def tcp_connect(TERMINATE, stop, rsock, peer_addr, peer_port):
+    sock = None
+    retrydelay = 1.0
+    # The peer has a firewall that prevents a response to the connect, we 
+    # will be forever blocked in the connect, so we put a reasonable timeout.
+    rsock.settimeout(10) 
+    # We wait for 
+    for i in xrange(30):
+        if stop:
+            break
+        if TERMINATE:
+            raise OSError, "Killed"
+        try:
+            rsock.connect((peer_addr, peer_port))
+            sock = rsock
+            break
+        except socket.error:
+            # wait a while, retry
+            print >>sys.stderr, "%s: Could not connect. Retrying in a sec..." % (time.strftime('%c'),)
+            time.sleep(min(30.0,retrydelay))
+            retrydelay *= 1.1
+    else:
+        rsock.connect((peer_addr, peer_port))
+        sock = rsock
+    if sock:
+        print >>sys.stderr, "tcp_connect: TCP sock connected to remote %s:%s" % (peer_addr, peer_port)
+        sock.settimeout(0) 
+        
+        print >>sys.stderr, "tcp_connect: disabling NAGLE"
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    return sock
+
+def tcp_listen(TERMINATE, stop, lsock, local_addr, local_port):
+    sock = None
+    retrydelay = 1.0
+    # We try to bind to the local virtual interface. 
+    # It might not exist yet so we wait in a loop.
+    for i in xrange(30):
+        if stop:
+            break
+        if TERMINATE:
+            raise OSError, "Killed"
+        try:
+            lsock.bind((local_addr, local_port))
+            break
+        except socket.error:
+            # wait a while, retry
+            print >>sys.stderr, "%s: Could not bind. Retrying in a sec..." % (time.strftime('%c'),)
+            time.sleep(min(30.0,retrydelay))
+            retrydelay *= 1.1
+    else:
+        lsock.bind((local_addr, local_port))
+
+    print >>sys.stderr, "tcp_listen: TCP sock listening in local sock %s:%s" % (local_addr, local_port)
+    # Now we wait until the other side connects. 
+    # The other side might not be ready yet, so we also wait in a loop for timeouts.
+    timeout = 1
+    lsock.listen(1)
+    for i in xrange(30):
+        if TERMINATE:
+            raise OSError, "Killed"
+        rlist, wlist, xlist = select.select([lsock], [], [], timeout)
+        if stop:
+            break
+        if lsock in rlist:
+            sock,raddr = lsock.accept()
+            print >>sys.stderr, "tcp_listen: TCP connection accepted in local sock %s:%s" % (local_addr, local_port)
+            break
+        timeout += 5
+    return sock
+
+def tcp_handshake(rsock, listen, hand):
+    # we are going to use a barrier algorithm to decide wich side listen.
+    # each side will "roll a dice" and send the resulting value to the other 
+    # side. 
+    win = False
+    rsock.settimeout(10)
+    try:
+        rsock.send(hand)
+        peer_hand = rsock.recv(1)
+        print >>sys.stderr, "tcp_handshake: hand %s, peer_hand %s" % (hand, peer_hand)
+        if hand < peer_hand:
+            if listen:
+                win = True
+        elif hand > peer_hand:
+            if not listen:
+                win = True
+    except socket.timeout:
+        pass
+    rsock.settimeout(0)
+    return win
+
+def tcp_establish(TERMINATE, local_addr, local_port, peer_addr, peer_port):
+    def listen(stop, hand, lsock, lresult):
+        win = False
+        rsock = tcp_listen(TERMINATE, stop, lsock, local_addr, local_port)
+        if rsock:
+            win = tcp_handshake(rsock, True, hand)
+            stop.append(True)
+        lresult.append((win, rsock))
+
+    def connect(stop, hand, rsock, rresult):
+        win = False
+        rsock = tcp_connect(TERMINATE, stop, rsock, peer_addr, peer_port)
+        if rsock:
+            win = tcp_handshake(rsock, False, hand)
+            stop.append(True)
+        rresult.append((win, rsock))
+  
+    end = False
+    sock = None
+    for i in xrange(0, 50):
+        if end:
+            break
+        if TERMINATE:
+            raise OSError, "Killed"
+        hand = str(random.randint(1, 6))
+        stop = []
+        lresult = []
+        rresult = []
+        lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+        rsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+        listen_thread = threading.Thread(target=listen, args=(stop, hand, lsock, lresult))
+        connect_thread = threading.Thread(target=connect, args=(stop, hand, rsock, rresult))
+        connect_thread.start()
+        listen_thread.start()
+        connect_thread.join()
+        listen_thread.join()
+        (lwin, lrsock) = lresult[0]
+        (rwin, rrsock) = rresult[0]
+        if not lrsock or not rrsock:
+            if not lrsock:
+                sock = rrsock
+            if not rrsock:
+                sock = lrsock
+            end = True
+        # both socket are connected
+        else:
+           if lwin:
+                sock = lrsock
+                end = True
+           elif rwin: 
+                sock = rrsock
+                end = True
+
+    if not sock:
+        raise OSError, "Error: tcp_establish could not establish connection."
+    return sock