TProxy UDP support, including DNS.

firewall.py

@@ -69,10 +69,12 @@ def _ipt_ttl(family, *args):
 # multiple copies shouldn't have overlapping subnets, or only the most-
 # recently-started one will win (because we use "-I OUTPUT 1" instead of
 # "-A OUTPUT").
-def do_iptables_nat(port, dnsport, family, subnets):
+def do_iptables_nat(port, dnsport, family, subnets, udp):
     # only ipv4 supported with NAT
     if family != socket.AF_INET:
         raise Exception('Address family "%s" unsupported by nat method'%family_to_string(family))
+    if udp:
+        raise Exception("UDP not supported by nat method")
 
     table = "nat"
     def ipt(*args):
@@ -122,7 +124,7 @@ def do_iptables_nat(port, dnsport, family, subnets):
                     '--to-ports', str(dnsport))
 
 
-def do_iptables_tproxy(port, dnsport, family, subnets):
+def do_iptables_tproxy(port, dnsport, family, subnets, udp):
     if family not in [socket.AF_INET, socket.AF_INET6]:
         raise Exception('Address family "%s" unsupported by tproxy method'%family_to_string(family))
 
@@ -164,6 +166,21 @@ def do_iptables_tproxy(port, dnsport, family, subnets):
         ipt('-A', divert_chain, '-j', 'ACCEPT')
         ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain,
              '-m', 'tcp', '-p', 'tcp')
+    if subnets and udp:
+        ipt('-A', tproxy_chain, '-m', 'socket', '-j', divert_chain,
+             '-m', 'udp', '-p', 'udp')
+
+    if dnsport:
+        nslist = resolvconf_nameservers()
+        for f,ip in filter(lambda i: i[0]==family, nslist):
+            ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1',
+                '--dest', '%s/32' % ip,
+                '-m', 'udp', '-p', 'udp', '--dport', '53')
+            ipt('-A', tproxy_chain, '-j', 'TPROXY', '--tproxy-mark', '0x1/0x1',
+                 '--dest', '%s/32' % ip,
+                 '-m', 'udp', '-p', 'udp', '--dport', '53',
+                 '--on-port', str(dnsport))
+
     if subnets:
         for f,swidth,sexclude,snet in sorted(subnets, key=lambda s: s[1], reverse=True):
             if sexclude:
@@ -182,6 +199,22 @@ def do_iptables_tproxy(port, dnsport, family, subnets):
                      '-m', 'tcp', '-p', 'tcp',
                      '--on-port', str(port))
 
+            if sexclude and udp:
+                ipt('-A', mark_chain, '-j', 'RETURN',
+                    '--dest', '%s/%s' % (snet,swidth),
+                    '-m', 'udp', '-p', 'udp')
+                ipt('-A', tproxy_chain, '-j', 'RETURN',
+                    '--dest', '%s/%s' % (snet,swidth),
+                    '-m', 'udp', '-p', 'udp')
+            elif udp:
+                ipt('-A', mark_chain, '-j', 'MARK', '--set-mark', '1',
+                     '--dest', '%s/%s' % (snet,swidth),
+                     '-m', 'udp', '-p', 'udp')
+                ipt('-A', tproxy_chain, '-j', 'TPROXY', '--tproxy-mark', '0x1/0x1',
+                     '--dest', '%s/%s' % (snet,swidth),
+                     '-m', 'udp', '-p', 'udp',
+                     '--on-port', str(port))
+
 
 def ipfw_rule_exists(n):
     argv = ['ipfw', 'list']
@@ -288,10 +321,12 @@ def ipfw(*args):
         raise Fatal('%r returned %d' % (argv, rv))
 
 
-def do_ipfw(port, dnsport, family, subnets):
+def do_ipfw(port, dnsport, family, subnets, udp):
     # IPv6 not supported
     if family not in [socket.AF_INET, ]:
         raise Exception('Address family "%s" unsupported by ipfw method'%family_to_string(family))
+    if udp:
+        raise Exception("UDP not supported by ipfw method")
 
     sport = str(port)
     xsport = str(port+1)
@@ -454,7 +489,7 @@ def restore_etc_hosts(port):
 # exit.  In case that fails, it's not the end of the world; future runs will
 # supercede it in the transproxy list, at least, so the leftover rules
 # are hopefully harmless.
-def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, syslog):
+def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, udp, syslog):
     assert(port_v6 >= 0)
     assert(port_v6 <= 65535)
     assert(port_v4 >= 0)
@@ -529,13 +564,13 @@ def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, syslog):
 
             subnets_v6 = filter(lambda i: i[0]==socket.AF_INET6, subnets)
             if port_v6:
-                do_wait = do_it(port_v6, dnsport_v6, socket.AF_INET6, subnets_v6)
+                do_wait = do_it(port_v6, dnsport_v6, socket.AF_INET6, subnets_v6, udp)
             elif len(subnets_v6) > 0:
                 debug1("IPv6 subnets defined but IPv6 disabled\n")
 
             subnets_v4 = filter(lambda i: i[0]==socket.AF_INET, subnets)
             if port_v4:
-                do_wait = do_it(port_v4, dnsport_v4, socket.AF_INET, subnets_v4)
+                do_wait = do_it(port_v4, dnsport_v4, socket.AF_INET, subnets_v4, udp)
             elif len(subnets_v4) > 0:
                 debug1('IPv4 subnets defined but IPv4 disabled\n')
 
@@ -567,6 +602,8 @@ def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, syslog):
             debug1('firewall manager: undoing changes.\n')
         except:
             pass
+        if port_v6:
+            do_it(port_v6, 0, socket.AF_INET6, [], udp)
         if port_v4:
-            do_it(port_v4, 0, socket.AF_INET, [])
+            do_it(port_v4, 0, socket.AF_INET, [], udp)
         restore_etc_hosts(port_v6 or port_v4)

main.py

@@ -123,11 +123,11 @@ try:
         server.latency_control = opt.latency_control
         sys.exit(server.main())
     elif opt.firewall:
-        if len(extra) != 5:
-            o.fatal('exactly five arguments expected')
+        if len(extra) != 6:
+            o.fatal('exactly six arguments expected')
         sys.exit(firewall.main(int(extra[0]), int(extra[1]),
                                int(extra[2]), int(extra[3]),
-                                extra[4], opt.syslog))
+                               extra[4], int(extra[5]), opt.syslog))
     elif opt.hostwatch:
         sys.exit(hostwatch.hw_main(extra))
     else:

server.py

@@ -163,6 +163,35 @@ class DnsProxy(Handler):
         self.ok = False
 
 
+class UdpProxy(Handler):
+    def __init__(self, mux, chan, family):
+        sock = socket.socket(family, socket.SOCK_DGRAM)
+        Handler.__init__(self, [sock])
+        self.timeout = time.time()+30
+        self.mux = mux
+        self.chan = chan
+        self.sock = sock
+        if family == socket.AF_INET:
+            self.sock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42)
+
+    def send(self, dstip, data):
+        debug2('UDP: sending to %r port %d\n' % dstip)
+        try:
+            self.sock.sendto(data,dstip)
+        except socket.error, e:
+            log('UDP send to %r port %d: %s\n' % (dstip[0], dstip[1], e))
+            return
+
+    def callback(self):
+        try:
+            data,peer = self.sock.recvfrom(4096)
+        except socket.error, e:
+            log('UDP recv from %r port %d: %s\n' % (peer[0], peer[1], e))
+            return
+        debug2('UDP response: %d bytes\n' % len(data))
+        hdr = "%s,%r,"%(peer[0], peer[1])
+        self.mux.send(self.chan, ssnet.CMD_UDP_DATA, hdr+data)
+
 def main():
     if helpers.verbose >= 1:
         helpers.logprefix = ' s: '
@@ -231,6 +260,34 @@ def main():
         dnshandlers[channel] = h
     mux.got_dns_req = dns_req
 
+    udphandlers = {}
+    def udp_req(channel, cmd, data):
+        debug2('Incoming UDP request channel=%d, cmd=%d\n' % (channel,cmd))
+        if cmd == ssnet.CMD_UDP_DATA:
+            (dstip,dstport,data) = data.split(",",2)
+            dstport = int(dstport)
+            debug2('is incoming UDP data. %r %d.\n' % (dstip,dstport))
+            h = udphandlers[channel]
+            h.send((dstip,dstport),data)
+        elif cmd == ssnet.CMD_UDP_CLOSE:
+            debug2('is incoming UDP close\n')
+            h = udphandlers[channel]
+            h.ok = False
+            del mux.channels[channel]
+
+    def udp_open(channel, data):
+        debug2('Incoming UDP open.\n')
+        family = int(data)
+        mux.channels[channel] = lambda cmd, data: udp_req(channel, cmd, data)
+        if channel in udphandlers:
+            raise Fatal('UDP connection channel %d already open'%channel)
+        else:
+            h = UdpProxy(mux, channel, family)
+            handlers.append(h)
+            udphandlers[channel] = h
+    mux.got_udp_open = udp_open
+
+
     while mux.ok:
         if hw.pid:
             assert(hw.pid > 0)
@@ -250,3 +307,8 @@ def main():
                     debug3('expiring dnsreqs channel=%d\n' % channel)
                     del dnshandlers[channel]
                     h.ok = False
+            for channel,h in udphandlers.items():
+                if not h.ok:
+                    debug3('expiring UDP channel=%d\n' % channel)
+                    del udphandlers[channel]
+                    h.ok = False

ssnet.py

@@ -25,6 +25,9 @@ CMD_HOST_REQ = 0x4208
 CMD_HOST_LIST = 0x4209
 CMD_DNS_REQ = 0x420a
 CMD_DNS_RESPONSE = 0x420b
+CMD_UDP_OPEN = 0x420c
+CMD_UDP_DATA = 0x420d
+CMD_UDP_CLOSE = 0x420e
 
 cmd_to_name = {
     CMD_EXIT: 'EXIT',
@@ -39,6 +42,9 @@ cmd_to_name = {
     CMD_HOST_LIST: 'HOST_LIST',
     CMD_DNS_REQ: 'DNS_REQ',
     CMD_DNS_RESPONSE: 'DNS_RESPONSE',
+    CMD_UDP_OPEN: 'UDP_OPEN',
+    CMD_UDP_DATA: 'UDP_DATA',
+    CMD_UDP_CLOSE: 'UDP_CLOSE',
 }
 
 
@@ -318,6 +324,7 @@ class Mux(Handler):
         self.rsock = rsock
         self.wsock = wsock
         self.new_channel = self.got_dns_req = self.got_routes = None
+        self.got_udp_open = self.got_udp_data = self.got_udp_close = None
         self.got_host_req = self.got_host_list = None
         self.channels = {}
         self.chani = 0
@@ -383,6 +390,10 @@ class Mux(Handler):
             assert(not self.channels.get(channel))
             if self.got_dns_req:
                 self.got_dns_req(channel, data)
+        elif cmd == CMD_UDP_OPEN:
+            assert(not self.channels.get(channel))
+            if self.got_udp_open:
+                self.got_udp_open(channel, data)
         elif cmd == CMD_ROUTES:
             if self.got_routes:
                 self.got_routes(data)