diff --git a/changes/3116.feature.md b/changes/3116.feature.md new file mode 100644 index 0000000000..455c5a4d9f --- /dev/null +++ b/changes/3116.feature.md @@ -0,0 +1 @@ +Implement Image Rescanning using Harbor Webhook API. diff --git a/src/ai/backend/manager/api/auth.py b/src/ai/backend/manager/api/auth.py index 094a364942..d933c8c57d 100644 --- a/src/ai/backend/manager/api/auth.py +++ b/src/ai/backend/manager/api/auth.py @@ -428,6 +428,16 @@ async def auth_middleware(request: web.Request, handler) -> web.StreamResponse: Fetches user information and sets up keypair, user, and is_authorized attributes. """ + allow_list = request.app["auth_middleware_allowlist"] + + if any(request.path.startswith(path) for path in allow_list): + request["is_authorized"] = False + request["is_admin"] = False + request["is_superadmin"] = False + request["keypair"] = None + request["user"] = None + return await handler(request) + # This is a global middleware: request.app is the root app. root_ctx: RootContext = request.app["_root.context"] request["is_authorized"] = False diff --git a/src/ai/backend/manager/api/container_registry.py b/src/ai/backend/manager/api/container_registry.py index c28d4c575f..3613817295 100644 --- a/src/ai/backend/manager/api/container_registry.py +++ b/src/ai/backend/manager/api/container_registry.py @@ -1,19 +1,26 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Iterable, Tuple +from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple import aiohttp_cors import sqlalchemy as sa import trafaret as t from aiohttp import web from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession from ai.backend.common import validators as tx from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.api.exceptions import ( + ContainerRegistryWebhookAuthorizationFailed, + InternalServerError, +) +from ai.backend.manager.container_registry.harbor import HarborRegistry_v2 from ai.backend.manager.models.association_container_registries_groups import ( AssociationContainerRegistriesGroupsRow, ) +from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType from .exceptions import ContainerRegistryNotFound, GenericBadRequest @@ -21,7 +28,7 @@ from .context import RootContext from .auth import superadmin_required -from .manager import READ_ALLOWED, server_status_required +from .manager import ALL_ALLOWED, READ_ALLOWED, server_status_required from .types import CORSOptions, WebMiddleware from .utils import check_api_params @@ -84,6 +91,103 @@ async def disassociate_with_group(request: web.Request, params: Any) -> web.Resp return web.json_response({}) +async def _get_registry_row_matching_url( + db_sess: AsyncSession, registry_url: str, project: str +) -> ContainerRegistryRow: + query = sa.select(ContainerRegistryRow).where( + (ContainerRegistryRow.type == ContainerRegistryType.HARBOR2) + & (ContainerRegistryRow.url.like(f"%{registry_url}%")) + & (ContainerRegistryRow.project == project) + ) + result = await db_sess.execute(query) + return result.scalars().one_or_none() + + +def _is_authorized_harbor_webhook_request( + auth_header: Optional[str], registry_row: ContainerRegistryRow +) -> bool: + if auth_header: + extra = registry_row.extra or {} + return extra.get("webhook_auth_header") == auth_header + return True + + +async def _handle_harbor_webhook_event( + root_ctx: RootContext, + event_type: str, + registry_row: ContainerRegistryRow, + project: str, + img_name: str, + tag: str, +) -> None: + match event_type: + # Perform image rescan only for events that require it. + case "PUSH_ARTIFACT": + await _handle_push_artifact_event(root_ctx, registry_row, project, img_name, tag) + case _: + log.debug( + 'Ignore harbor webhook event: "{}". Recommended to modify the webhook config to not subscribe to this event type.', + event_type, + ) + + +async def _handle_push_artifact_event( + root_ctx: RootContext, registry_row: ContainerRegistryRow, project: str, img_name: str, tag: str +) -> None: + scanner = HarborRegistry_v2(root_ctx.db, registry_row.registry_name, registry_row) + await scanner.scan_single_ref(f"{project}/{img_name}:{tag}") + + +@server_status_required(ALL_ALLOWED) +@check_api_params( + t.Dict({ + "type": t.String, + "event_data": t.Dict({ + "resources": t.List( + t.Dict({ + "resource_url": t.String, + "tag": t.String, + }).allow_extra("*") + ), + "repository": t.Dict({ + "namespace": t.String, + "name": t.String, + }).allow_extra("*"), + }).allow_extra("*"), + }).allow_extra("*") +) +async def harbor_webhook_handler(request: web.Request, params: Any) -> web.Response: + auth_header = request.headers.get("Authorization", None) + event_type = params["type"] + resources = params["event_data"]["resources"] + project = params["event_data"]["repository"]["namespace"] + img_name = params["event_data"]["repository"]["name"] + log.info("HARBOR_WEBHOOK_HANDLER (event_type:{})", event_type) + + root_ctx: RootContext = request.app["_root.context"] + async with root_ctx.db.begin_session() as db_sess: + for resource in resources: + resource_url = resource["resource_url"] + registry_url = resource_url.split("/")[0] + + registry_row = await _get_registry_row_matching_url(db_sess, registry_url, project) + if not registry_row: + raise InternalServerError( + extra_msg=f"Harbor webhook triggered, but the matching container registry row not found! (registry_url: {registry_url}, project: {project})", + ) + + if not _is_authorized_harbor_webhook_request(auth_header, registry_row): + raise ContainerRegistryWebhookAuthorizationFailed( + extra_msg=f"Unauthorized webhook request (registry: {registry_row.registry_name}, project: {project})", + ) + + await _handle_harbor_webhook_event( + root_ctx, event_type, registry_row, project, img_name, resource["tag"] + ) + + return web.json_response({}) + + def create_app( default_cors_options: CORSOptions, ) -> Tuple[web.Application, Iterable[WebMiddleware]]: @@ -91,6 +195,8 @@ def create_app( app["api_versions"] = (1, 2, 3, 4, 5) app["prefix"] = "container-registries" cors = aiohttp_cors.setup(app, defaults=default_cors_options) + + cors.add(app.router.add_route("POST", "/webhook/harbor", harbor_webhook_handler)) cors.add(app.router.add_route("POST", "/associate-with-group", associate_with_group)) cors.add(app.router.add_route("POST", "/disassociate-with-group", disassociate_with_group)) return app, [] diff --git a/src/ai/backend/manager/api/exceptions.py b/src/ai/backend/manager/api/exceptions.py index 8db4ee7db9..02097e9be1 100644 --- a/src/ai/backend/manager/api/exceptions.py +++ b/src/ai/backend/manager/api/exceptions.py @@ -96,6 +96,7 @@ def __init__( extra_msg: Optional[str] = None, extra_data: Optional[Any] = None, *, + status_code: int = 404, object_name: Optional[str] = None, **kwargs, ) -> None: @@ -194,6 +195,11 @@ class GraphQLError(BackendError, web.HTTPBadRequest): error_title = "GraphQL-generated error." +class ContainerRegistryWebhookAuthorizationFailed(BackendError, web.HTTPUnauthorized): + error_type = "https://api.backend.ai/probs/webhook/auth-failed" + error_title = "Container Registry Webhook authorization failed." + + class InstanceNotFound(ObjectNotFound): object_name = "agent instance" diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py index f10ff47e20..321d6abc14 100644 --- a/src/ai/backend/manager/server.py +++ b/src/ai/backend/manager/server.py @@ -801,6 +801,14 @@ def build_root_app( loop = asyncio.get_running_loop() loop.set_exception_handler(global_exception_handler) app["_root.context"] = root_ctx + + # If the request path starts with the following route, the auth_middleware is bypassed. + # In this case, all authentication flags are turned off. + # Used in special cases where the request headers cannot be modified. + app["auth_middleware_allowlist"] = [ + "/container-registries/webhook", + ] + root_ctx.local_config = local_config root_ctx.pidx = pidx root_ctx.cors_options = {