From e77ab483df656b292852b14c9ecc781de0c1ddd2 Mon Sep 17 00:00:00 2001
From: Disassembler <disassembler@dasm.cz>
Date: Fri, 2 Nov 2018 15:25:48 +0100
Subject: [PATCH] Update hosts file with exclusive lock

---
 usr/lib/python3.6/vmmgr/__init__.py |  5 ++--
 usr/lib/python3.6/vmmgr/tools.py    | 42 ++++++++++++++---------------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/usr/lib/python3.6/vmmgr/__init__.py b/usr/lib/python3.6/vmmgr/__init__.py
index d6782f9..1bfdbc8 100644
--- a/usr/lib/python3.6/vmmgr/__init__.py
+++ b/usr/lib/python3.6/vmmgr/__init__.py
@@ -216,15 +216,14 @@ class VMMgr:
         app = os.environ['LXC_NAME']
         pid = os.environ['LXC_PID']
         # Lease the first unused IP to the container
-        ip = tools.get_unused_ip()
-        tools.update_hosts_lease(ip, app)
+        ip = tools.update_hosts_lease(app, True)
         tools.set_container_ip(pid, ip)
 
     def unregister_container(self):
         # Extract the variables from values given via lxc.hook.post-stop hook
         app = os.environ['LXC_NAME']
         # Release the container IP
-        tools.update_hosts_lease(None, app)
+        tools.update_hosts_lease(app, False)
         # Remove ephemeral layer data
         tools.clean_ephemeral_layer(app)
 
diff --git a/usr/lib/python3.6/vmmgr/tools.py b/usr/lib/python3.6/vmmgr/tools.py
index c3853cc..9decb39 100644
--- a/usr/lib/python3.6/vmmgr/tools.py
+++ b/usr/lib/python3.6/vmmgr/tools.py
@@ -3,6 +3,7 @@
 import bcrypt
 import dns.exception
 import dns.resolver
+import fcntl
 import os
 import requests
 import shutil
@@ -109,30 +110,27 @@ def shutdown_vm():
 def reboot_vm():
     subprocess.run(['/sbin/reboot'])
 
-def get_unused_ip():
+def update_hosts_lease(app, is_request):
     # This is a poor man's DHCP server which uses /etc/hosts as lease database
     # Leases the first unused IP from range 172.17.0.0/16
-    leased = []
-    with open('/etc/hosts', 'r') as f:
-        for line in f.read().splitlines():
-            if line.startswith('172.17'):
-                ip = line.split()[0].split('.')
-                leased.append(int(ip[2]) * 256 + int(ip[3]))
-    for i in range(1, 65534):
-        if i not in leased:
-            break
-    return '172.17.{}.{}'. format(i // 256, i % 256)
-
-def update_hosts_lease(ip, app):
-    hosts = []
-    with open('/etc/hosts', 'r') as f:
-        for line in f:
-            if not line.strip().endswith(' {}'.format(app)):
-                hosts.append(line)
-    if ip:
-        hosts.append('{} {}\n'.format(ip, app))
-    with open('/etc/hosts', 'w') as f:
-        f.writelines(hosts)
+    ip = None
+    with open('/var/lock/hosts.lock', 'w') as lock:
+        fcntl.lockf(lock, fcntl.LOCK_EX)
+        with open('/etc/hosts', 'r') as f:
+            leases = [l.strip().split(' ', 1) for l in f]
+        if is_request:
+            used_ips = [l[0] lor l in leases]
+            for i in range(2, 65534):
+                ip = '172.17.{}.{}'. format(i // 256, i % 256)
+                if ip not in used_ips:
+                    leases.append([ip, app])
+                    break
+        else:
+            leases = [l for l in leases if l[1] != app]
+        with open('/etc/hosts', 'w') as f:
+            for lease in leases:
+                f.write('{} {}\n'.format(lease[0], lease[1]))
+    return ip
 
 def set_container_ip(pid, ip):
     # Set IP in container based on PID given via lxc.hook.start-host hook