Implemented support for limiting number of clients per user

The setting is global for all users and optional.

api/__init__.py

@@ -6,8 +6,9 @@ import re
 
 from flask import jsonify, Flask, redirect, request, Response
 
-from .wireguard import WireGuard
+from .exceptions import ClientLimitError
 from .oauth import Oauth
+from .wireguard import WireGuard
 
 
 login_path = '/login'
@@ -119,6 +120,8 @@ def create_config(config_id: str) -> dict:
                                      creation_time)
     except FileExistsError:
         return fail('Id already in use')
+    except ClientLimitError as e:
+        return fail(e.message)
     return get_config(config_id)
 
 @app.route('/configs/<config_id>/update', methods=['POST'])

api/exceptions.py

@@ -0,0 +1,3 @@
+class ClientLimitError(Exception):
+    def __init__(self, message):
+        self.message = message

api/wireguard.py

@@ -7,6 +7,8 @@ import ipaddress
 import json
 import subprocess
 
+from .exceptions import ClientLimitError
+
 
 confsuffix = '.conf'
 serversuffix = '.serverconf'
@@ -37,9 +39,13 @@ def call_with_lock(func: callable, args: list=[], timeout: float=None):
             with open(lockfile, 'x') as lf:
                 print(f'Lock successful after {slept}s, '
                       f'tried on {sleep_time}s intervals.')
-                result = func(*args)
-                lockfile.unlink()
-                return result
+                try:
+                    result = func(*args)
+                    return result
+                except Exception as e:
+                    raise e
+                finally:
+                    lockfile.unlink()
         except FileExistsError:
             slept += sleep_time
             sleep(sleep_time)
@@ -123,6 +129,7 @@ class WireGuard:
         self.dns_server = ipaddress.ip_address(config['dns_server'])
         self.client_network = ipaddress.ip_network(config['client_network'])
         self.configs_base = Path(config['configs_base'])
+        self.max_clients = config.getint('user_client_limit', fallback=0)
 
         self.server_config_base = None
         if 'server_extra_config' in config.keys():
@@ -163,6 +170,11 @@ class WireGuard:
                 return addr
         raise Exception('No addresses available')
 
+    def get_user_client_count(self) -> int:
+        # Path.glob() returns a generator, so we have to count by iteration
+        metafiles = self.user_base.glob(f'*{metasuffix}')
+        return sum(1 for f in metafiles)
+
     def filepath(self, config_filename: str) -> Path:
         if self.user_base is None:
             raise Exception('Not properly initialized')
@@ -210,6 +222,9 @@ class WireGuard:
                                       name: str,
                                       description: str,
                                       creation_time: datetime) -> ipaddress:
+        if self.max_clients \
+           and self.get_user_client_count() >= self.max_clients:
+            raise ClientLimitError('client limit reached')
         client_privkey, client_pubkey = generate_keypair()
         client_ip = self.get_free_ip()
         with open(self.config_filepath(config_id), 'x') as cf, \

config.ini.example

@@ -24,6 +24,9 @@ client_network = a.network.in.cidr/notation
 # Optional.
 server_extra_config = path/to/a/conf/fragment
 
+# The maximum number of clients to allow per user.
+# Optional, defaults to unlimited, equivalent to setting this value to 0.
+user_client_limit = 3
 
 
 [security]