|
| 1 | +# Copyright 2025 The Feast Authors |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +""" |
| 16 | +Structured audit logging for the Feast feature server. |
| 17 | +
|
| 18 | +Emits JSONL audit events for MCP tool calls, REST requests, and |
| 19 | +authentication/authorization decisions. Sensitive payloads (tokens, |
| 20 | +entity rows, feature values) are never included. |
| 21 | +""" |
| 22 | + |
| 23 | +import abc |
| 24 | +import contextvars |
| 25 | +import logging |
| 26 | +import sys |
| 27 | +import threading |
| 28 | +import uuid |
| 29 | +from datetime import datetime, timezone |
| 30 | +from typing import TYPE_CHECKING, Any, Dict, List, Optional |
| 31 | + |
| 32 | +if TYPE_CHECKING: |
| 33 | + from feast.infra.feature_servers.base_config import AuditLoggingConfig |
| 34 | + |
| 35 | +from pydantic import BaseModel, Field |
| 36 | + |
| 37 | +logger = logging.getLogger(__name__) |
| 38 | + |
| 39 | +# ContextVar used by the MCP audit wrapper to propagate a single request_id |
| 40 | +# into the internal REST call so that ``mcp.tools.call`` and |
| 41 | +# ``http.request`` events share the same identifier in SIEM. |
| 42 | +mcp_audit_request_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( |
| 43 | + "mcp_audit_request_id", default=None |
| 44 | +) |
| 45 | + |
| 46 | + |
| 47 | +# --------------------------------------------------------------------------- |
| 48 | +# Audit event schema |
| 49 | +# --------------------------------------------------------------------------- |
| 50 | + |
| 51 | + |
| 52 | +class AuditPrincipal(BaseModel): |
| 53 | + username: str = "" |
| 54 | + roles: List[str] = Field(default_factory=list) |
| 55 | + auth_type: str = "" |
| 56 | + |
| 57 | + |
| 58 | +class AuditSource(BaseModel): |
| 59 | + ip: str = "" |
| 60 | + transport: str = "" |
| 61 | + |
| 62 | + |
| 63 | +class AuditAction(BaseModel): |
| 64 | + mcp_tool: str = "" |
| 65 | + path: str = "" |
| 66 | + |
| 67 | + |
| 68 | +class AuditResource(BaseModel): |
| 69 | + type: str = "" |
| 70 | + name: str = "" |
| 71 | + actions: List[str] = Field(default_factory=list) |
| 72 | + |
| 73 | + |
| 74 | +class AuditEvent(BaseModel): |
| 75 | + """A single structured audit log entry. |
| 76 | +
|
| 77 | + Fields follow the schema proposed in feast-dev/feast#6452. |
| 78 | + No sensitive payloads (tokens, entity rows, feature values) are stored. |
| 79 | + When OpenTelemetry is active, ``trace_id`` and ``span_id`` are populated |
| 80 | + automatically for correlation with distributed traces. |
| 81 | + """ |
| 82 | + |
| 83 | + event_type: str |
| 84 | + timestamp: str = "" |
| 85 | + request_id: str = "" |
| 86 | + trace_id: Optional[str] = None |
| 87 | + span_id: Optional[str] = None |
| 88 | + jsonrpc_id: Optional[str] = None |
| 89 | + principal: AuditPrincipal = Field(default_factory=AuditPrincipal) |
| 90 | + source: AuditSource = Field(default_factory=AuditSource) |
| 91 | + action: AuditAction = Field(default_factory=AuditAction) |
| 92 | + resource: AuditResource = Field(default_factory=AuditResource) |
| 93 | + outcome: str = "" |
| 94 | + duration_ms: Optional[float] = None |
| 95 | + detail: str = "" |
| 96 | + |
| 97 | + def to_jsonl(self) -> str: |
| 98 | + return self.model_dump_json(exclude_none=True) |
| 99 | + |
| 100 | + |
| 101 | +# --------------------------------------------------------------------------- |
| 102 | +# Sink abstraction |
| 103 | +# --------------------------------------------------------------------------- |
| 104 | + |
| 105 | + |
| 106 | +class AuditSink(abc.ABC): |
| 107 | + """Base class for audit event sinks.""" |
| 108 | + |
| 109 | + @abc.abstractmethod |
| 110 | + def emit(self, event: AuditEvent) -> None: ... |
| 111 | + |
| 112 | + def close(self) -> None: |
| 113 | + pass |
| 114 | + |
| 115 | + |
| 116 | +class StdoutAuditSink(AuditSink): |
| 117 | + """Write JSONL events to stdout (the default).""" |
| 118 | + |
| 119 | + def emit(self, event: AuditEvent) -> None: |
| 120 | + sys.stdout.write(event.to_jsonl() + "\n") |
| 121 | + sys.stdout.flush() |
| 122 | + |
| 123 | + |
| 124 | +class FileAuditSink(AuditSink): |
| 125 | + """Append JSONL events to a local file.""" |
| 126 | + |
| 127 | + def __init__(self, file_path: str) -> None: |
| 128 | + self._file_path = file_path |
| 129 | + self._fh = open(file_path, "a", buffering=1) |
| 130 | + |
| 131 | + def emit(self, event: AuditEvent) -> None: |
| 132 | + self._fh.write(event.to_jsonl() + "\n") |
| 133 | + |
| 134 | + def close(self) -> None: |
| 135 | + self._fh.close() |
| 136 | + |
| 137 | + |
| 138 | +class LoggerAuditSink(AuditSink): |
| 139 | + """Emit audit events through Python's ``logging`` module at INFO level.""" |
| 140 | + |
| 141 | + def __init__(self, logger_name: str = "feast.audit") -> None: |
| 142 | + self._logger = logging.getLogger(logger_name) |
| 143 | + |
| 144 | + def emit(self, event: AuditEvent) -> None: |
| 145 | + self._logger.info(event.to_jsonl()) |
| 146 | + |
| 147 | + |
| 148 | +# --------------------------------------------------------------------------- |
| 149 | +# AuditLogger — the main entry point |
| 150 | +# --------------------------------------------------------------------------- |
| 151 | + |
| 152 | + |
| 153 | +class AuditLogger: |
| 154 | + """Central audit logger that routes events to the configured sink. |
| 155 | +
|
| 156 | + Instantiate once during feature-server startup and share via |
| 157 | + ``app.state.audit_logger``. |
| 158 | + """ |
| 159 | + |
| 160 | + def __init__( |
| 161 | + self, |
| 162 | + sink: AuditSink, |
| 163 | + *, |
| 164 | + log_successful_reads: bool = True, |
| 165 | + ) -> None: |
| 166 | + self._sink = sink |
| 167 | + self._log_successful_reads = log_successful_reads |
| 168 | + self._lock = threading.Lock() |
| 169 | + |
| 170 | + # -- helpers ----------------------------------------------------------- |
| 171 | + |
| 172 | + @staticmethod |
| 173 | + def new_request_id() -> str: |
| 174 | + return str(uuid.uuid4()) |
| 175 | + |
| 176 | + @staticmethod |
| 177 | + def _utcnow_iso() -> str: |
| 178 | + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" |
| 179 | + |
| 180 | + @staticmethod |
| 181 | + def _inject_otel_context(event: AuditEvent) -> None: |
| 182 | + """Populate trace_id/span_id from the active OpenTelemetry span, if any.""" |
| 183 | + if event.trace_id is not None: |
| 184 | + return |
| 185 | + try: |
| 186 | + from opentelemetry import trace as otel_trace |
| 187 | + |
| 188 | + span = otel_trace.get_current_span() |
| 189 | + ctx = span.get_span_context() |
| 190 | + if ctx and ctx.is_valid: |
| 191 | + event.trace_id = format(ctx.trace_id, "032x") |
| 192 | + event.span_id = format(ctx.span_id, "016x") |
| 193 | + except ImportError: |
| 194 | + pass |
| 195 | + except Exception: |
| 196 | + pass |
| 197 | + |
| 198 | + # MCP tool names and REST paths that correspond to read operations. |
| 199 | + _READ_TOOL_NAMES = frozenset( |
| 200 | + { |
| 201 | + "get_online_features", |
| 202 | + "retrieve_online_documents", |
| 203 | + "get_historical_features", |
| 204 | + } |
| 205 | + ) |
| 206 | + _READ_ACTIONS = frozenset({"READ_ONLINE", "READ_OFFLINE"}) |
| 207 | + |
| 208 | + @classmethod |
| 209 | + def _is_read_event(cls, event: AuditEvent) -> bool: |
| 210 | + """Return ``True`` when *event* represents a successful read. |
| 211 | +
|
| 212 | + Checks both ``resource.actions`` (populated by REST middleware) and |
| 213 | + ``action.mcp_tool`` (populated by the MCP handler wrapper) so that |
| 214 | + suppression via ``log_successful_reads=False`` works for both layers. |
| 215 | + """ |
| 216 | + if event.event_type not in {"mcp.tools.call", "http.request"}: |
| 217 | + return False |
| 218 | + resource_actions = set(event.resource.actions) |
| 219 | + if resource_actions and resource_actions.issubset(cls._READ_ACTIONS): |
| 220 | + return True |
| 221 | + if event.action.mcp_tool in cls._READ_TOOL_NAMES: |
| 222 | + return True |
| 223 | + return False |
| 224 | + |
| 225 | + # -- public API -------------------------------------------------------- |
| 226 | + |
| 227 | + def log(self, event: AuditEvent) -> None: |
| 228 | + if not event.timestamp: |
| 229 | + event.timestamp = self._utcnow_iso() |
| 230 | + if not event.request_id: |
| 231 | + event.request_id = self.new_request_id() |
| 232 | + |
| 233 | + self._inject_otel_context(event) |
| 234 | + |
| 235 | + if event.outcome == "success" and not self._log_successful_reads: |
| 236 | + if self._is_read_event(event): |
| 237 | + return |
| 238 | + |
| 239 | + try: |
| 240 | + with self._lock: |
| 241 | + self._sink.emit(event) |
| 242 | + except Exception: |
| 243 | + logger.exception("Failed to emit audit event") |
| 244 | + |
| 245 | + def log_mcp_call( |
| 246 | + self, |
| 247 | + *, |
| 248 | + request_id: str, |
| 249 | + tool_name: str, |
| 250 | + path: str = "", |
| 251 | + principal: Optional[AuditPrincipal] = None, |
| 252 | + source: Optional[AuditSource] = None, |
| 253 | + resource: Optional[AuditResource] = None, |
| 254 | + outcome: str = "success", |
| 255 | + duration_ms: Optional[float] = None, |
| 256 | + detail: str = "", |
| 257 | + ) -> None: |
| 258 | + self.log( |
| 259 | + AuditEvent( |
| 260 | + event_type="mcp.tools.call", |
| 261 | + request_id=request_id, |
| 262 | + principal=principal or AuditPrincipal(), |
| 263 | + source=source or AuditSource(transport="mcp-http"), |
| 264 | + action=AuditAction(mcp_tool=tool_name, path=path), |
| 265 | + resource=resource or AuditResource(), |
| 266 | + outcome=outcome, |
| 267 | + duration_ms=duration_ms, |
| 268 | + detail=detail, |
| 269 | + ) |
| 270 | + ) |
| 271 | + |
| 272 | + def log_http_request( |
| 273 | + self, |
| 274 | + *, |
| 275 | + request_id: str, |
| 276 | + method: str, |
| 277 | + path: str, |
| 278 | + principal: Optional[AuditPrincipal] = None, |
| 279 | + source: Optional[AuditSource] = None, |
| 280 | + resource: Optional[AuditResource] = None, |
| 281 | + outcome: str = "success", |
| 282 | + duration_ms: Optional[float] = None, |
| 283 | + status_code: int = 200, |
| 284 | + ) -> None: |
| 285 | + self.log( |
| 286 | + AuditEvent( |
| 287 | + event_type="http.request", |
| 288 | + request_id=request_id, |
| 289 | + principal=principal or AuditPrincipal(), |
| 290 | + source=source or AuditSource(transport="http"), |
| 291 | + action=AuditAction(path=path), |
| 292 | + resource=resource or AuditResource(), |
| 293 | + outcome=outcome, |
| 294 | + duration_ms=duration_ms, |
| 295 | + detail=f"{method} {path} -> {status_code}", |
| 296 | + ) |
| 297 | + ) |
| 298 | + |
| 299 | + def log_authn( |
| 300 | + self, |
| 301 | + *, |
| 302 | + request_id: str, |
| 303 | + outcome: str, |
| 304 | + principal: Optional[AuditPrincipal] = None, |
| 305 | + source: Optional[AuditSource] = None, |
| 306 | + detail: str = "", |
| 307 | + ) -> None: |
| 308 | + event_type = "authn.success" if outcome == "success" else "authn.failure" |
| 309 | + self.log( |
| 310 | + AuditEvent( |
| 311 | + event_type=event_type, |
| 312 | + request_id=request_id, |
| 313 | + principal=principal or AuditPrincipal(), |
| 314 | + source=source or AuditSource(), |
| 315 | + outcome=outcome, |
| 316 | + detail=detail, |
| 317 | + ) |
| 318 | + ) |
| 319 | + |
| 320 | + def log_authz( |
| 321 | + self, |
| 322 | + *, |
| 323 | + request_id: str, |
| 324 | + outcome: str, |
| 325 | + principal: Optional[AuditPrincipal] = None, |
| 326 | + resource: Optional[AuditResource] = None, |
| 327 | + detail: str = "", |
| 328 | + ) -> None: |
| 329 | + self.log( |
| 330 | + AuditEvent( |
| 331 | + event_type="authz.decision", |
| 332 | + request_id=request_id, |
| 333 | + principal=principal or AuditPrincipal(), |
| 334 | + resource=resource or AuditResource(), |
| 335 | + outcome=outcome, |
| 336 | + detail=detail, |
| 337 | + ) |
| 338 | + ) |
| 339 | + |
| 340 | + def close(self) -> None: |
| 341 | + self._sink.close() |
| 342 | + |
| 343 | + |
| 344 | +# --------------------------------------------------------------------------- |
| 345 | +# Factory |
| 346 | +# --------------------------------------------------------------------------- |
| 347 | + |
| 348 | +_SINK_FACTORIES = { |
| 349 | + "stdout": lambda cfg: StdoutAuditSink(), |
| 350 | + "file": lambda cfg: FileAuditSink(cfg.get("file_path", "feast_audit.log")), |
| 351 | + "logger": lambda cfg: LoggerAuditSink(cfg.get("logger_name", "feast.audit")), |
| 352 | +} |
| 353 | + |
| 354 | + |
| 355 | +def create_audit_logger_from_config( |
| 356 | + audit_cfg: Optional["AuditLoggingConfig"], |
| 357 | +) -> Optional[AuditLogger]: |
| 358 | + """Build an ``AuditLogger`` from an ``AuditLoggingConfig`` pydantic model. |
| 359 | +
|
| 360 | + Returns ``None`` when audit logging is disabled. |
| 361 | + """ |
| 362 | + if audit_cfg is None or not getattr(audit_cfg, "enabled", False): |
| 363 | + return None |
| 364 | + |
| 365 | + sink_type = getattr(audit_cfg, "sink", "stdout") |
| 366 | + raw: Dict[str, Any] = {} |
| 367 | + if hasattr(audit_cfg, "file_path"): |
| 368 | + raw["file_path"] = audit_cfg.file_path |
| 369 | + if hasattr(audit_cfg, "logger_name"): |
| 370 | + raw["logger_name"] = audit_cfg.logger_name |
| 371 | + |
| 372 | + factory = _SINK_FACTORIES.get(sink_type) |
| 373 | + if factory is None: |
| 374 | + logger.warning("Unknown audit sink %r, falling back to stdout", sink_type) |
| 375 | + factory = _SINK_FACTORIES["stdout"] |
| 376 | + |
| 377 | + sink = factory(raw) |
| 378 | + return AuditLogger( |
| 379 | + sink, |
| 380 | + log_successful_reads=getattr(audit_cfg, "log_successful_reads", True), |
| 381 | + ) |
0 commit comments