feat: Implement Image Rescanning using Harbor Webhook API

changes/3116.feature.md

@@ -0,0 +1 @@
+Implement Image Rescanning using Harbor Webhook API.

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

src/ai/backend/manager/api/container_registry.py

@@ -1,27 +1,34 @@
 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
 
 if TYPE_CHECKING:
     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,13 +91,112 @@ 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]]:
     app = web.Application()
     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, []

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"
 

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 = {