From 5424d7a29425ee930b3b297821bccb5e98737b2b Mon Sep 17 00:00:00 2001 From: Morten Hansen Date: Fri, 23 Jan 2026 11:51:50 +0100 Subject: [PATCH] feat: add service key authentication for registration Add optional X-Service-Key header authentication to the registration module, enabling services to register with orchestrators that require authentication. - Add service_key and service_key_env parameters to with_registration() - Include X-Service-Key header in register, keepalive, and deregister requests - Default env var: SERVICEKIT_REGISTRATION_KEY - Direct parameter takes precedence over environment variable - Backwards compatible: no header sent if key not configured - Bump version to 0.7.0 --- examples/registration/README.md | 48 ++++++ pyproject.toml | 2 +- src/servicekit/api/registration.py | 51 ++++++- src/servicekit/api/service_builder.py | 12 ++ tests/test_api_registration.py | 202 ++++++++++++++++++++++++++ uv.lock | 2 +- 6 files changed, 311 insertions(+), 6 deletions(-) diff --git a/examples/registration/README.md b/examples/registration/README.md index a91425c..08891a7 100644 --- a/examples/registration/README.md +++ b/examples/registration/README.md @@ -203,6 +203,7 @@ Services can configure keepalive behavior: | `SERVICEKIT_ORCHESTRATOR_URL` | (required) | Orchestrator registration endpoint | | `SERVICEKIT_HOST` | auto-detected | Service hostname (override auto-detection) | | `SERVICEKIT_PORT` | 8000 | Service port | +| `SERVICEKIT_REGISTRATION_KEY` | (optional) | Service key for authenticated registration | ### Builder Configuration @@ -218,9 +219,56 @@ Services can configure keepalive behavior: retry_delay=2.0, fail_on_error=False, # Don't abort on failure timeout=10.0, + service_key=None, # Service key for authentication + service_key_env="SERVICEKIT_REGISTRATION_KEY", ) ``` +### Service Key Authentication + +When registering with an orchestrator that requires authentication, you can configure a service key that will be sent as an `X-Service-Key` header with all registration requests (register, keepalive pings, deregister). + +**Using environment variable (recommended for production):** + +```yaml +services: + my-service: + image: my-service:latest + environment: + SERVICEKIT_ORCHESTRATOR_URL: http://orchestrator:9000/services/$register + SERVICEKIT_REGISTRATION_KEY: ${REGISTRATION_SECRET} +``` + +**Using direct parameter (for testing):** + +```python +app = ( + BaseServiceBuilder(info=ServiceInfo(display_name="My Service")) + .with_registration( + orchestrator_url="http://orchestrator:9000/services/$register", + service_key="my-secret-key", + ) + .build() +) +``` + +**Using custom environment variable name:** + +```python +app = ( + BaseServiceBuilder(info=ServiceInfo(display_name="My Service")) + .with_registration( + service_key_env="MY_APP_REGISTRATION_KEY", + ) + .build() +) +``` + +The service key is included in: +- Initial registration (`POST /services/$register`) +- Keepalive pings (`PUT /services/{id}/$ping`) +- Deregistration (`DELETE /services/{id}`) + ## Examples ### Basic Registration (main.py) diff --git a/pyproject.toml b/pyproject.toml index fef2bb1..99b80b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "servicekit" -version = "0.6.0" +version = "0.7.0" description = "Async SQLAlchemy framework with FastAPI integration - reusable foundation for building data services" readme = "README.md" authors = [{ name = "Morten Hansen", email = "morten@winterop.com" }] diff --git a/src/servicekit/api/registration.py b/src/servicekit/api/registration.py index 3075da4..5c8fc53 100644 --- a/src/servicekit/api/registration.py +++ b/src/servicekit/api/registration.py @@ -16,6 +16,14 @@ _keepalive_task: asyncio.Task | None = None _service_id: str | None = None _ping_url: str | None = None +_service_key: str | None = None + + +def _resolve_service_key(service_key: str | None, service_key_env: str) -> str | None: + """Resolve service key from parameter or environment variable.""" + if service_key: + return service_key + return os.getenv(service_key_env) async def register_service( @@ -31,8 +39,17 @@ async def register_service( retry_delay: float = 2.0, fail_on_error: bool = False, timeout: float = 10.0, + service_key: str | None = None, + service_key_env: str = "SERVICEKIT_REGISTRATION_KEY", ) -> dict[str, Any] | None: """Register service with orchestrator for service discovery and return registration info.""" + # Resolve service key for authentication + resolved_service_key = _resolve_service_key(service_key, service_key_env) + + # Store globally for keepalive + global _service_key + _service_key = resolved_service_key + # Resolve orchestrator URL resolved_orchestrator_url = orchestrator_url or os.getenv(orchestrator_url_env) if not resolved_orchestrator_url: @@ -114,12 +131,18 @@ async def register_service( # Registration with retry logic last_error: Exception | None = None + # Build headers with optional service key + headers: dict[str, str] = {} + if resolved_service_key: + headers["X-Service-Key"] = resolved_service_key + for attempt in range(1, max_retries + 1): try: async with httpx.AsyncClient(timeout=timeout) as client: response = await client.post( resolved_orchestrator_url, json=payload, + headers=headers if headers else None, ) response.raise_for_status() @@ -187,16 +210,21 @@ async def register_service( return None -async def _keepalive_loop(ping_url: str, interval: float, timeout: float) -> None: +async def _keepalive_loop(ping_url: str, interval: float, timeout: float, service_key: str | None) -> None: """Background task to periodically ping the orchestrator.""" logger.info("keepalive.started", ping_url=ping_url, interval_seconds=interval) + # Build headers with optional service key + headers: dict[str, str] = {} + if service_key: + headers["X-Service-Key"] = service_key + while True: try: await asyncio.sleep(interval) async with httpx.AsyncClient(timeout=timeout) as client: - response = await client.put(ping_url) + response = await client.put(ping_url, headers=headers if headers else None) response.raise_for_status() response_data = response.json() @@ -226,6 +254,8 @@ async def start_keepalive( ping_url: str, interval: float = 10.0, timeout: float = 10.0, + service_key: str | None = None, + service_key_env: str = "SERVICEKIT_REGISTRATION_KEY", ) -> None: """Start background keepalive task to ping orchestrator.""" global _keepalive_task @@ -234,7 +264,10 @@ async def start_keepalive( logger.warning("keepalive.already_running") return - _keepalive_task = asyncio.create_task(_keepalive_loop(ping_url, interval, timeout)) + # Resolve service key (use global from registration if not provided) + resolved_service_key = _resolve_service_key(service_key, service_key_env) or _service_key + + _keepalive_task = asyncio.create_task(_keepalive_loop(ping_url, interval, timeout, resolved_service_key)) logger.info("keepalive.task_started", ping_url=ping_url, interval_seconds=interval) @@ -257,6 +290,8 @@ async def deregister_service( service_id: str, orchestrator_url: str, timeout: float = 10.0, + service_key: str | None = None, + service_key_env: str = "SERVICEKIT_REGISTRATION_KEY", ) -> None: """Explicitly deregister service from orchestrator.""" # Build deregister URL from orchestrator base URL @@ -265,9 +300,17 @@ async def deregister_service( base_url = orchestrator_url.replace("/$register", "") deregister_url = f"{base_url}/{service_id}" + # Resolve service key (use global from registration if not provided) + resolved_service_key = _resolve_service_key(service_key, service_key_env) or _service_key + + # Build headers with optional service key + headers: dict[str, str] = {} + if resolved_service_key: + headers["X-Service-Key"] = resolved_service_key + try: async with httpx.AsyncClient(timeout=timeout) as client: - response = await client.delete(deregister_url) + response = await client.delete(deregister_url, headers=headers if headers else None) response.raise_for_status() logger.info( diff --git a/src/servicekit/api/service_builder.py b/src/servicekit/api/service_builder.py index 225cd7e..f9a45ff 100644 --- a/src/servicekit/api/service_builder.py +++ b/src/servicekit/api/service_builder.py @@ -93,6 +93,8 @@ class _RegistrationOptions: enable_keepalive: bool keepalive_interval: float auto_deregister: bool + service_key: str | None + service_key_env: str class ServiceInfo(BaseModel): @@ -321,6 +323,8 @@ def with_registration( enable_keepalive: bool = True, keepalive_interval: float = 10.0, auto_deregister: bool = True, + service_key: str | None = None, + service_key_env: str = "SERVICEKIT_REGISTRATION_KEY", ) -> Self: """Enable service registration with orchestrator for service discovery.""" self._registration_options = _RegistrationOptions( @@ -337,6 +341,8 @@ def with_registration( enable_keepalive=enable_keepalive, keepalive_interval=keepalive_interval, auto_deregister=auto_deregister, + service_key=service_key, + service_key_env=service_key_env, ) return self @@ -666,6 +672,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: retry_delay=registration_options.retry_delay, fail_on_error=registration_options.fail_on_error, timeout=registration_options.timeout, + service_key=registration_options.service_key, + service_key_env=registration_options.service_key_env, ) # Start keepalive if registration succeeded and enabled @@ -676,6 +684,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: ping_url=ping_url, interval=registration_options.keepalive_interval, timeout=registration_options.timeout, + service_key=registration_options.service_key, + service_key_env=registration_options.service_key_env, ) try: @@ -698,6 +708,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: service_id=service_id, orchestrator_url=orchestrator_url, timeout=registration_options.timeout, + service_key=registration_options.service_key, + service_key_env=registration_options.service_key_env, ) for hook in shutdown_hooks: diff --git a/tests/test_api_registration.py b/tests/test_api_registration.py index bf3356a..7e9c7b6 100644 --- a/tests/test_api_registration.py +++ b/tests/test_api_registration.py @@ -642,3 +642,205 @@ async def test_deregister_service_handles_errors(): orchestrator_url="http://orchestrator:9000/services/$register", timeout=5.0, ) + + +@pytest.mark.asyncio +async def test_registration_with_service_key_parameter(): + """Test registration sends X-Service-Key header when service_key parameter is provided.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"id": "01K83B5V85PQZ1HTH4DQ7NC9JM", "status": "registered"}) + + with patch("httpx.AsyncClient") as mock_client: + mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) + + info = ServiceInfo(display_name="Test Service", version="1.0.0") + + await register_service( + orchestrator_url="http://orchestrator:9000/services/$register", + host="test-service", + port=8000, + info=info, + service_key="my-secret-key", + ) + + # Verify POST was called with X-Service-Key header + call_args = mock_client.return_value.__aenter__.return_value.post.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "my-secret-key" + + +@pytest.mark.asyncio +async def test_registration_with_service_key_from_env(): + """Test registration sends X-Service-Key header from default environment variable.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"id": "01K83B5V85PQZ1HTH4DQ7NC9JM", "status": "registered"}) + + with ( + patch("httpx.AsyncClient") as mock_client, + patch.dict(os.environ, {"SERVICEKIT_REGISTRATION_KEY": "env-secret-key"}), + ): + mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) + + info = ServiceInfo(display_name="Test Service", version="1.0.0") + + await register_service( + orchestrator_url="http://orchestrator:9000/services/$register", + host="test-service", + port=8000, + info=info, + ) + + # Verify POST was called with X-Service-Key header from env + call_args = mock_client.return_value.__aenter__.return_value.post.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "env-secret-key" + + +@pytest.mark.asyncio +async def test_registration_with_custom_service_key_env(): + """Test registration uses custom environment variable name for service key.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"id": "01K83B5V85PQZ1HTH4DQ7NC9JM", "status": "registered"}) + + with ( + patch("httpx.AsyncClient") as mock_client, + patch.dict(os.environ, {"MY_CUSTOM_SERVICE_KEY": "custom-env-key"}), + ): + mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) + + info = ServiceInfo(display_name="Test Service", version="1.0.0") + + await register_service( + orchestrator_url="http://orchestrator:9000/services/$register", + host="test-service", + port=8000, + info=info, + service_key_env="MY_CUSTOM_SERVICE_KEY", + ) + + # Verify POST was called with X-Service-Key header from custom env + call_args = mock_client.return_value.__aenter__.return_value.post.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "custom-env-key" + + +@pytest.mark.asyncio +async def test_registration_without_service_key(): + """Test registration works without service key (backwards compatibility).""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"id": "01K83B5V85PQZ1HTH4DQ7NC9JM", "status": "registered"}) + + with patch("httpx.AsyncClient") as mock_client, patch.dict(os.environ, {}, clear=True): + mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) + + info = ServiceInfo(display_name="Test Service", version="1.0.0") + + await register_service( + orchestrator_url="http://orchestrator:9000/services/$register", + host="test-service", + port=8000, + info=info, + ) + + # Verify POST was called without headers (or with None) + call_args = mock_client.return_value.__aenter__.return_value.post.call_args + headers = call_args[1].get("headers") + assert headers is None + + +@pytest.mark.asyncio +async def test_service_key_parameter_overrides_env(): + """Test that service_key parameter takes precedence over environment variable.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"id": "01K83B5V85PQZ1HTH4DQ7NC9JM", "status": "registered"}) + + with ( + patch("httpx.AsyncClient") as mock_client, + patch.dict(os.environ, {"SERVICEKIT_REGISTRATION_KEY": "env-key"}), + ): + mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) + + info = ServiceInfo(display_name="Test Service", version="1.0.0") + + await register_service( + orchestrator_url="http://orchestrator:9000/services/$register", + host="test-service", + port=8000, + info=info, + service_key="param-key", # Should override env + ) + + # Verify POST was called with parameter value, not env value + call_args = mock_client.return_value.__aenter__.return_value.post.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "param-key" + + +@pytest.mark.asyncio +async def test_keepalive_sends_service_key(): + """Test that keepalive pings include X-Service-Key header.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value={"status": "alive", "last_ping_at": "2025-01-01T00:00:00Z"}) + + with patch("httpx.AsyncClient") as mock_client: + mock_client.return_value.__aenter__.return_value.put = AsyncMock(return_value=mock_response) + + # Start keepalive with service key + await start_keepalive( + ping_url="http://orchestrator:9000/services/test/$ping", + interval=0.1, + timeout=5.0, + service_key="keepalive-secret-key", + ) + + # Wait for at least one ping + await asyncio.sleep(0.15) + + # Verify PUT was called with X-Service-Key header + call_args = mock_client.return_value.__aenter__.return_value.put.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "keepalive-secret-key" + + # Clean up + await stop_keepalive() + + +@pytest.mark.asyncio +async def test_deregister_sends_service_key(): + """Test that deregistration includes X-Service-Key header.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + + with patch("httpx.AsyncClient") as mock_client: + mock_client.return_value.__aenter__.return_value.delete = AsyncMock(return_value=mock_response) + + await deregister_service( + service_id="01K83B5V85PQZ1HTH4DQ7NC9JM", + orchestrator_url="http://orchestrator:9000/services/$register", + timeout=5.0, + service_key="deregister-secret-key", + ) + + # Verify DELETE was called with X-Service-Key header + call_args = mock_client.return_value.__aenter__.return_value.delete.call_args + headers = call_args[1].get("headers") + assert headers is not None + assert headers.get("X-Service-Key") == "deregister-secret-key" diff --git a/uv.lock b/uv.lock index dd8de44..8bbbf39 100644 --- a/uv.lock +++ b/uv.lock @@ -1129,7 +1129,7 @@ wheels = [ [[package]] name = "servicekit" -version = "0.6.0" +version = "0.7.0" source = { editable = "." } dependencies = [ { name = "aiosqlite" },