specify private addresses to assign to VETHs in the config

3 files changed, 136 insertions, 33 deletions

diff --git a/db_connection_config.yml b/db_connection_config.yml
index 680584c..08e952f 100644
--- a/db_connection_config.yml
+++ b/db_connection_config.yml
@@ -6,3 +6,4 @@ database: "ztdns"
 enabled: no

 handled_vpns: [1, 2]

 parallel_vpns: 20

+private_addresses: ["10.25.25.0 - 10.25.25.59", "10.25.25.120 - 10.25.25.139"]
\ No newline at end of file
diff --git a/src/hourly.py b/src/hourly.py
index f830bb7..209b577 100755
--- a/src/hourly.py
+++ b/src/hourly.py
@@ -4,6 +4,7 @@ from sys import argv
 import subprocess
 from os import path, waitpid, unlink
 from time import gmtime, strftime, sleep
+import re
 
 # our own module used by several scripts in the project
 from ztdnslib import start_db_connection, \
@@ -54,11 +55,91 @@ def unlock_on_file():
     except FileNotFoundError:
         return False
 
+address_range_regex = re.compile(r'''
+([\d]+\.[\d]+\.[\d]+\.[\d]+) # first IPv4 address in the range
+
+[\s]*-[\s]*                  # dash (with optional whitespace around)
+
+([\d]+\.[\d]+\.[\d]+\.[\d]+) # last IPv4 address in the range
+''', re.VERBOSE)
+
+address_regex = re.compile(r'([\d]+)\.([\d]+)\.([\d]+)\.([\d]+)')
+
+def ip_address_to_number(address):
+    match = address_regex.match(address)
+    if not match:
+        return None
+    number = 0
+    for byte in match.groups():
+        byteval = int(byte)
+        if byteval > 256:
+            return None
+        number = number * 256 + byteval
+    return number
+
+def number_to_ip_address(number):
+    byte1 = number % 256
+    number = number // 256
+    byte2 = number % 256
+    number = number // 256
+    byte3 = number % 256
+    number = number // 256
+    byte4 = number % 256
+    return "{}.{}.{}.{}".format(byte4, byte3, byte2, byte1)
+
+# this functions accepts list of IPv4 address ranges like:
+#     ['10.25.25.0 - 10.25.25.59', '10.25.25.120 - 10.25.25.135']
+# and returns a set of /30 subnetworks; each subnetwork is represented
+# by a tuple of 2 usable addresses within that subnetwork.
+# E.g. for subnetwork 10.25.25.16/30 it would be ('10.25.25.17', '10.25.25.18');
+# Addressess ending with .16 (subnet address)
+# and .19 (broadcast in the subnet) are considered unusable in this case.
+# The returned set will contain up to count elements.
+def get_available_subnetworks(count, address_ranges, logfile):
+    available_subnetworks = set()
+
+    for address_range in address_ranges:
+        match = address_range_regex.match(address_range)
+        ok_flag = True
+
+        if not match:
+            ok_flag = False
+
+        if ok_flag:
+            start_addr_number = ip_address_to_number(match.groups()[0])
+            end_addr_number = ip_address_to_number(match.groups()[1])
+            if not start_addr_number or not end_addr_number:
+                ok_flag = False
+
+        if ok_flag:
+            # round so that start_addr is first ip address in a /30 network
+            # and end_addr is last ip address in a /30 network
+            while start_addr_number % 4 != 0:
+                start_addr_number += 1
+            while end_addr_number % 4 != 3:
+                end_addr_number -= 1
+
+            if start_addr_number >= end_addr_number:
+                logfile.write("address range '{}' doesn't contain any"
+                              " /30 subnetworks\n".format(address_range))
+            else:
+                while len(available_subnetworks) < count and \
+                      start_addr_number < end_addr_number:
+                    usable_addr1 = number_to_ip_address(start_addr_number + 1)
+                    usable_addr2 = number_to_ip_address(start_addr_number + 2)
+                    available_subnetworks.add((usable_addr1, usable_addr2))
+                    start_addr_number += 4
+        else:
+            logfile.write("'{}' is not a valid address range\n"\
+                          .format(address_range))
+
+    return available_subnetworks
+
 def do_hourly_work(hour, logfile):
     ztdns_config = get_ztdns_config()
     if ztdns_config['enabled'] != 'yes':
         logfile.write("0tdns not enabled in the config - exiting\n")
-        exit()
+        return
 
     connection = start_db_connection(ztdns_config)
     cursor = connection.cursor()
@@ -73,6 +154,23 @@ def do_hourly_work(hour, logfile):
         # if not specfied in the config, all vpns are handled
         hadled_vpns = [vpn[0] for vpn in vpns]
 
+    parallel_vpns = ztdns_config['parallel_vpns'] # we need this many subnets
+    subnets = get_available_subnetworks(parallel_vpns,
+                                        ztdns_config['private_addresses'],
+                                        logfile)
+
+    if not subnets:
+        logfile.write("couldn't get ANY /30 subnet of private addresses from"
+                      " the 0tdns config file - exiting\n");
+        return # TODO close cursor and connection here
+
+    if len(subnets) < parallel_vpns:
+        logfile.write("configuration allows running {0} parallel vpn"
+                      " connections, but provided private ip addresses give"
+                      " only {1} /30 subnets, which limits parallel connections"
+                      " to {1}\n".format(parallel_vpns, len(subnets)))
+        parallel_vpns = len(subnets)
+    
     for vpn_id, config_hash in vpns:
         config_path = "/var/lib/0tdns/{}.ovpn".format(config_hash)
         if not path.isfile(config_path):
@@ -80,41 +178,48 @@ def do_hourly_work(hour, logfile):
                           .format(vpn_id, config_hash))
             sync_ovpn_config(cursor, vpn_id, config_path, config_hash)
 
-    parallel_vpns = ztdns_config['parallel_vpns']
-    pids_vpns = {} # map each wrapper pid to id of the vpn it connects to
+    # map of each wrapper pid to tuple containing id of the vpn it connects to
+    # and subnet (represented as tuple of addresses) it uses for veth device
+    pids_wrappers = {}
 
     def wait_for_wrapper_process():
         while True:
             pid, exit_status = waitpid(0, 0)
-            # make sure 
-            if pids_vpns.get(pid) is not None:
+            # make sure it's one of our wrapper processes
+            vpn_id, subnet, _ = pids_wrappers.get(pid, (None, None, None))
+            if subnet:
                 break
+
         if exit_status == 2:
             # this means our perform_queries.py crashed... not good
             logfile.write('performing queries through vpn {} failed\n'\
-                          .format(pids_vpns[pid]))
+                          .format(vpn_id))
         elif exit_status != 0:
             # vpn server is probably not responding
             logfile.write('connection to vpn {} failed\n'\
-                          .format(pids_vpns[pid]))
-        pids_vpns.pop(pid)
+                          .format(vpn_id))
+        pids_wrappers.pop(pid)
+        subnets.add(subnet)
 
     for vpn_id, config_hash in vpns:
-        if len(pids_vpns) == parallel_vpns:
+        if len(pids_wrappers) == parallel_vpns:
             wait_for_wrapper_process()
 
         config_path = "/var/lib/0tdns/{}.ovpn".format(config_hash)
         physical_ip = get_default_host_address(ztdns_config['host'])
+        subnet = subnets.pop()
+        veth_addr1, veth_addr2 = subnet
         route_through_veth = ztdns_config['host'] + "/32"
         command_in_namespace = [perform_queries, hour, str(vpn_id)]
         logfile.write("Running connection for vpn {}\n".format(vpn_id))
-        
-        p = subprocess.Popen([wrapper, config_path, physical_ip,
-                              route_through_veth] + command_in_namespace)
 
-        pids_vpns[p.pid] = vpn_id
+        p = subprocess.Popen([wrapper, config_path, physical_ip, veth_addr1,
+                              veth_addr2, route_through_veth] +
+                             command_in_namespace)
+
+        pids_wrappers[p.pid] = (vpn_id, subnet, p)
 
-    while len(pids_vpns) > 0:
+    while len(pids_wrappers) > 0:
         wait_for_wrapper_process()
 
     cursor.execute('''
diff --git a/src/vpn_wrapper.sh b/src/vpn_wrapper.sh
index 8dbb702..2dbb821 100755
--- a/src/vpn_wrapper.sh
+++ b/src/vpn_wrapper.sh
@@ -5,28 +5,25 @@
 OVPN_COMMAND="/usr/sbin/openvpn"
 
 OPENVPN_CONFIG="$1"
-PHYSICAL_IP="$2"
-ROUTE_THROUGH_VETH="$3"
-# rest of args is the command to run in network namespace
-shift
-shift
-shift
-
 # for routing some traffic from within the namespace to physical
 # network (e.g. database connection) we need to create a veth pair;
+# ip datagrams routed through veth pair are going to have veth's private address
+# as their source address - we need to change it to the address of our physical
+# network device using iptables' SNAT. This address is provided by the caller.
+PHYSICAL_IP="$2"
 # as we want multiple instances of vpn_wrapper.sh to be able to
-# run simultaneously, we need unique ip addresses for them;
-# the solution is to derive an ip address from current shell's
-# PID (which is unique within a system)
-NUMBER=$((($$ - 1) * 4))
-WORD0HOST0=$(($NUMBER % 256 + 1))
-WORD0HOST1=$(($NUMBER % 256 + 2))
-NUMBER=$(($NUMBER / 256))
-WORD1=$(($NUMBER % 256))
-NUMBER=$(($NUMBER / 256))
-WORD2=$(($NUMBER % 256))
-VETH_HOST0=10.$WORD2.$WORD1.$WORD0HOST0
-VETH_HOST1=10.$WORD2.$WORD1.$WORD0HOST1
+# run simultaneously, we need unique ip addresses for veth devices, which
+# caller provides to us in command line arguments
+VETH_HOST0="$3"
+VETH_HOST1="$4"
+# caller specifies space-delimited subnets, traffic to which should not be
+# routed through the vpn (<database_ip>/32 is going to be here)
+ROUTE_THROUGH_VETH="$5"
+
+# rest of args is the command to run in network namespace
+for _ in `seq 5`; do
+    shift
+done
 
 # to enable multiple instances of this script to run simultaneously,
 # we tag namespace name with this shell's PID