From 375a345820b4d7b6a64b8af3775bc818facd908a Mon Sep 17 00:00:00 2001
From: Maxim Devaev <mdevaev@gmail.com>
Date: Sun, 9 Feb 2025 00:40:48 +0200
Subject: [PATCH] pikvm/pikvm#1204: Configurable global expiration policy

---
 kvmd/apps/__init__.py                |  2 ++
 kvmd/apps/kvmd/__init__.py           |  1 +
 kvmd/apps/kvmd/auth.py               | 42 ++++++++++++++++++++++------
 testenv/tests/apps/kvmd/test_auth.py |  2 ++
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/kvmd/apps/__init__.py b/kvmd/apps/__init__.py
index 7c587c3f..e48079a5 100644
--- a/kvmd/apps/__init__.py
+++ b/kvmd/apps/__init__.py
@@ -65,6 +65,7 @@ from ..validators.basic import valid_string_list
 
 from ..validators.auth import valid_user
 from ..validators.auth import valid_users_list
+from ..validators.auth import valid_expire
 
 from ..validators.os import valid_abs_path
 from ..validators.os import valid_abs_file
@@ -357,6 +358,7 @@ def _get_config_scheme() -> dict:
 
             "auth": {
                 "enabled": Option(True, type=valid_bool),
+                "expire":  Option(0,    type=valid_expire),
 
                 "internal": {
                     "type":        Option("htpasswd"),
diff --git a/kvmd/apps/kvmd/__init__.py b/kvmd/apps/kvmd/__init__.py
index 088a62ef..b87d2dd0 100644
--- a/kvmd/apps/kvmd/__init__.py
+++ b/kvmd/apps/kvmd/__init__.py
@@ -76,6 +76,7 @@ def main(argv: (list[str] | None)=None) -> None:
     KvmdServer(
         auth_manager=AuthManager(
             enabled=config.auth.enabled,
+            expire=config.auth.expire,
             unauth_paths=([] if config.prometheus.auth.enabled else ["/export/prometheus/metrics"]),
 
             internal_type=config.auth.internal.type,
diff --git a/kvmd/apps/kvmd/auth.py b/kvmd/apps/kvmd/auth.py
index 706e770b..3e70a9b7 100644
--- a/kvmd/apps/kvmd/auth.py
+++ b/kvmd/apps/kvmd/auth.py
@@ -43,15 +43,16 @@ class _Session:
     expire_ts: int
 
     def __post_init__(self) -> None:
-        assert self.user.strip()
+        assert self.user == self.user.strip()
         assert self.user
         assert self.expire_ts >= 0
 
 
-class AuthManager:
+class AuthManager:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         enabled: bool,
+        expire: int,
         unauth_paths: list[str],
 
         internal_type: str,
@@ -68,6 +69,11 @@ class AuthManager:
         if not enabled:
             get_logger().warning("AUTHORIZATION IS DISABLED")
 
+        assert expire >= 0
+        self.__expire = expire
+        if expire > 0:
+            get_logger().warning("Maximum user session time is limited: %d seconds", expire)
+
         self.__unauth_paths = frozenset(unauth_paths)  # To speed up
         for path in self.__unauth_paths:
             get_logger().warning("Authorization is disabled for API %r", path)
@@ -132,24 +138,46 @@ class AuthManager:
         assert user
         assert expire >= 0
         assert self.__enabled
+
         if (await self.authorize(user, passwd)):
             token = self.__make_new_token()
             session = _Session(
                 user=user,
-                expire_ts=(0 if expire <= 0 else (self.__get_now_ts() + expire)),
+                expire_ts=self.__make_expire_ts(expire),
             )
             self.__sessions[token] = session
             get_logger().info("Logged in user %r (expire_ts=%d)", session.user, session.expire_ts)
             return token
-        else:
-            return None
+
+        return None
 
     def __make_new_token(self) -> str:
         for _ in range(10):
             token = secrets.token_hex(32)
             if token not in self.__sessions:
                 return token
-        raise AssertionError("Can't generate new unique token")
+        raise RuntimeError("Can't generate new unique token")
+
+    def __make_expire_ts(self, expire: int) -> int:
+        assert expire >= 0
+        assert self.__expire >= 0
+
+        if expire == 0:
+            # The user requested infinite session: apply global expire.
+            # It will allow this (0) or set a limit.
+            expire = self.__expire
+        else:
+            # The user wants a limited session
+            if self.__expire > 0:
+                # If we have a global limit, override the user limit
+                assert expire > 0
+                expire = min(expire, self.__expire)
+
+        if expire > 0:
+            return (self.__get_now_ts() + expire)
+
+        assert expire == 0
+        return 0
 
     def __get_now_ts(self) -> int:
         return int(time.monotonic())
@@ -171,12 +199,10 @@ class AuthManager:
         if session is not None:
             if session.expire_ts <= 0:
                 # Infinite session
-                assert session.user
                 return session.user
             else:
                 # Limited session
                 if self.__get_now_ts() < session.expire_ts:
-                    assert session.user
                     return session.user
                 else:
                     del self.__sessions[token]
diff --git a/testenv/tests/apps/kvmd/test_auth.py b/testenv/tests/apps/kvmd/test_auth.py
index d6183a39..43b1c530 100644
--- a/testenv/tests/apps/kvmd/test_auth.py
+++ b/testenv/tests/apps/kvmd/test_auth.py
@@ -61,6 +61,7 @@ async def _get_configured_manager(
 
     manager = AuthManager(
         enabled=True,
+        expire=0,
         unauth_paths=unauth_paths,
 
         internal_type="htpasswd",
@@ -250,6 +251,7 @@ async def test_ok__disabled() -> None:
     try:
         manager = AuthManager(
             enabled=False,
+            expire=0,
             unauth_paths=[],
 
             internal_type="foobar",