From 718d11b4c0f1baf30afc6da4ec33fa37b0bfdd58 Mon Sep 17 00:00:00 2001 From: Walter Oggioni Date: Sun, 25 Oct 2026 21:24:54 +0000 Subject: [PATCH] Offload certificate key generation to worker threads --- .../packages/kca-app/src/kca/routes/keys.py | 17 ++-- .../kca-app/src/kca/routes/requests.py | 16 ++-- server/packages/kca-app/src/kca/service.py | 22 ++++- server/tests/test_routes_extra.py | 86 +++++++++++++++++++ 4 files changed, 117 insertions(+), 24 deletions(-) diff --git a/server/packages/kca-app/src/kca/routes/keys.py b/server/packages/kca-app/src/kca/routes/keys.py index 58eb592..09e1c00 100644 --- a/server/packages/kca-app/src/kca/routes/keys.py +++ b/server/packages/kca-app/src/kca/routes/keys.py @@ -1,7 +1,6 @@ """Key generation, detail, rename/delete and export.""" from __future__ import annotations -import asyncio import base64 from typing import TYPE_CHECKING, Optional @@ -18,6 +17,7 @@ from ..models import Certificate, Key, Request from ..openapi import error_response, json_response, schema_ref from ..service import ( ValidationError, + generate_key as generate_private_key, parse_key_generation, parse_uuid, save_key, @@ -30,9 +30,6 @@ from ..service import ( if TYPE_CHECKING: from ..mixin import Kca -# Key generation is CPU-bound; this bounds concurrent RSA/DSA jobs. -_keygen_semaphore = asyncio.Semaphore(2) - async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str): user = await authenticated_user(kca.oidc, ctx) @@ -85,13 +82,11 @@ def register(app: KayaApp, kca: "Kca") -> None: await send_error(ctx, 400, str(exc)) return try: - async with _keygen_semaphore: - generated = await asyncio.to_thread( - key_core.generate_key, - params["key_type"], - params["size"], - params["curve"], - ) + generated = await generate_private_key( + params["key_type"], + params["size"], + params["curve"], + ) key = await save_key( vault, generated, diff --git a/server/packages/kca-app/src/kca/routes/requests.py b/server/packages/kca-app/src/kca/routes/requests.py index e322e71..b17ec0a 100644 --- a/server/packages/kca-app/src/kca/routes/requests.py +++ b/server/packages/kca-app/src/kca/routes/requests.py @@ -1,7 +1,6 @@ """Certificate signing requests: list, create, detail, export and transforms.""" from __future__ import annotations -import asyncio import base64 from datetime import datetime, timezone from typing import TYPE_CHECKING, Optional @@ -21,6 +20,7 @@ from ..models import Key, Request, Template from ..openapi import error_response, json_response, schema_ref from ..service import ( ValidationError, + generate_key, parse_extensions, parse_subject, parse_uuid, @@ -36,8 +36,6 @@ from ..service import ( if TYPE_CHECKING: from ..mixin import Kca -_keygen_semaphore = asyncio.Semaphore(2) - async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str): user = await authenticated_user(kca.oidc, ctx) @@ -98,13 +96,11 @@ def register(app: KayaApp, kca: "Kca") -> None: key_obj = await load_private_key(key_model) elif isinstance(body.get("new_key"), dict): new_key = body["new_key"] - async with _keygen_semaphore: - generated = await asyncio.to_thread( - key_core.generate_key, - str(new_key.get("key_type") or "rsa"), - int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE), - new_key.get("curve"), - ) + generated = await generate_key( + str(new_key.get("key_type") or "rsa"), + int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE), + new_key.get("curve"), + ) key_model = await save_key( vault, generated, diff --git a/server/packages/kca-app/src/kca/service.py b/server/packages/kca-app/src/kca/service.py index 9f16d6d..a337126 100644 --- a/server/packages/kca-app/src/kca/service.py +++ b/server/packages/kca-app/src/kca/service.py @@ -1,6 +1,7 @@ """Business logic: validation, persistence, serialization and CA operations.""" from __future__ import annotations +import asyncio import hashlib import secrets from datetime import datetime, timedelta, timezone @@ -471,6 +472,20 @@ async def serialize_vault_detail(vault: Vault) -> Dict[str, Any]: # --------------------------------------------------------------------------- +# Key generation is CPU-bound; bound concurrent jobs across all routes. +_keygen_semaphore = asyncio.Semaphore(2) + + +async def generate_key( + key_type: str, + size: int = key_core.DEFAULT_KEY_SIZE, + curve: Optional[str] = None, +) -> key_core.PrivateKey: + """Generate a private key without blocking the event loop.""" + async with _keygen_semaphore: + return await asyncio.to_thread(key_core.generate_key, key_type, size, curve) + + def parse_key_generation(payload: Mapping[str, Any]) -> Dict[str, Any]: key_type = str(payload.get("key_type") or payload.get("type") or "rsa").lower() if key_type not in key_core.KEY_TYPES: @@ -708,7 +723,7 @@ async def issue_certificate( key_obj = await load_private_key(key_model) elif new_key_payload: try: - generated = key_core.generate_key( + generated = await generate_key( str(new_key_payload.get("key_type") or "rsa"), int(new_key_payload.get("size") or key_core.DEFAULT_KEY_SIZE), new_key_payload.get("curve"), @@ -873,13 +888,13 @@ async def preview_certificate(vault: Vault, payload: Mapping[str, Any]) -> Dict[ else: new_key = payload.get("new_key") if isinstance(new_key, Mapping): - generated = key_core.generate_key( + generated = await generate_key( str(new_key.get("key_type") or "ec"), int(new_key.get("size") or 2048), new_key.get("curve") or "secp256r1", ) else: - generated = key_core.generate_key("ec", curve="secp256r1") + generated = await generate_key("ec", curve="secp256r1") public_key = generated.public_key() extensions_spec = parse_extensions(payload.get("extensions")) @@ -1446,6 +1461,7 @@ __all__ = [ "find_key", "find_request", "generate_crl", + "generate_key", "import_parsed_items", "iso", "issue_certificate", diff --git a/server/tests/test_routes_extra.py b/server/tests/test_routes_extra.py index 00f6b00..e3bd8a8 100644 --- a/server/tests/test_routes_extra.py +++ b/server/tests/test_routes_extra.py @@ -55,6 +55,92 @@ async def _cert(client, vault_id: str, key_id: str, name: str, ca: bool = False, return response.json() +class KeyGenerationTests(unittest.TestCase): + @async_test + async def test_shared_generation_is_offloaded_and_bounded(self) -> None: + import asyncio + import threading + from unittest.mock import patch + + from kca import service + + active = 0 + max_active = 0 + worker_threads: set[int] = set() + lock = threading.Lock() + barrier = threading.Barrier(2) + + def fake_generate(*args: object) -> object: + nonlocal active, max_active + with lock: + active += 1 + max_active = max(max_active, active) + worker_threads.add(threading.get_ident()) + try: + barrier.wait(timeout=5) + finally: + with lock: + active -= 1 + return object() + + with ( + patch.object(service, "_keygen_semaphore", asyncio.Semaphore(2)), + patch.object(service.key_core, "generate_key", new=fake_generate), + ): + results = await asyncio.gather(*(service.generate_key("rsa", 8192) for _ in range(6))) + + self.assertEqual(len(results), 6) + self.assertEqual(max_active, 2) + self.assertNotIn(threading.get_ident(), worker_threads) + + @async_test + async def test_certificate_generation_is_offloaded(self) -> None: + import threading + from unittest.mock import patch + + from kca.core import keys as key_core + + event_loop_thread = threading.get_ident() + generation_threads: list[int] = [] + generated_key = key_core.generate_key("rsa", 1024) + + def fake_generate(*args: object) -> object: + generation_threads.append(threading.get_ident()) + return generated_key + + async with api_client() as client: + with oidc_user("alice"): + vault_id = await _vault(client) + with patch.object(key_core, "generate_key", new=fake_generate): + response = await client.post( + f"/api/vaults/{vault_id}/certificates/preview", + json={ + "subject": _subject("preview"), + "new_key": {"key_type": "rsa", "size": 8192}, + }, + ) + assert response.status_code == 200, response.text + + response = await client.post( + f"/api/vaults/{vault_id}/certificates/preview", + json={"subject": _subject("default preview")}, + ) + assert response.status_code == 200, response.text + + response = await client.post( + f"/api/vaults/{vault_id}/certificates", + json={ + "name": "issued", + "subject": _subject("issued"), + "new_key": {"key_type": "rsa", "size": 8192, "name": "issued key"}, + }, + ) + assert response.status_code == 201, response.text + + self.assertEqual(len(generation_threads), 3) + self.assertNotIn(event_loop_thread, generation_threads) + + class RouteExtraTests(unittest.TestCase): @async_test async def test_csr_creation_and_signing(self) -> None: