mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 10:38:32 +08:00
feat(dataset): proxy KnowledgeFS Console requests (#39158)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Stephen Zhou
GitHub
parent
c8ccfba960
commit
43d725a1ce
@ -681,6 +681,14 @@ AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30
|
||||
AGENT_BACKEND_STREAM_MAX_RECONNECTS=3
|
||||
AGENT_BACKEND_RUN_TIMEOUT_SECONDS=1200
|
||||
|
||||
# KnowledgeFS (Dataset 2.0)
|
||||
KNOWLEDGE_FS_ENABLED=false
|
||||
KNOWLEDGE_FS_BASE_URL=
|
||||
# Shared with KnowledgeFS; use at least 32 random characters.
|
||||
KNOWLEDGE_FS_JWT_SECRET=
|
||||
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=300
|
||||
KNOWLEDGE_FS_TIMEOUT_SECONDS=10
|
||||
|
||||
# Marketplace configuration
|
||||
MARKETPLACE_ENABLED=true
|
||||
MARKETPLACE_API_URL=https://marketplace.dify.ai
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from configs.extra.agent_backend_config import AgentBackendConfig
|
||||
from configs.extra.archive_config import ArchiveStorageConfig
|
||||
from configs.extra.knowledge_fs_config import KnowledgeFSConfig
|
||||
from configs.extra.notion_config import NotionConfig
|
||||
from configs.extra.sentry_config import SentryConfig
|
||||
|
||||
@ -8,6 +9,7 @@ class ExtraServiceConfig(
|
||||
# place the configs in alphabet order
|
||||
AgentBackendConfig,
|
||||
ArchiveStorageConfig,
|
||||
KnowledgeFSConfig,
|
||||
NotionConfig,
|
||||
SentryConfig,
|
||||
):
|
||||
|
||||
64
api/configs/extra/knowledge_fs_config.py
Normal file
64
api/configs/extra/knowledge_fs_config.py
Normal file
@ -0,0 +1,64 @@
|
||||
"""Configuration for the optional KnowledgeFS Console bridge."""
|
||||
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from pydantic import Field, PositiveFloat, SecretStr, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class KnowledgeFSConfig(BaseSettings):
|
||||
"""Server-only settings for the KnowledgeFS production connection."""
|
||||
|
||||
KNOWLEDGE_FS_ENABLED: bool = Field(
|
||||
default=False,
|
||||
description="Enable the private KnowledgeFS Console bridge.",
|
||||
)
|
||||
KNOWLEDGE_FS_BASE_URL: str | None = Field(default=None, description="KnowledgeFS gateway base URL.")
|
||||
KNOWLEDGE_FS_JWT_SECRET: SecretStr | None = Field(
|
||||
default=None,
|
||||
min_length=32,
|
||||
description="Shared secret used to sign short-lived KnowledgeFS service JWTs.",
|
||||
)
|
||||
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS: PositiveFloat = Field(default=300.0, le=3600.0, allow_inf_nan=False)
|
||||
KNOWLEDGE_FS_TIMEOUT_SECONDS: PositiveFloat = Field(default=10.0, le=60.0, allow_inf_nan=False)
|
||||
|
||||
@field_validator(
|
||||
"KNOWLEDGE_FS_BASE_URL",
|
||||
"KNOWLEDGE_FS_JWT_SECRET",
|
||||
mode="before",
|
||||
)
|
||||
@classmethod
|
||||
def normalize_optional_string(cls, value: object) -> object:
|
||||
if isinstance(value, SecretStr):
|
||||
normalized = value.get_secret_value().strip()
|
||||
return SecretStr(normalized) if normalized else None
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip()
|
||||
return normalized or None
|
||||
return value
|
||||
|
||||
@field_validator("KNOWLEDGE_FS_BASE_URL")
|
||||
@classmethod
|
||||
def validate_base_url(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
parsed = urlsplit(value)
|
||||
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must be an absolute HTTP(S) URL")
|
||||
try:
|
||||
_ = parsed.port
|
||||
except ValueError as exc:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must include a valid port") from exc
|
||||
if parsed.username or parsed.password or parsed.query or parsed.fragment:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must not include credentials, query, or fragment")
|
||||
return value.rstrip("/")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_enabled_connection(self) -> "KnowledgeFSConfig":
|
||||
if not self.KNOWLEDGE_FS_ENABLED:
|
||||
return self
|
||||
if bool(self.KNOWLEDGE_FS_BASE_URL) != bool(self.KNOWLEDGE_FS_JWT_SECRET):
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL and KNOWLEDGE_FS_JWT_SECRET must be configured together")
|
||||
if not self.KNOWLEDGE_FS_BASE_URL:
|
||||
raise ValueError("KnowledgeFS connection settings are required when the integration is enabled")
|
||||
return self
|
||||
@ -38,6 +38,7 @@ from . import (
|
||||
feature,
|
||||
human_input_form,
|
||||
init_validate,
|
||||
knowledge_fs_proxy,
|
||||
notification,
|
||||
ping,
|
||||
setup,
|
||||
@ -195,6 +196,7 @@ __all__ = [
|
||||
"human_input_form",
|
||||
"init_validate",
|
||||
"installed_app",
|
||||
"knowledge_fs_proxy",
|
||||
"load_balancing_config",
|
||||
"login",
|
||||
"mcp_server",
|
||||
|
||||
332
api/controllers/console/knowledge_fs_proxy.py
Normal file
332
api/controllers/console/knowledge_fs_proxy.py
Normal file
@ -0,0 +1,332 @@
|
||||
"""Authenticated transport adapter for the Console-to-KnowledgeFS proxy.
|
||||
|
||||
These raw Blueprint routes deliberately stay outside Dify's OpenAPI surface:
|
||||
KnowledgeFS owns the wire contract consumed by the frontend. The catch-all path
|
||||
avoids resource-specific Dify controllers, while the forwarding module consumes
|
||||
only the operations explicitly enabled by Dify's product registry. The registry
|
||||
can be validated explicitly against the pinned KnowledgeFS contract during development.
|
||||
Console auth and contract-specific dataset RBAC run before forwarding. Request
|
||||
bodies are capped at 64 MiB, JSON and binary responses have separate bounds,
|
||||
SSE responses remain streaming with a bounded idle read timeout, and only safe
|
||||
response headers are exposed. Upstream 401 responses become 502 so they cannot
|
||||
trigger Dify browser-session recovery; resource-level 403 responses remain 403.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Iterator
|
||||
from functools import wraps
|
||||
from http import HTTPStatus
|
||||
from typing import NoReturn, cast
|
||||
|
||||
import httpx
|
||||
from flask import Response, request, stream_with_context
|
||||
from flask.typing import ResponseReturnValue
|
||||
from werkzeug.exceptions import (
|
||||
BadGateway,
|
||||
Forbidden,
|
||||
GatewayTimeout,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.console import api, bp
|
||||
from controllers.console.wraps import (
|
||||
account_initialization_required,
|
||||
cloud_edition_billing_rate_limit_check,
|
||||
setup_required,
|
||||
)
|
||||
from core.helper import ssrf_proxy
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from services.knowledge_fs_proxy import (
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSMethod,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSTimeoutError,
|
||||
KnowledgeFSTransportError,
|
||||
KnowledgeFSUpstreamResponse,
|
||||
authorize_knowledge_fs_request,
|
||||
get_knowledge_fs_operation,
|
||||
proxy_knowledge_fs_request,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MAX_PROXY_BODY_BYTES = 64 * 1024 * 1024
|
||||
_RESPONSE_HEADER_ALLOWLIST = (
|
||||
"Cache-Control",
|
||||
"Content-Disposition",
|
||||
"Content-Type",
|
||||
"Retry-After",
|
||||
"X-Trace-Id",
|
||||
)
|
||||
_RESPONSE_HEADER_DENYLIST = frozenset(
|
||||
{
|
||||
"authorization",
|
||||
"connection",
|
||||
"cookie",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"set-cookie",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _console_api_errors[**P](
|
||||
view: Callable[P, ResponseReturnValue],
|
||||
) -> Callable[P, ResponseReturnValue]:
|
||||
"""Route raw Blueprint exceptions through the Console API JSON handlers."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
|
||||
try:
|
||||
return view(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
return api.handle_error(exc)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _knowledge_fs_enabled[**P](
|
||||
view: Callable[P, ResponseReturnValue],
|
||||
) -> Callable[P, ResponseReturnValue]:
|
||||
"""Hide the complete KnowledgeFS route surface while the bridge is disabled."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
|
||||
if not dify_config.KNOWLEDGE_FS_ENABLED:
|
||||
raise NotFound()
|
||||
return view(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _translate_proxy_error(exc: Exception, *, tenant_id: str) -> NoReturn:
|
||||
"""Map forwarding failures to the stable Console HTTP error surface."""
|
||||
if isinstance(exc, KnowledgeFSRouteNotAllowedError):
|
||||
raise NotFound() from exc
|
||||
if isinstance(exc, KnowledgeFSAccessDeniedError):
|
||||
raise Forbidden() from exc
|
||||
if isinstance(exc, KnowledgeFSConfigurationError):
|
||||
logger.error("KnowledgeFS request was blocked by invalid configuration for tenant_id=%s", tenant_id)
|
||||
raise ServiceUnavailable("KnowledgeFS integration is misconfigured") from exc
|
||||
if isinstance(exc, KnowledgeFSTimeoutError):
|
||||
raise GatewayTimeout("KnowledgeFS request timed out") from exc
|
||||
if isinstance(exc, KnowledgeFSTransportError):
|
||||
logger.warning("KnowledgeFS transport request failed for tenant_id=%s", tenant_id)
|
||||
raise BadGateway("KnowledgeFS is unavailable") from exc
|
||||
raise exc
|
||||
|
||||
|
||||
def _knowledge_fs_operation_access_required(
|
||||
view: Callable[[KnowledgeFSMethod, str], ResponseReturnValue],
|
||||
) -> Callable[[KnowledgeFSMethod, str], ResponseReturnValue]:
|
||||
"""Authorize one declared operation before billing and request-body work."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(method: KnowledgeFSMethod, upstream_path: str) -> ResponseReturnValue:
|
||||
try:
|
||||
operation = get_knowledge_fs_operation(method, upstream_path)
|
||||
except KnowledgeFSRouteNotAllowedError as exc:
|
||||
raise NotFound() from exc
|
||||
|
||||
current_user, tenant_id = current_account_with_tenant()
|
||||
try:
|
||||
authorize_knowledge_fs_request(
|
||||
account=current_user,
|
||||
tenant_id=tenant_id,
|
||||
operation=operation,
|
||||
)
|
||||
except KnowledgeFSAccessDeniedError as exc:
|
||||
_translate_proxy_error(exc, tenant_id=tenant_id)
|
||||
return view(method, upstream_path)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _request_body() -> bytes:
|
||||
"""Read the raw body up to the proxy limit or raise RequestEntityTooLarge."""
|
||||
body = request.stream.read(_MAX_PROXY_BODY_BYTES + 1)
|
||||
if len(body) > _MAX_PROXY_BODY_BYTES:
|
||||
raise RequestEntityTooLarge("KnowledgeFS proxy request body is too large")
|
||||
return body
|
||||
|
||||
|
||||
def _stream_response_body(
|
||||
upstream: httpx.Response,
|
||||
*,
|
||||
tenant_id: str,
|
||||
max_response_bytes: int,
|
||||
) -> Iterator[bytes]:
|
||||
"""Yield one bounded SSE response and always release its pooled connection."""
|
||||
total_bytes = 0
|
||||
try:
|
||||
for chunk in upstream.iter_bytes():
|
||||
total_bytes += len(chunk)
|
||||
if total_bytes > max_response_bytes:
|
||||
logger.warning("KnowledgeFS stream exceeded the proxy limit for tenant_id=%s", tenant_id)
|
||||
raise ssrf_proxy.ResponseTooLargeError(f"response exceeded {max_response_bytes} bytes")
|
||||
yield chunk
|
||||
finally:
|
||||
upstream.close()
|
||||
|
||||
|
||||
def _proxy_response(
|
||||
upstream_result: KnowledgeFSUpstreamResponse,
|
||||
*,
|
||||
tenant_id: str,
|
||||
contract_response_headers: tuple[str, ...],
|
||||
max_response_bytes: int,
|
||||
) -> Response:
|
||||
"""Expose raw content, status, and allowlisted headers from KnowledgeFS.
|
||||
|
||||
Raises:
|
||||
BadGateway: KnowledgeFS rejects the configured server credential.
|
||||
Forbidden: KnowledgeFS denies the account access to the requested resource.
|
||||
"""
|
||||
upstream = upstream_result.response
|
||||
if upstream.status_code == HTTPStatus.UNAUTHORIZED:
|
||||
upstream.close()
|
||||
logger.error(
|
||||
"KnowledgeFS rejected the Dify server credential with HTTP %s for tenant_id=%s",
|
||||
upstream.status_code,
|
||||
tenant_id,
|
||||
)
|
||||
raise BadGateway("KnowledgeFS authentication failed")
|
||||
if upstream.status_code == HTTPStatus.FORBIDDEN:
|
||||
upstream.close()
|
||||
raise Forbidden()
|
||||
|
||||
allowed_header_names = dict.fromkeys(
|
||||
name.lower() for name in (*_RESPONSE_HEADER_ALLOWLIST, *contract_response_headers)
|
||||
)
|
||||
headers = {
|
||||
name: value
|
||||
for name in allowed_header_names
|
||||
if name not in _RESPONSE_HEADER_DENYLIST
|
||||
if (value := upstream.headers.get(name)) is not None
|
||||
}
|
||||
if upstream_result.response_kind == "stream":
|
||||
response = Response(
|
||||
stream_with_context( # pyrefly: ignore[no-matching-overload]
|
||||
_stream_response_body(
|
||||
upstream,
|
||||
tenant_id=tenant_id,
|
||||
max_response_bytes=max_response_bytes,
|
||||
)
|
||||
),
|
||||
status=upstream.status_code,
|
||||
headers=headers,
|
||||
)
|
||||
response.call_on_close(upstream.close)
|
||||
return response
|
||||
|
||||
try:
|
||||
content = upstream.content
|
||||
finally:
|
||||
upstream.close()
|
||||
return Response(content, status=upstream.status_code, headers=headers)
|
||||
|
||||
|
||||
def _proxy_request(method: KnowledgeFSMethod, upstream_path: str) -> Response:
|
||||
"""Forward the current raw request and return its filtered upstream response.
|
||||
|
||||
The call performs one outbound KnowledgeFS request. Integration failures are
|
||||
converted to Console HTTP exceptions for the outer JSON error adapter.
|
||||
"""
|
||||
if not dify_config.KNOWLEDGE_FS_ENABLED:
|
||||
raise NotFound()
|
||||
current_user, tenant_id = current_account_with_tenant()
|
||||
try:
|
||||
proxy_result = proxy_knowledge_fs_request(
|
||||
account=current_user,
|
||||
method=method,
|
||||
path=upstream_path,
|
||||
tenant_id=tenant_id,
|
||||
accept=request.headers.get("Accept"),
|
||||
content_type=request.content_type,
|
||||
query=request.query_string or None,
|
||||
body=_request_body() if method != "GET" else None,
|
||||
request_headers=request.headers,
|
||||
)
|
||||
except (
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSTimeoutError,
|
||||
KnowledgeFSTransportError,
|
||||
) as exc:
|
||||
_translate_proxy_error(exc, tenant_id=tenant_id)
|
||||
return _proxy_response(
|
||||
proxy_result,
|
||||
tenant_id=tenant_id,
|
||||
contract_response_headers=proxy_result.operation.response_headers,
|
||||
max_response_bytes=proxy_result.operation.max_response_bytes,
|
||||
)
|
||||
|
||||
|
||||
@_knowledge_fs_enabled
|
||||
@_knowledge_fs_operation_access_required
|
||||
@cloud_edition_billing_rate_limit_check("knowledge")
|
||||
def _proxy_knowledge_fs_non_get(
|
||||
method: KnowledgeFSMethod,
|
||||
upstream_path: str,
|
||||
) -> ResponseReturnValue:
|
||||
"""Apply knowledge billing checks to one allowlisted non-GET operation."""
|
||||
return _proxy_request(method, upstream_path)
|
||||
|
||||
|
||||
@bp.route(
|
||||
"/knowledge-fs/<path:upstream_path>",
|
||||
methods=["GET", "OPTIONS"],
|
||||
provide_automatic_options=False,
|
||||
)
|
||||
@_console_api_errors
|
||||
@_knowledge_fs_enabled
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def proxy_knowledge_fs_get(upstream_path: str) -> ResponseReturnValue:
|
||||
"""Forward one authenticated, dataset-readable GET request.
|
||||
|
||||
Args:
|
||||
upstream_path: Relative KFS path captured after the Console proxy prefix.
|
||||
|
||||
Returns:
|
||||
The filtered raw KnowledgeFS response or a Console JSON error response.
|
||||
"""
|
||||
if request.method != "GET":
|
||||
raise NotFound()
|
||||
return _proxy_request("GET", upstream_path)
|
||||
|
||||
|
||||
@bp.route(
|
||||
"/knowledge-fs/<path:upstream_path>",
|
||||
methods=["DELETE", "PATCH", "POST", "PUT"],
|
||||
provide_automatic_options=False,
|
||||
)
|
||||
@_console_api_errors
|
||||
@_knowledge_fs_enabled
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def proxy_knowledge_fs_write(upstream_path: str) -> ResponseReturnValue:
|
||||
"""Forward one authenticated non-GET request under its contract access policy.
|
||||
|
||||
Args:
|
||||
upstream_path: Relative KFS path captured after the Console proxy prefix.
|
||||
|
||||
Returns:
|
||||
The filtered raw KnowledgeFS response or a Console JSON error response.
|
||||
"""
|
||||
method = cast(KnowledgeFSMethod, request.method)
|
||||
return _proxy_knowledge_fs_non_get(method, upstream_path)
|
||||
@ -47,6 +47,24 @@ class MaxRetriesExceededError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class ResponseLimitError(ValueError):
|
||||
"""Base error for responses that cannot be safely bounded."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ResponseTooLargeError(ResponseLimitError):
|
||||
"""Raised when an identity response exceeds the configured byte limit."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class UnsupportedResponseEncodingError(ResponseLimitError):
|
||||
"""Raised when response encoding prevents safe decoded-size enforcement."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
request_error = httpx.RequestError
|
||||
max_retries_exceeded_error = MaxRetriesExceededError
|
||||
|
||||
@ -142,7 +160,31 @@ def _inject_trace_headers(headers: Headers | None) -> Headers:
|
||||
return headers
|
||||
|
||||
|
||||
def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETRIES, **kwargs: Any) -> httpx.Response:
|
||||
def make_request(
|
||||
method: str,
|
||||
url: str,
|
||||
max_retries: int = SSRF_DEFAULT_MAX_RETRIES,
|
||||
stream_response: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> httpx.Response:
|
||||
"""Send one SSRF-protected request with optional streaming.
|
||||
|
||||
Args:
|
||||
method: HTTP method sent through the configured SSRF client.
|
||||
url: Absolute request URL.
|
||||
max_retries: Number of retry attempts after the initial request.
|
||||
stream_response: Return an open streaming response that the caller must close.
|
||||
**kwargs: Additional keyword arguments forwarded to ``httpx.Client``.
|
||||
|
||||
Returns:
|
||||
A buffered response, or an open response when ``stream_response`` is true.
|
||||
|
||||
Raises:
|
||||
ToolSSRFError: The configured SSRF proxy rejects the destination.
|
||||
MaxRetriesExceededError: All configured request attempts fail.
|
||||
httpx.RequestError: A request fails while retries are disabled.
|
||||
ValueError: The SSL verification option or request headers are invalid.
|
||||
"""
|
||||
# Convert requests-style allow_redirects to httpx-style follow_redirects
|
||||
if "allow_redirects" in kwargs:
|
||||
allow_redirects = kwargs.pop("allow_redirects")
|
||||
@ -175,6 +217,11 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
# When using a forward proxy, httpx may override the Host header based on the URL.
|
||||
# We extract and preserve any explicitly set Host header to support virtual hosting.
|
||||
user_provided_host = _get_user_provided_host_header(headers)
|
||||
send_kwargs: dict[str, Any] = {}
|
||||
if "auth" in kwargs:
|
||||
send_kwargs["auth"] = kwargs.pop("auth")
|
||||
if "follow_redirects" in kwargs:
|
||||
send_kwargs["follow_redirects"] = kwargs.pop("follow_redirects")
|
||||
|
||||
retries = 0
|
||||
while retries <= max_retries:
|
||||
@ -185,7 +232,11 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
if user_provided_host is not None:
|
||||
headers["host"] = user_provided_host
|
||||
kwargs["headers"] = headers
|
||||
response = client.request(method=method, url=url, **kwargs)
|
||||
request = client.build_request(method=method, url=url, **kwargs)
|
||||
if stream_response:
|
||||
response = client.send(request, stream=True, **send_kwargs)
|
||||
else:
|
||||
response = client.send(request, **send_kwargs)
|
||||
|
||||
# Check for SSRF protection by Squid proxy
|
||||
if response.status_code in (401, 403):
|
||||
@ -195,6 +246,7 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
|
||||
# Squid typically identifies itself in Server or Via headers
|
||||
if "squid" in server_header or "squid" in via_header:
|
||||
response.close()
|
||||
raise ToolSSRFError(
|
||||
f"Access to '{url}' was blocked by SSRF protection. "
|
||||
f"The URL may point to a private or local network address. "
|
||||
@ -208,6 +260,7 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
response.status_code,
|
||||
url,
|
||||
)
|
||||
response.close()
|
||||
|
||||
except httpx.RequestError as e:
|
||||
logger.warning("Request to URL %s failed on attempt %s: %s", url, retries + 1, e)
|
||||
@ -220,6 +273,42 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
raise MaxRetriesExceededError(f"Reached maximum retries ({max_retries}) for URL {url}")
|
||||
|
||||
|
||||
def buffer_response(response: httpx.Response, *, max_response_bytes: int) -> httpx.Response:
|
||||
"""Consume one open identity response under a decoded byte limit and close its stream."""
|
||||
if max_response_bytes <= 0:
|
||||
raise ValueError("max_response_bytes must be positive")
|
||||
|
||||
try:
|
||||
content_encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||||
if content_encoding not in {"", "identity"}:
|
||||
raise UnsupportedResponseEncodingError(f"content encoding {content_encoding} cannot be safely bounded")
|
||||
content = bytearray()
|
||||
for chunk in response.iter_bytes():
|
||||
if len(content) + len(chunk) > max_response_bytes:
|
||||
raise ResponseTooLargeError(f"response exceeded {max_response_bytes} bytes")
|
||||
content.extend(chunk)
|
||||
decoded_headers = {
|
||||
name: value
|
||||
for name, value in response.headers.items()
|
||||
if name.lower() not in {"content-encoding", "content-length", "transfer-encoding"}
|
||||
}
|
||||
try:
|
||||
request = response.request
|
||||
except RuntimeError:
|
||||
request = None
|
||||
return httpx.Response(
|
||||
response.status_code,
|
||||
headers=decoded_headers,
|
||||
content=bytes(content),
|
||||
request=request,
|
||||
extensions=response.extensions,
|
||||
history=response.history,
|
||||
default_encoding=response.default_encoding,
|
||||
)
|
||||
finally:
|
||||
response.close()
|
||||
|
||||
|
||||
def get(url: str, max_retries: int = SSRF_DEFAULT_MAX_RETRIES, **kwargs: Any) -> httpx.Response:
|
||||
return make_request("GET", url, max_retries=max_retries, **kwargs)
|
||||
|
||||
|
||||
255
api/dev/generate_knowledge_fs_contract.py
Normal file
255
api/dev/generate_knowledge_fs_contract.py
Normal file
@ -0,0 +1,255 @@
|
||||
"""Validate Dify Console KnowledgeFS declarations against a pinned OpenAPI document.
|
||||
|
||||
The OpenAPI document is exported only during explicit development validation. Runtime declarations live with Dify
|
||||
product policy; this module validates their transport metadata without generating a complete operation catalog.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, TypedDict
|
||||
|
||||
API_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(API_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(API_ROOT))
|
||||
|
||||
WORKSPACE_ROOT = API_ROOT.parent
|
||||
LOCK_PATH = API_ROOT / "knowledge-fs-contract.lock.json"
|
||||
DEFAULT_REPOSITORY = WORKSPACE_ROOT.parent / "knowledge-fs"
|
||||
OPENAPI_METHODS = ("delete", "get", "head", "options", "patch", "post", "put", "trace")
|
||||
PROXY_METHODS = frozenset({"delete", "get", "patch", "post", "put"})
|
||||
|
||||
|
||||
class ContractDeclaration(TypedDict):
|
||||
"""KnowledgeFS transport contract declared by one Dify Console registry entry."""
|
||||
|
||||
operation_id: str
|
||||
method: str
|
||||
path: str
|
||||
required_scope: str | None
|
||||
response_kind: str
|
||||
max_response_bytes: int
|
||||
request_headers: tuple[str, ...]
|
||||
response_headers: tuple[str, ...]
|
||||
response_media_types: tuple[str, ...]
|
||||
|
||||
|
||||
type DeclarationField = Literal[
|
||||
"method",
|
||||
"path",
|
||||
"required_scope",
|
||||
"response_kind",
|
||||
"max_response_bytes",
|
||||
"request_headers",
|
||||
"response_headers",
|
||||
"response_media_types",
|
||||
]
|
||||
|
||||
DECLARATION_FIELDS: tuple[DeclarationField, ...] = (
|
||||
"method",
|
||||
"path",
|
||||
"required_scope",
|
||||
"response_kind",
|
||||
"max_response_bytes",
|
||||
"request_headers",
|
||||
"response_headers",
|
||||
"response_media_types",
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Update or verify the pin and validate Console declarations against its OpenAPI document."""
|
||||
parser = argparse.ArgumentParser()
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument("--check", action="store_true")
|
||||
mode.add_argument("--update-lock", action="store_true")
|
||||
parser.add_argument("--repository", type=Path, default=DEFAULT_REPOSITORY)
|
||||
args = parser.parse_args()
|
||||
|
||||
repository = args.repository.resolve()
|
||||
lock = json.loads(LOCK_PATH.read_text())
|
||||
tracked_changes = run("git", "status", "--porcelain", "--untracked-files=no", cwd=repository).strip()
|
||||
if tracked_changes:
|
||||
raise RuntimeError("KnowledgeFS checkout must not contain tracked changes during contract export")
|
||||
|
||||
commit = run("git", "rev-parse", "HEAD", cwd=repository).strip()
|
||||
if not args.update_lock and commit != lock["commit"]:
|
||||
raise RuntimeError(
|
||||
f"KnowledgeFS checkout mismatch: expected {lock['commit']}, received {commit}. "
|
||||
"Use the pinned commit or pass --update-lock intentionally."
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="dify-knowledge-fs-contract-") as directory:
|
||||
openapi_path = Path(directory) / "knowledge-fs.openapi.json"
|
||||
subprocess.run(
|
||||
["pnpm", "openapi:export", "--", "--output", str(openapi_path)],
|
||||
cwd=repository,
|
||||
check=True,
|
||||
)
|
||||
openapi_content = openapi_path.read_bytes()
|
||||
|
||||
openapi_sha256 = sha256(openapi_content)
|
||||
if not args.update_lock and openapi_sha256 != lock["openapiSha256"]:
|
||||
raise RuntimeError(
|
||||
f"KnowledgeFS OpenAPI hash mismatch: expected {lock['openapiSha256']}, received {openapi_sha256}"
|
||||
)
|
||||
|
||||
document: dict[str, Any] = json.loads(openapi_content)
|
||||
validate_declarations(document, console_contract_declarations())
|
||||
|
||||
if args.update_lock:
|
||||
LOCK_PATH.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"commit": commit,
|
||||
"openapiSha256": openapi_sha256,
|
||||
"repository": lock["repository"],
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
|
||||
def validate_declarations(document: dict[str, Any], declarations: tuple[ContractDeclaration, ...]) -> None:
|
||||
"""Validate Dify Console declarations against matching pinned OpenAPI operations."""
|
||||
operations_by_id: dict[str, list[tuple[str, str, dict[str, Any], dict[str, Any]]]] = {}
|
||||
for path, path_item in document.get("paths", {}).items():
|
||||
for method in OPENAPI_METHODS:
|
||||
operation = path_item.get(method)
|
||||
if operation is None:
|
||||
continue
|
||||
operation_id = operation.get("operationId")
|
||||
if isinstance(operation_id, str) and operation_id:
|
||||
operations_by_id.setdefault(operation_id, []).append((method, path, path_item, operation))
|
||||
|
||||
declared_ids: set[str] = set()
|
||||
for declaration in declarations:
|
||||
operation_id = declaration["operation_id"]
|
||||
if operation_id in declared_ids:
|
||||
raise ValueError(f"Dify Console registry has duplicate operationId: {operation_id}")
|
||||
declared_ids.add(operation_id)
|
||||
|
||||
matches = operations_by_id.get(operation_id, [])
|
||||
if not matches:
|
||||
raise ValueError(f"KnowledgeFS OpenAPI has no operationId: {operation_id}")
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"KnowledgeFS OpenAPI has duplicate operationId: {operation_id}")
|
||||
|
||||
method, path, path_item, operation = matches[0]
|
||||
if not path.startswith("/"):
|
||||
raise ValueError(f"KnowledgeFS OpenAPI path must be absolute: {path}")
|
||||
if method not in PROXY_METHODS:
|
||||
raise ValueError(f"KnowledgeFS proxy does not support {method.upper()} {path}")
|
||||
expected: ContractDeclaration = {
|
||||
"operation_id": operation_id,
|
||||
"method": method.upper(),
|
||||
"path": path[1:],
|
||||
"required_scope": required_scope(operation),
|
||||
"response_kind": response_kind(operation),
|
||||
"max_response_bytes": required_max_response_bytes(operation),
|
||||
"request_headers": request_header_names(path_item, operation),
|
||||
"response_headers": response_header_names(operation),
|
||||
"response_media_types": response_media_types(operation),
|
||||
}
|
||||
for field in DECLARATION_FIELDS:
|
||||
expected_value = expected[field]
|
||||
received_value = declaration[field]
|
||||
if received_value != expected_value:
|
||||
raise ValueError(
|
||||
f"KnowledgeFS operation {operation_id} field {field} drifted: "
|
||||
f"expected {expected_value!r}, received {received_value!r}"
|
||||
)
|
||||
|
||||
|
||||
def console_contract_declarations() -> tuple[ContractDeclaration, ...]:
|
||||
"""Return transport declarations from the runtime Console operation registry."""
|
||||
from services.knowledge_fs_proxy import KNOWLEDGE_FS_CONSOLE_OPERATIONS
|
||||
|
||||
return tuple(
|
||||
{
|
||||
"operation_id": operation.operation_id,
|
||||
"method": operation.method,
|
||||
"path": operation.path,
|
||||
"required_scope": operation.required_scope,
|
||||
"response_kind": operation.response_kind,
|
||||
"max_response_bytes": operation.max_response_bytes,
|
||||
"request_headers": operation.request_headers,
|
||||
"response_headers": operation.response_headers,
|
||||
"response_media_types": operation.response_media_types,
|
||||
}
|
||||
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS
|
||||
)
|
||||
|
||||
|
||||
def response_kind(operation: dict[str, Any]) -> str:
|
||||
media_types = response_media_types(operation)
|
||||
if "text/event-stream" in media_types:
|
||||
return "stream"
|
||||
if "application/octet-stream" in media_types:
|
||||
return "binary"
|
||||
return "buffered"
|
||||
|
||||
|
||||
def response_media_types(operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
media_types: set[str] = set()
|
||||
for status, response in operation.get("responses", {}).items():
|
||||
if status == "2XX" or (len(status) == 3 and status.startswith("2") and status.isdigit()):
|
||||
media_types.update(response.get("content", {}))
|
||||
return tuple(sorted(media_types))
|
||||
|
||||
|
||||
def required_scope(operation: dict[str, Any]) -> str | None:
|
||||
scope = operation.get("x-knowledge-fs-required-scope")
|
||||
if scope in ("knowledge-spaces:read", "knowledge-spaces:write"):
|
||||
return scope
|
||||
if operation.get("security") == []:
|
||||
return None
|
||||
raise ValueError(f"KnowledgeFS operation has no supported required scope: {scope}")
|
||||
|
||||
|
||||
def required_max_response_bytes(operation: dict[str, Any]) -> int:
|
||||
value = operation.get("x-knowledge-fs-max-response-bytes")
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise ValueError(f"KnowledgeFS operation has no valid response byte limit: {value}")
|
||||
return value
|
||||
|
||||
|
||||
def request_header_names(path_item: dict[str, Any], operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
names: set[str] = set()
|
||||
for parameter in [*path_item.get("parameters", []), *operation.get("parameters", [])]:
|
||||
if "$ref" in parameter:
|
||||
raise ValueError(f"KnowledgeFS request header references are not supported: {parameter['$ref']}")
|
||||
if parameter.get("in") == "header":
|
||||
names.add(parameter["name"].lower())
|
||||
return tuple(sorted(names))
|
||||
|
||||
|
||||
def response_header_names(operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
{
|
||||
name.lower()
|
||||
for response in operation.get("responses", {}).values()
|
||||
for name in response.get("headers", {})
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def sha256(content: bytes) -> str:
|
||||
return hashlib.sha256(content).hexdigest()
|
||||
|
||||
|
||||
def run(*command: str, cwd: Path) -> str:
|
||||
return subprocess.run(command, cwd=cwd, check=True, capture_output=True, text=True).stdout
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
5
api/knowledge-fs-contract.lock.json
Normal file
5
api/knowledge-fs-contract.lock.json
Normal file
@ -0,0 +1,5 @@
|
||||
{
|
||||
"commit": "4310e2d582d25e7de58183f27720afab01e123cf",
|
||||
"openapiSha256": "5827ca930ce38462bfd1b2bef387efbf37eb7ffcaedde4558af2fbaeccbfbc4b",
|
||||
"repository": "https://github.com/langgenius/knowledge-fs"
|
||||
}
|
||||
358
api/services/knowledge_fs_proxy.py
Normal file
358
api/services/knowledge_fs_proxy.py
Normal file
@ -0,0 +1,358 @@
|
||||
"""Transport-only forwarding for the explicitly enabled KnowledgeFS Console operations.
|
||||
|
||||
KnowledgeFS owns the request and response contract. This module binds short-lived
|
||||
account and workspace identities, enforces Dify's coarse workspace policy, and
|
||||
normalizes transport failures. Dify deliberately maintains a small product-facing
|
||||
operation registry instead of exposing the full upstream OpenAPI surface. The
|
||||
dedicated request path uses Dify's shared SSRF policy, never follows redirects,
|
||||
bounds buffered responses, and rejects compressed responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from http import HTTPStatus
|
||||
from typing import Final, Literal, NamedTuple, Protocol
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
|
||||
from configs import dify_config
|
||||
from core.helper import ssrf_proxy
|
||||
from core.rbac import RBACPermission, RBACResourceScope
|
||||
from core.tools.errors import ToolSSRFError
|
||||
from models import Account
|
||||
from services.enterprise.rbac_service import RBACService
|
||||
|
||||
type KnowledgeFSMethod = Literal["DELETE", "GET", "PATCH", "POST", "PUT"]
|
||||
type KnowledgeFSResponseKind = Literal["binary", "buffered", "stream"]
|
||||
type KnowledgeFSRequiredScope = Literal["knowledge-spaces:read", "knowledge-spaces:write"]
|
||||
|
||||
_JWT_AUDIENCE = "knowledge-fs"
|
||||
_JWT_ISSUER = "dify"
|
||||
_JWT_TTL_SECONDS = 60
|
||||
_MAX_BUFFERED_RESPONSE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class KnowledgeFSOperation(NamedTuple):
|
||||
operation_id: str
|
||||
method: KnowledgeFSMethod
|
||||
path: str
|
||||
response_kind: KnowledgeFSResponseKind
|
||||
required_scope: KnowledgeFSRequiredScope
|
||||
rbac_permission: RBACPermission
|
||||
requires_dataset_editor: bool
|
||||
max_response_bytes: int
|
||||
request_headers: tuple[str, ...]
|
||||
response_headers: tuple[str, ...]
|
||||
response_media_types: tuple[str, ...]
|
||||
|
||||
|
||||
KNOWLEDGE_FS_CONSOLE_OPERATIONS: Final[tuple[KnowledgeFSOperation, ...]] = (
|
||||
KnowledgeFSOperation(
|
||||
operation_id="listKnowledgeSpaces",
|
||||
method="GET",
|
||||
path="knowledge-spaces",
|
||||
response_kind="buffered",
|
||||
required_scope="knowledge-spaces:read",
|
||||
rbac_permission=RBACPermission.DATASET_READONLY,
|
||||
requires_dataset_editor=False,
|
||||
max_response_bytes=1_048_576,
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
response_media_types=("application/json",),
|
||||
),
|
||||
KnowledgeFSOperation(
|
||||
operation_id="createKnowledgeSpace",
|
||||
method="POST",
|
||||
path="knowledge-spaces",
|
||||
response_kind="buffered",
|
||||
required_scope="knowledge-spaces:write",
|
||||
rbac_permission=RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
requires_dataset_editor=True,
|
||||
max_response_bytes=1_048_576,
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
response_media_types=("application/json",),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class KnowledgeFSUpstreamResponse(NamedTuple):
|
||||
response: httpx.Response
|
||||
response_kind: KnowledgeFSResponseKind
|
||||
operation: KnowledgeFSOperation
|
||||
|
||||
|
||||
class _RequestHeaders(Protocol):
|
||||
def items(self) -> Iterable[tuple[str, str]]: ...
|
||||
|
||||
|
||||
class KnowledgeFSConfigurationError(RuntimeError):
|
||||
"""KnowledgeFS is incompletely configured or blocked by outbound policy."""
|
||||
|
||||
|
||||
class KnowledgeFSTimeoutError(RuntimeError):
|
||||
"""KnowledgeFS exceeded the configured request timeout."""
|
||||
|
||||
|
||||
class KnowledgeFSTransportError(RuntimeError):
|
||||
"""KnowledgeFS could not be reached or returned a response outside safety bounds."""
|
||||
|
||||
|
||||
class KnowledgeFSRouteNotAllowedError(RuntimeError):
|
||||
"""The requested path is outside the Console-visible KnowledgeFS surface."""
|
||||
|
||||
|
||||
class KnowledgeFSAccessDeniedError(RuntimeError):
|
||||
"""The Dify account lacks the workspace permission required by the operation."""
|
||||
|
||||
|
||||
def authorize_knowledge_fs_request(
|
||||
*,
|
||||
account: Account,
|
||||
tenant_id: str,
|
||||
operation: KnowledgeFSOperation,
|
||||
) -> None:
|
||||
"""Enforce Dify's workspace policy before KFS performs resource authorization.
|
||||
|
||||
Args:
|
||||
account: Authenticated Dify account with its current workspace role.
|
||||
tenant_id: Current Dify workspace identifier.
|
||||
operation: Dify-maintained KnowledgeFS operation and policy metadata.
|
||||
|
||||
Raises:
|
||||
KnowledgeFSAccessDeniedError: The account lacks a required legacy or enterprise permission.
|
||||
"""
|
||||
if operation.requires_dataset_editor and not account.is_dataset_editor:
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS mutations require dataset edit access")
|
||||
if not RBACService.CheckAccess.check(
|
||||
tenant_id,
|
||||
account.id,
|
||||
scene=operation.rbac_permission.value,
|
||||
resource_type=RBACResourceScope.DATASET.value,
|
||||
):
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS operation is denied by workspace RBAC")
|
||||
|
||||
|
||||
def proxy_knowledge_fs_request(
|
||||
*,
|
||||
account: Account,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
tenant_id: str,
|
||||
accept: str | None = None,
|
||||
content_type: str | None = None,
|
||||
query: bytes | None = None,
|
||||
body: bytes | None = None,
|
||||
request_headers: _RequestHeaders | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
"""Authorize and forward one allowlisted KnowledgeFS request as a single use case."""
|
||||
operation = get_knowledge_fs_operation(method, path)
|
||||
authorize_knowledge_fs_request(
|
||||
account=account,
|
||||
tenant_id=tenant_id,
|
||||
operation=operation,
|
||||
)
|
||||
incoming_request_headers = {name.lower(): value for name, value in (request_headers or {}).items()}
|
||||
contract_request_headers = {
|
||||
name: incoming_request_headers[name] for name in operation.request_headers if name in incoming_request_headers
|
||||
}
|
||||
return _forward_knowledge_fs_request(
|
||||
account_id=account.id,
|
||||
method=method,
|
||||
path=path,
|
||||
tenant_id=tenant_id,
|
||||
accept=accept,
|
||||
content_type=content_type,
|
||||
query=query,
|
||||
body=body,
|
||||
request_headers=contract_request_headers,
|
||||
)
|
||||
|
||||
|
||||
def _forward_knowledge_fs_request(
|
||||
*,
|
||||
account_id: str,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
tenant_id: str,
|
||||
accept: str | None = None,
|
||||
content_type: str | None = None,
|
||||
query: bytes | None = None,
|
||||
body: bytes | None = None,
|
||||
request_headers: Mapping[str, str] | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
"""Forward one fixed-route request without parsing its KnowledgeFS payload.
|
||||
|
||||
Args:
|
||||
account_id: Current Dify account used as the KFS member identity.
|
||||
method: Allowlisted upstream HTTP method.
|
||||
path: Relative KnowledgeFS path under an allowlisted product surface.
|
||||
tenant_id: Current Dify workspace used as the KFS tenant identity.
|
||||
accept: Original Accept header, when present.
|
||||
content_type: Original request Content-Type header, when present.
|
||||
query: Original encoded query string from the Console request.
|
||||
body: Original request body, when present.
|
||||
request_headers: Contract-declared request headers forwarded by the Console adapter.
|
||||
|
||||
Returns:
|
||||
The KnowledgeFS response and its actual transport kind. Non-success responses are buffered.
|
||||
|
||||
Raises:
|
||||
KnowledgeFSConfigurationError: The connection is incomplete or blocked by outbound policy.
|
||||
KnowledgeFSRouteNotAllowedError: The path is outside the allowlisted product surface.
|
||||
KnowledgeFSTimeoutError: KnowledgeFS exceeds the configured timeout.
|
||||
KnowledgeFSTransportError: The request fails or its response cannot be safely bounded.
|
||||
|
||||
Each request is bound to stable Dify account and workspace principals with a short expiration.
|
||||
"""
|
||||
operation = get_knowledge_fs_operation(method, path)
|
||||
base_url = dify_config.KNOWLEDGE_FS_BASE_URL
|
||||
jwt_secret = dify_config.KNOWLEDGE_FS_JWT_SECRET
|
||||
if base_url is None or jwt_secret is None:
|
||||
raise KnowledgeFSConfigurationError("KnowledgeFS connection configuration is incomplete")
|
||||
now = datetime.now(UTC)
|
||||
token = jwt.encode(
|
||||
{
|
||||
"aud": _JWT_AUDIENCE,
|
||||
"caller_kind": "interactive",
|
||||
"dify_account_id": f"dify-account:{account_id}",
|
||||
"exp": now + timedelta(seconds=_JWT_TTL_SECONDS),
|
||||
"iat": now,
|
||||
"iss": _JWT_ISSUER,
|
||||
"scopes": [operation.required_scope],
|
||||
"sub": f"dify-workspace:{tenant_id}",
|
||||
"tenant_id": tenant_id,
|
||||
},
|
||||
jwt_secret.get_secret_value(),
|
||||
algorithm="HS256",
|
||||
)
|
||||
headers = {
|
||||
"Accept": accept or "application/json",
|
||||
"Accept-Encoding": "identity",
|
||||
"Authorization": f"Bearer {token}",
|
||||
}
|
||||
if body is not None:
|
||||
headers["Content-Type"] = content_type or "application/json"
|
||||
allowed_request_headers = set(operation.request_headers)
|
||||
for name, value in (request_headers or {}).items():
|
||||
normalized_name = name.lower()
|
||||
if normalized_name not in allowed_request_headers:
|
||||
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS request header is not allowed")
|
||||
headers[normalized_name] = value
|
||||
|
||||
try:
|
||||
upstream_url = httpx.URL(f"{base_url}/").join(operation.path)
|
||||
response = ssrf_proxy.make_request(
|
||||
method=operation.method,
|
||||
url=str(upstream_url),
|
||||
params=query,
|
||||
content=body,
|
||||
headers=headers,
|
||||
timeout=dify_config.KNOWLEDGE_FS_TIMEOUT_SECONDS,
|
||||
follow_redirects=False,
|
||||
max_retries=0,
|
||||
stream_response=True,
|
||||
)
|
||||
response_kind = _classify_response(operation, response)
|
||||
if response_kind == "stream":
|
||||
content_encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||||
if content_encoding not in {"", "identity"}:
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS streaming response used an unsupported encoding")
|
||||
_set_response_read_timeout(response, dify_config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS)
|
||||
return KnowledgeFSUpstreamResponse(response, response_kind, operation)
|
||||
|
||||
max_response_bytes = (
|
||||
operation.max_response_bytes
|
||||
if HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
|
||||
else _MAX_BUFFERED_RESPONSE_BYTES
|
||||
)
|
||||
buffered_response = ssrf_proxy.buffer_response(response, max_response_bytes=max_response_bytes)
|
||||
if buffered_response.content and not buffered_response.headers.get("content-type", "").strip():
|
||||
buffered_response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
|
||||
return KnowledgeFSUpstreamResponse(buffered_response, response_kind, operation)
|
||||
except ssrf_proxy.ResponseLimitError as exc:
|
||||
raise KnowledgeFSTransportError("KnowledgeFS response violated the proxy limit") from exc
|
||||
except ToolSSRFError as exc:
|
||||
raise KnowledgeFSConfigurationError("KnowledgeFS origin was blocked by outbound policy") from exc
|
||||
except httpx.TimeoutException as exc:
|
||||
raise KnowledgeFSTimeoutError("KnowledgeFS request timed out") from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise KnowledgeFSTransportError("KnowledgeFS transport request failed") from exc
|
||||
|
||||
|
||||
def get_knowledge_fs_operation(method: KnowledgeFSMethod, path: str) -> KnowledgeFSOperation:
|
||||
"""Resolve an exact operation and its transport/access contract metadata."""
|
||||
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS:
|
||||
if method == operation.method and _matches_route_template(operation.path, path):
|
||||
return operation._replace(path=path)
|
||||
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS route is not allowed")
|
||||
|
||||
|
||||
def _classify_response(operation: KnowledgeFSOperation, response: httpx.Response) -> KnowledgeFSResponseKind:
|
||||
"""Resolve the actual response kind from status and Content-Type before reading its body."""
|
||||
content_type = response.headers.get("content-type", "").partition(";")[0].strip().lower()
|
||||
is_success = HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
|
||||
if not is_success:
|
||||
if content_type and not _is_json_content_type(content_type):
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS error response used an unsupported media type")
|
||||
return "buffered"
|
||||
|
||||
if operation.response_kind == "stream":
|
||||
if content_type != "text/event-stream":
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS stream response used an unsupported media type")
|
||||
return "stream"
|
||||
if operation.response_kind == "binary":
|
||||
if content_type not in operation.response_media_types:
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS binary response used an unsupported media type")
|
||||
return "binary"
|
||||
if content_type and not _is_json_content_type(content_type):
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
|
||||
return "buffered"
|
||||
|
||||
|
||||
def _is_json_content_type(content_type: str) -> bool:
|
||||
return content_type == "application/json" or content_type.endswith("+json")
|
||||
|
||||
|
||||
def _set_response_read_timeout(response: httpx.Response, timeout_seconds: float | None) -> None:
|
||||
"""Set the body-read timeout after headers identify a valid SSE response."""
|
||||
try:
|
||||
request = response.request
|
||||
except RuntimeError:
|
||||
return
|
||||
timeout = request.extensions.get("timeout")
|
||||
if isinstance(timeout, dict):
|
||||
timeout["read"] = timeout_seconds
|
||||
|
||||
|
||||
def _matches_route_template(template: str, path: str) -> bool:
|
||||
"""Match path parameters without permitting encoded or traversal-like segments."""
|
||||
template_segments = template.split("/")
|
||||
path_segments = path.split("/")
|
||||
if len(template_segments) != len(path_segments):
|
||||
return False
|
||||
|
||||
for template_segment, path_segment in zip(template_segments, path_segments, strict=True):
|
||||
if template_segment.startswith("{") and template_segment.endswith("}"):
|
||||
if (
|
||||
not path_segment
|
||||
or path_segment in {".", ".."}
|
||||
or "\\" in path_segment
|
||||
or "%" in path_segment
|
||||
or "?" in path_segment
|
||||
or "#" in path_segment
|
||||
):
|
||||
return False
|
||||
continue
|
||||
if template_segment != path_segment:
|
||||
return False
|
||||
return True
|
||||
142
api/tests/unit_tests/configs/test_knowledge_fs_config.py
Normal file
142
api/tests/unit_tests/configs/test_knowledge_fs_config.py
Normal file
@ -0,0 +1,142 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic import SecretStr, ValidationError
|
||||
|
||||
from configs.extra.knowledge_fs_config import KnowledgeFSConfig
|
||||
|
||||
_REPOSITORY_ROOT = Path(__file__).resolve().parents[4]
|
||||
_KNOWLEDGE_FS_DOCKER_VARIABLES = (
|
||||
"KNOWLEDGE_FS_ENABLED",
|
||||
"KNOWLEDGE_FS_BASE_URL",
|
||||
"KNOWLEDGE_FS_JWT_SECRET",
|
||||
"KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS",
|
||||
"KNOWLEDGE_FS_TIMEOUT_SECONDS",
|
||||
)
|
||||
|
||||
|
||||
def test_knowledge_fs_config_normalizes_complete_connection() -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=True,
|
||||
KNOWLEDGE_FS_BASE_URL=" https://knowledge-fs.test/ ",
|
||||
KNOWLEDGE_FS_JWT_SECRET=" production-secret-with-at-least-32-bytes ",
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_BASE_URL == "https://knowledge-fs.test"
|
||||
assert isinstance(config.KNOWLEDGE_FS_JWT_SECRET, SecretStr)
|
||||
assert config.KNOWLEDGE_FS_JWT_SECRET.get_secret_value() == "production-secret-with-at-least-32-bytes"
|
||||
assert "production-secret" not in repr(config)
|
||||
assert "production-secret" not in config.model_dump_json()
|
||||
assert config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS == 300.0
|
||||
|
||||
|
||||
def test_knowledge_fs_config_treats_blank_connection_as_disabled() -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=" ",
|
||||
KNOWLEDGE_FS_JWT_SECRET="",
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_BASE_URL is None
|
||||
assert config.KNOWLEDGE_FS_JWT_SECRET is None
|
||||
assert config.KNOWLEDGE_FS_ENABLED is False
|
||||
|
||||
|
||||
def test_knowledge_fs_config_requires_connection_when_enabled() -> None:
|
||||
with pytest.raises(ValidationError, match="connection settings are required"):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_ENABLED=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_url", "jwt_secret"),
|
||||
[
|
||||
("https://knowledge-fs.test", None),
|
||||
(None, "production-secret-with-at-least-32-bytes"),
|
||||
],
|
||||
)
|
||||
def test_disabled_knowledge_fs_config_allows_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=False,
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_ENABLED is False
|
||||
|
||||
|
||||
def test_knowledge_fs_docker_config_is_not_shadowed_by_root_env() -> None:
|
||||
root_env_example = (_REPOSITORY_ROOT / "docker/.env.example").read_text(encoding="utf-8")
|
||||
api_env_example = (_REPOSITORY_ROOT / "docker/envs/core-services/api.env.example").read_text(encoding="utf-8")
|
||||
|
||||
for variable in _KNOWLEDGE_FS_DOCKER_VARIABLES:
|
||||
assert f"{variable}=" not in root_env_example
|
||||
assert f"{variable}=" in api_env_example
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_url", "jwt_secret"),
|
||||
[
|
||||
("https://knowledge-fs.test", None),
|
||||
(None, "production-secret-with-at-least-32-bytes"),
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
|
||||
with pytest.raises(ValidationError, match="must be configured together"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=True,
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_url", ["knowledge-fs.test", "ftp://knowledge-fs.test", "http:///missing-host"])
|
||||
def test_knowledge_fs_config_rejects_non_http_absolute_urls(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="absolute HTTP\\(S\\) URL"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_url",
|
||||
[
|
||||
"https://knowledge-fs.test:notaport",
|
||||
"https://knowledge-fs.test:65536",
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_invalid_ports(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="valid port"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_url",
|
||||
[
|
||||
"https://user:password@knowledge-fs.test",
|
||||
"https://knowledge-fs.test?region=us",
|
||||
"https://knowledge-fs.test#gateway",
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_unsafe_base_url_components(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="must not include credentials, query, or fragment"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 60.0001])
|
||||
def test_knowledge_fs_config_rejects_unbounded_timeouts(timeout_seconds: float) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_TIMEOUT_SECONDS=timeout_seconds)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 3600.0001])
|
||||
def test_knowledge_fs_config_rejects_unbounded_sse_read_timeouts(timeout_seconds: float) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=timeout_seconds)
|
||||
@ -0,0 +1,639 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
from inspect import unwrap
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from flask import Flask, Response
|
||||
from werkzeug.exceptions import (
|
||||
BadGateway,
|
||||
Forbidden,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
from controllers.console import bp
|
||||
from controllers.console.knowledge_fs_proxy import (
|
||||
_console_api_errors,
|
||||
_proxy_knowledge_fs_non_get,
|
||||
_proxy_request,
|
||||
_proxy_response,
|
||||
proxy_knowledge_fs_get,
|
||||
proxy_knowledge_fs_write,
|
||||
)
|
||||
from controllers.console.wraps import RBACPermission
|
||||
from services.knowledge_fs_proxy import (
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSMethod,
|
||||
KnowledgeFSOperation,
|
||||
KnowledgeFSResponseKind,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSUpstreamResponse,
|
||||
get_knowledge_fs_operation,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", True)
|
||||
|
||||
|
||||
def _upstream(
|
||||
response: httpx.Response,
|
||||
kind: KnowledgeFSResponseKind = "buffered",
|
||||
*,
|
||||
max_response_bytes: int | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
operation = KnowledgeFSOperation(
|
||||
operation_id="testOperation",
|
||||
method="GET",
|
||||
path="test",
|
||||
response_kind=kind,
|
||||
required_scope="knowledge-spaces:read",
|
||||
rbac_permission=RBACPermission.DATASET_READONLY,
|
||||
requires_dataset_editor=False,
|
||||
max_response_bytes=max_response_bytes
|
||||
or (64 * 1024 * 1024 if kind == "stream" else 25 * 1024 * 1024 if kind == "binary" else 1024 * 1024),
|
||||
request_headers=(),
|
||||
response_headers=(
|
||||
"content-security-policy",
|
||||
"x-content-type-options",
|
||||
"x-query-run-id",
|
||||
"x-session-id",
|
||||
),
|
||||
response_media_types=(),
|
||||
)
|
||||
return KnowledgeFSUpstreamResponse(response, kind, operation)
|
||||
|
||||
|
||||
def _set_current_workspace(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
editor: bool = True,
|
||||
has_edit_permission: bool = True,
|
||||
admin_or_owner: bool = True,
|
||||
) -> None:
|
||||
account = MagicMock(
|
||||
id="account-1",
|
||||
has_edit_permission=has_edit_permission,
|
||||
is_admin_or_owner=admin_or_owner,
|
||||
is_dataset_editor=editor,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.current_account_with_tenant",
|
||||
lambda: (account, "tenant-1"),
|
||||
)
|
||||
|
||||
|
||||
def _bypass_policy_wrappers(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy._proxy_knowledge_fs_non_get",
|
||||
unwrap(_proxy_knowledge_fs_non_get),
|
||||
)
|
||||
|
||||
|
||||
def test_console_blueprint_registers_generic_knowledge_fs_routes() -> None:
|
||||
app = Flask("knowledge-fs-route-registration")
|
||||
app.register_blueprint(bp)
|
||||
adapter = app.url_map.bind("localhost")
|
||||
|
||||
get_endpoint, get_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="GET",
|
||||
)
|
||||
assert get_endpoint.endswith("proxy_knowledge_fs_get")
|
||||
assert get_values == {"upstream_path": "knowledge-spaces"}
|
||||
for method in ("DELETE", "PATCH", "POST", "PUT"):
|
||||
write_endpoint, write_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces/space-1",
|
||||
method=method,
|
||||
)
|
||||
assert write_endpoint.endswith("proxy_knowledge_fs_write")
|
||||
assert write_values == {"upstream_path": "knowledge-spaces/space-1"}
|
||||
options_endpoint, options_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="OPTIONS",
|
||||
)
|
||||
assert options_endpoint.endswith("proxy_knowledge_fs_get")
|
||||
assert options_values == {"upstream_path": "knowledge-spaces"}
|
||||
|
||||
|
||||
def test_proxy_is_hidden_when_knowledge_fs_is_disabled(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", False)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="GET"):
|
||||
with pytest.raises(NotFound):
|
||||
_proxy_request("GET", "knowledge-spaces")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method"),
|
||||
[
|
||||
(proxy_knowledge_fs_get, "GET"),
|
||||
(proxy_knowledge_fs_write, "POST"),
|
||||
],
|
||||
)
|
||||
def test_proxy_routes_are_hidden_before_downstream_work_when_disabled(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route,
|
||||
method: KnowledgeFSMethod,
|
||||
) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", False)
|
||||
proxy_request = MagicMock()
|
||||
proxy_non_get = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy_request)
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_knowledge_fs_non_get", proxy_non_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method=method):
|
||||
response = app.make_response(route("knowledge-spaces"))
|
||||
|
||||
assert response.status_code == 404
|
||||
proxy_request.assert_not_called()
|
||||
proxy_non_get.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["HEAD", "OPTIONS"])
|
||||
def test_generic_get_hides_unregistered_methods(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: str,
|
||||
) -> None:
|
||||
proxy = MagicMock(return_value=Response(status=200))
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method=method,
|
||||
):
|
||||
with pytest.raises(NotFound):
|
||||
route("knowledge-spaces")
|
||||
|
||||
proxy.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method", "path", "permission"),
|
||||
[
|
||||
(proxy_knowledge_fs_get, "GET", "knowledge-spaces", RBACPermission.DATASET_READONLY),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"POST",
|
||||
"knowledge-spaces",
|
||||
RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_generic_routes_delegate_to_the_authorized_service_use_case(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
permission: RBACPermission,
|
||||
) -> None:
|
||||
response_kind: KnowledgeFSResponseKind = "buffered"
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset" if response_kind == "binary" else b"{}",
|
||||
headers={"Content-Type": "application/octet-stream" if response_kind == "binary" else "application/json"},
|
||||
)
|
||||
proxy = MagicMock(return_value=_upstream(upstream, response_kind))
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
proxy,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
raw_route = unwrap(route)
|
||||
|
||||
with app.test_request_context(f"/console/api/knowledge-fs/{path}", method=method, data=b"{}"):
|
||||
response = raw_route(path)
|
||||
|
||||
assert isinstance(response, Response)
|
||||
assert proxy.call_args.kwargs["account"].id == "account-1"
|
||||
assert proxy.call_args.kwargs["tenant_id"] == "tenant-1"
|
||||
assert proxy.call_args.kwargs["method"] == method
|
||||
assert proxy.call_args.kwargs["path"] == path
|
||||
assert get_knowledge_fs_operation(method, path).rbac_permission == permission
|
||||
|
||||
|
||||
def test_read_post_applies_knowledge_rate_limit_once(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr("controllers.common.wraps.dify_config.RBAC_ENABLED", False)
|
||||
account = MagicMock(id="account-1", is_dataset_editor=True)
|
||||
|
||||
def current_workspace() -> tuple[MagicMock, str]:
|
||||
return account, "tenant-1"
|
||||
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr(
|
||||
"services.knowledge_fs_proxy.RBACService.CheckAccess.check",
|
||||
MagicMock(return_value=True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
|
||||
MagicMock(return_value=MagicMock(enabled=True, limit=10)),
|
||||
)
|
||||
zadd = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zadd", zadd)
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zremrangebyscore", MagicMock())
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zcard", MagicMock(return_value=1))
|
||||
proxy = MagicMock(return_value=Response(status=200))
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="POST"):
|
||||
response = _proxy_knowledge_fs_non_get("POST", "knowledge-spaces")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
zadd.assert_called_once()
|
||||
proxy.assert_called_once_with("POST", "knowledge-spaces")
|
||||
|
||||
|
||||
def test_denied_write_does_not_consume_the_workspace_rate_limit(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=False)
|
||||
|
||||
def current_workspace() -> tuple[MagicMock, str]:
|
||||
return account, "tenant-1"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.current_account_with_tenant",
|
||||
current_workspace,
|
||||
)
|
||||
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
|
||||
MagicMock(return_value=MagicMock(enabled=True, limit=10)),
|
||||
)
|
||||
zadd = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zadd", zadd)
|
||||
upstream_request = MagicMock()
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", upstream_request)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=b"{}",
|
||||
content_type="application/json",
|
||||
):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
zadd.assert_not_called()
|
||||
upstream_request.assert_not_called()
|
||||
|
||||
|
||||
def test_generic_get_forwards_path_query_and_raw_response(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
200,
|
||||
content=gzip.compress(b'{"items":[],"nextCursor":null}'),
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"Content-Encoding": "gzip",
|
||||
"Content-Disposition": 'attachment; filename="result.json"',
|
||||
"Content-Type": "application/json",
|
||||
"Retry-After": "3",
|
||||
"Set-Cookie": "kfs=secret",
|
||||
"X-Trace-Id": "trace-1",
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
query_string=[("limit", "20"), ("cursor", "first"), ("cursor", "second")],
|
||||
):
|
||||
response = route("knowledge-spaces")
|
||||
|
||||
request = forward.call_args.kwargs
|
||||
assert request["account"].id == "account-1"
|
||||
assert request["method"] == "GET"
|
||||
assert request["path"] == "knowledge-spaces"
|
||||
assert request["tenant_id"] == "tenant-1"
|
||||
assert request["accept"] is None
|
||||
assert request["content_type"] is None
|
||||
assert request["query"] == b"limit=20&cursor=first&cursor=second"
|
||||
assert request["body"] is None
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 200
|
||||
assert response.get_json() == {"items": [], "nextCursor": None}
|
||||
assert response.headers["Cache-Control"] == "no-store"
|
||||
assert response.headers["Content-Disposition"] == 'attachment; filename="result.json"'
|
||||
assert response.headers["Retry-After"] == "3"
|
||||
assert response.headers["X-Trace-Id"] == "trace-1"
|
||||
assert "Content-Encoding" not in response.headers
|
||||
assert "Set-Cookie" not in response.headers
|
||||
|
||||
|
||||
def test_generic_write_forwards_path_raw_body_and_current_tenant(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
method = "POST"
|
||||
path = "knowledge-spaces"
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
201,
|
||||
content=b'{"id":"space-1","tenantId":"tenant-1"}',
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
body = b'{"idempotencyKey":"create-product-docs","name":"Product docs"}'
|
||||
|
||||
with app.test_request_context(
|
||||
f"/console/api/knowledge-fs/{path}",
|
||||
method=method,
|
||||
data=body,
|
||||
content_type="application/json",
|
||||
):
|
||||
response = route(path)
|
||||
|
||||
request = forward.call_args.kwargs
|
||||
assert request["account"].id == "account-1"
|
||||
assert request["method"] == method
|
||||
assert request["path"] == path
|
||||
assert request["tenant_id"] == "tenant-1"
|
||||
assert request["accept"] is None
|
||||
assert request["content_type"] == "application/json"
|
||||
assert request["query"] is None
|
||||
assert request["body"] == body
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 201
|
||||
assert response.get_json()["tenantId"] == "tenant-1"
|
||||
|
||||
|
||||
def test_generic_write_forwards_contract_declared_request_headers(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(202, content=b'{"status":"accepted"}', headers={"Content-Type": "application/json"})
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
body = b'{"name":"Product docs"}'
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=body,
|
||||
content_type="application/json",
|
||||
headers={"X-Trace-Id": "trace-1"},
|
||||
):
|
||||
response = route("knowledge-spaces")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 202
|
||||
assert forward.call_args.kwargs["request_headers"].get("X-Trace-Id") == "trace-1"
|
||||
|
||||
|
||||
def test_contract_response_headers_cannot_bypass_the_proxy_denylist() -> None:
|
||||
denied_headers = (
|
||||
"connection",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
)
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"{}",
|
||||
headers=dict.fromkeys(denied_headers, "blocked"),
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=denied_headers,
|
||||
max_response_bytes=1024 * 1024,
|
||||
)
|
||||
|
||||
for name in denied_headers:
|
||||
assert name not in response.headers
|
||||
|
||||
|
||||
def test_contract_response_headers_forward_binary_hardening_headers() -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset",
|
||||
headers={
|
||||
"Content-Security-Policy": "sandbox; default-src 'none'",
|
||||
"Content-Type": "image/png",
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
},
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream, "binary"),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=("content-security-policy", "x-content-type-options"),
|
||||
max_response_bytes=25 * 1024 * 1024,
|
||||
)
|
||||
|
||||
assert response.headers["Content-Security-Policy"] == "sandbox; default-src 'none'"
|
||||
assert response.headers["X-Content-Type-Options"] == "nosniff"
|
||||
|
||||
|
||||
def test_authorized_service_denial_is_exposed_as_forbidden(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSAccessDeniedError("workspace access denied")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_generic_post_rejects_oversized_body(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
forward = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._MAX_PROXY_BODY_BYTES", 8)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=b"x" * 9,
|
||||
content_type="application/json",
|
||||
):
|
||||
with pytest.raises(RequestEntityTooLarge):
|
||||
route("knowledge-spaces")
|
||||
|
||||
forward.assert_not_called()
|
||||
|
||||
|
||||
def test_server_credential_rejection_is_not_exposed_as_browser_auth_failure(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
401,
|
||||
content=b'{"error":"invalid server credential"}',
|
||||
headers={"Content-Type": "application/json", "WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(BadGateway, match="authentication failed"):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_resource_authorization_rejection_is_exposed_as_forbidden(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
403,
|
||||
content=b'{"error":"resource access denied"}',
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_contract_response_headers_are_deduplicated_case_insensitively() -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset",
|
||||
headers={
|
||||
"Cache-Control": "private",
|
||||
"Content-Disposition": 'inline; filename="asset.png"',
|
||||
"Content-Type": "image/png",
|
||||
},
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream, "binary"),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=("cache-control", "content-disposition"),
|
||||
max_response_bytes=25 * 1024 * 1024,
|
||||
)
|
||||
|
||||
assert response.headers.getlist("Cache-Control") == ["private"]
|
||||
assert response.headers.getlist("Content-Disposition") == ['inline; filename="asset.png"']
|
||||
|
||||
|
||||
def test_configuration_error_is_reported_as_unavailable(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSConfigurationError("missing token")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(ServiceUnavailable, match="misconfigured"):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_disallowed_kfs_route_is_hidden_as_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSRouteNotAllowedError("blocked")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/openapi.json"):
|
||||
with pytest.raises(NotFound):
|
||||
route("openapi.json")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["PATCH", "POST"])
|
||||
def test_disallowed_non_get_route_is_hidden_as_not_found(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: KnowledgeFSMethod,
|
||||
) -> None:
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/not-a-route", method=method):
|
||||
with pytest.raises(NotFound):
|
||||
route("not-a-route")
|
||||
|
||||
|
||||
def test_raw_route_uses_console_json_error_handler(app: Flask) -> None:
|
||||
def forbidden(_upstream_path: str) -> Response:
|
||||
raise Forbidden("blocked")
|
||||
|
||||
route = _console_api_errors(forbidden)
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
response = app.make_response(route("knowledge-spaces"))
|
||||
|
||||
assert response.status_code == 403
|
||||
assert response.is_json
|
||||
assert response.get_json() == {"code": "forbidden", "message": "blocked", "status": 403}
|
||||
@ -1,3 +1,5 @@
|
||||
import gzip
|
||||
from typing import override
|
||||
from unittest.mock import ANY, MagicMock, call, patch
|
||||
|
||||
import httpx
|
||||
@ -5,10 +7,13 @@ import pytest
|
||||
|
||||
from core.helper.ssrf_proxy import (
|
||||
SSRF_DEFAULT_MAX_RETRIES,
|
||||
ResponseTooLargeError,
|
||||
SSRFProxy,
|
||||
UnsupportedResponseEncodingError,
|
||||
_build_ssrf_client,
|
||||
_get_user_provided_host_header,
|
||||
_to_graphon_http_response,
|
||||
buffer_response,
|
||||
graphon_ssrf_proxy,
|
||||
make_request,
|
||||
max_retries_exceeded_error,
|
||||
@ -21,12 +26,91 @@ def test_successful_request(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com")
|
||||
assert response.status_code == 200
|
||||
mock_client.request.assert_called_once()
|
||||
mock_client.build_request.assert_called_once()
|
||||
mock_client.send.assert_called_once()
|
||||
|
||||
|
||||
def test_buffer_response_rejects_encoded_response_before_decoding() -> None:
|
||||
payload = b"x" * (8 * 1024 * 1024)
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
content=gzip.compress(payload),
|
||||
headers={"Content-Encoding": "gzip"},
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
with pytest.raises(UnsupportedResponseEncodingError, match="content encoding gzip"):
|
||||
buffer_response(response, max_response_bytes=1024 * 1024)
|
||||
|
||||
|
||||
def test_buffer_response_returns_response_within_decoded_byte_limit() -> None:
|
||||
payload = b"response-body"
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
content=payload,
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
streaming_response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
response = buffer_response(streaming_response, max_response_bytes=32)
|
||||
|
||||
assert response.content == payload
|
||||
assert str(response.request.url) == "http://example.com"
|
||||
|
||||
|
||||
def test_buffer_response_rejects_identity_response_exceeding_byte_limit() -> None:
|
||||
payload = b"response-body"
|
||||
transport = httpx.MockTransport(lambda request: httpx.Response(200, content=payload, request=request))
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
with pytest.raises(ResponseTooLargeError, match="response exceeded 8 bytes"):
|
||||
buffer_response(response, max_response_bytes=8)
|
||||
|
||||
|
||||
def test_request_can_return_an_open_stream_the_caller_closes() -> None:
|
||||
class EventStream(httpx.SyncByteStream):
|
||||
@override
|
||||
def __iter__(self):
|
||||
yield b"event: delta\ndata: first\n\n"
|
||||
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
stream=EventStream(),
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
httpx.Client(transport=transport) as client,
|
||||
patch("core.helper.ssrf_proxy._get_ssrf_client", return_value=client),
|
||||
):
|
||||
response = make_request(
|
||||
"GET",
|
||||
"http://example.com/events",
|
||||
max_retries=0,
|
||||
stream_response=True,
|
||||
)
|
||||
|
||||
assert response.is_stream_consumed is False
|
||||
assert response.is_closed is False
|
||||
assert b"".join(response.iter_bytes()) == b"event: delta\ndata: first\n\n"
|
||||
response.close()
|
||||
|
||||
assert response.is_closed
|
||||
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@ -34,7 +118,7 @@ def test_retry_exceed_max_retries(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with pytest.raises(Exception) as e:
|
||||
@ -47,13 +131,13 @@ def test_force_list_response_returns_when_retries_disabled(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com", max_retries=0)
|
||||
|
||||
assert response is mock_response
|
||||
mock_client.request.assert_called_once()
|
||||
mock_client.send.assert_called_once()
|
||||
|
||||
|
||||
def test_build_ssrf_client_passes_ssl_verify_to_proxy_mount_transports():
|
||||
@ -126,15 +210,15 @@ def test_host_header_preservation_with_user_header(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
custom_host = "custom.example.com:8080"
|
||||
response = make_request("GET", "http://example.com", headers={"Host": custom_host})
|
||||
|
||||
assert response.status_code == 200
|
||||
# Verify client.request was called with the host header preserved (lowercase)
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
# Verify the request was built with the host header preserved (lowercase)
|
||||
call_kwargs = mock_client.build_request.call_args.kwargs
|
||||
assert call_kwargs["headers"]["host"] == custom_host
|
||||
|
||||
|
||||
@ -145,36 +229,36 @@ def test_host_header_preservation_case_insensitive(mock_get_client, host_key):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com", headers={host_key: "api.example.com"})
|
||||
|
||||
assert response.status_code == 200
|
||||
# Host header should be normalized to lowercase "host"
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.build_request.call_args.kwargs
|
||||
assert call_kwargs["headers"]["host"] == "api.example.com"
|
||||
|
||||
|
||||
class TestFollowRedirectsParameter:
|
||||
"""Tests for follow_redirects parameter handling.
|
||||
|
||||
These tests verify that follow_redirects is correctly passed to client.request().
|
||||
These tests verify that follow_redirects is correctly passed to client.send().
|
||||
"""
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
def test_follow_redirects_passed_to_request(self, mock_get_client):
|
||||
"""Verify follow_redirects IS passed to client.request()."""
|
||||
"""Verify follow_redirects IS passed to client.send()."""
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
make_request("GET", "http://example.com", follow_redirects=True)
|
||||
|
||||
# Verify follow_redirects was passed to request
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
# Verify follow_redirects was passed to send
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@ -183,14 +267,14 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# Use allow_redirects (requests-style parameter)
|
||||
make_request("GET", "http://example.com", allow_redirects=True)
|
||||
|
||||
# Verify it was converted to follow_redirects
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
assert "allow_redirects" not in call_kwargs
|
||||
|
||||
@ -200,13 +284,13 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
make_request("GET", "http://example.com")
|
||||
|
||||
# follow_redirects should not be in kwargs, letting httpx use its default
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert "follow_redirects" not in call_kwargs
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@ -215,13 +299,13 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# Both specified - follow_redirects should take precedence
|
||||
make_request("GET", "http://example.com", allow_redirects=False, follow_redirects=True)
|
||||
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
|
||||
|
||||
|
||||
356
api/tests/unit_tests/dev/test_generate_knowledge_fs_contract.py
Normal file
356
api/tests/unit_tests/dev/test_generate_knowledge_fs_contract.py
Normal file
@ -0,0 +1,356 @@
|
||||
"""Tests for pinned KnowledgeFS declaration validation."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
from dev import generate_knowledge_fs_contract as contract_validator
|
||||
from dev.generate_knowledge_fs_contract import ContractDeclaration, validate_declarations
|
||||
from services.knowledge_fs_proxy import KNOWLEDGE_FS_CONSOLE_OPERATIONS, KnowledgeFSOperation
|
||||
|
||||
|
||||
def test_contract_cli_updates_checks_and_detects_openapi_drift(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
repository = tmp_path / "knowledge-fs"
|
||||
repository.mkdir()
|
||||
subprocess.run(["git", "init", "--quiet"], cwd=repository, check=True)
|
||||
subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"-c",
|
||||
"user.email=contract-test@example.com",
|
||||
"-c",
|
||||
"user.name=Contract Test",
|
||||
"commit",
|
||||
"--allow-empty",
|
||||
"--quiet",
|
||||
"-m",
|
||||
"fixture",
|
||||
],
|
||||
cwd=repository,
|
||||
check=True,
|
||||
)
|
||||
commit = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=repository, check=True, capture_output=True, text=True
|
||||
).stdout.strip()
|
||||
|
||||
document = console_registry_document()
|
||||
executable_directory = tmp_path / "bin"
|
||||
executable_directory.mkdir()
|
||||
fake_pnpm = executable_directory / "pnpm"
|
||||
write_fake_pnpm(fake_pnpm, document)
|
||||
|
||||
lock_path = tmp_path / "knowledge-fs-contract.lock.json"
|
||||
lock_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"commit": "",
|
||||
"openapiSha256": "",
|
||||
"repository": "https://github.com/langgenius/knowledge-fs",
|
||||
}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(contract_validator, "LOCK_PATH", lock_path)
|
||||
monkeypatch.setenv("PATH", f"{executable_directory}{os.pathsep}{os.environ['PATH']}")
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["generate_knowledge_fs_contract.py", "--repository", str(repository), "--update-lock"],
|
||||
)
|
||||
contract_validator.main()
|
||||
|
||||
updated_lock = json.loads(lock_path.read_text())
|
||||
assert updated_lock["commit"] == commit
|
||||
assert set(updated_lock) == {"commit", "openapiSha256", "repository"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["generate_knowledge_fs_contract.py", "--repository", str(repository), "--check"],
|
||||
)
|
||||
contract_validator.main()
|
||||
|
||||
with monkeypatch.context() as registry_drift:
|
||||
registry_drift.setattr(
|
||||
contract_validator,
|
||||
"console_contract_declarations",
|
||||
lambda: (declaration(method="DELETE"),),
|
||||
raising=False,
|
||||
)
|
||||
with pytest.raises(ValueError, match="listKnowledgeSpaces.*method.*expected.*received"):
|
||||
contract_validator.main()
|
||||
|
||||
write_fake_pnpm(fake_pnpm, {"paths": {}})
|
||||
with pytest.raises(RuntimeError, match="OpenAPI hash mismatch"):
|
||||
contract_validator.main()
|
||||
|
||||
|
||||
def test_contract_script_loads_runtime_registry_outside_api_directory(tmp_path: Path) -> None:
|
||||
script_path = Path(contract_validator.__file__).resolve()
|
||||
command = (
|
||||
"import runpy; "
|
||||
f"namespace = runpy.run_path({str(script_path)!r}); "
|
||||
"print(len(namespace['console_contract_declarations']()))"
|
||||
)
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", command],
|
||||
cwd=tmp_path,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
assert result.stdout.strip() == "2"
|
||||
|
||||
|
||||
def test_validate_declarations_accepts_matching_contract() -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
document = {"paths": {"/knowledge-spaces": {"get": route}}}
|
||||
|
||||
validate_declarations(
|
||||
document,
|
||||
(
|
||||
declaration(
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_console_operation_registry_matches_contract() -> None:
|
||||
list_route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
create_route = operation("knowledge-spaces:write", "createKnowledgeSpace")
|
||||
for route in (list_route, create_route):
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
|
||||
validate_declarations(
|
||||
{
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": list_route,
|
||||
"post": create_route,
|
||||
}
|
||||
}
|
||||
},
|
||||
tuple(_contract_declaration(operation) for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS),
|
||||
)
|
||||
|
||||
|
||||
def console_registry_document() -> dict[str, object]:
|
||||
list_route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
create_route = operation("knowledge-spaces:write", "createKnowledgeSpace")
|
||||
for route in (list_route, create_route):
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
return {"paths": {"/knowledge-spaces": {"get": list_route, "post": create_route}}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("method", "POST"),
|
||||
("path", "spaces"),
|
||||
("required_scope", "knowledge-spaces:write"),
|
||||
("response_kind", "stream"),
|
||||
("max_response_bytes", 2_097_152),
|
||||
("request_headers", ("authorization",)),
|
||||
("response_headers", ("cache-control",)),
|
||||
("response_media_types", ("text/event-stream",)),
|
||||
],
|
||||
)
|
||||
def test_validate_declarations_reports_contract_field_drift(field: str, value: object) -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match=rf"listKnowledgeSpaces.*{field}.*expected.*received"):
|
||||
validate_declarations(document, (declaration(**{field: value}),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_unknown_operation_id() -> None:
|
||||
with pytest.raises(ValueError, match="no operationId: listKnowledgeSpaces"):
|
||||
validate_declarations({"paths": {}}, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_duplicate_declared_operation_ids() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="registry has duplicate operationId: listKnowledgeSpaces"):
|
||||
validate_declarations(document, (declaration(), declaration()))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_duplicate_upstream_operation_ids() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
"/spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="OpenAPI has duplicate operationId: listKnowledgeSpaces"):
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_ignores_undeclared_operations() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
"/internal-maintenance": {
|
||||
"head": {"responses": {"200": {}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_preserves_public_operation_scope() -> None:
|
||||
document = {"paths": {"/health": {"get": operation(None, "getHealth", security=[])}}}
|
||||
|
||||
validate_declarations(
|
||||
document,
|
||||
(
|
||||
declaration(
|
||||
operation_id="getHealth",
|
||||
path="health",
|
||||
required_scope=None,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_unsupported_declared_method() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"head": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="does not support HEAD /knowledge-spaces"):
|
||||
validate_declarations(document, (declaration(method="HEAD"),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_non_absolute_upstream_path() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="path must be absolute: knowledge-spaces"):
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [None, True, 0, "1048576"])
|
||||
def test_validate_declarations_rejects_invalid_response_byte_limits(value: object) -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["x-knowledge-fs-max-response-bytes"] = value
|
||||
|
||||
with pytest.raises(ValueError, match="no valid response byte limit"):
|
||||
validate_declarations({"paths": {"/knowledge-spaces": {"get": route}}}, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_request_header_references() -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["parameters"] = [{"$ref": "#/components/parameters/TraceId"}]
|
||||
|
||||
with pytest.raises(ValueError, match="request header references are not supported"):
|
||||
validate_declarations({"paths": {"/knowledge-spaces": {"get": route}}}, (declaration(),))
|
||||
|
||||
|
||||
def operation(scope: str | None, operation_id: str, **overrides: object) -> dict[str, object]:
|
||||
value: dict[str, object] = {
|
||||
"operationId": operation_id,
|
||||
"responses": {"200": {"content": {"application/json": {}}}},
|
||||
"x-knowledge-fs-max-response-bytes": 1_048_576,
|
||||
}
|
||||
if scope is not None:
|
||||
value["x-knowledge-fs-required-scope"] = scope
|
||||
value.update(overrides)
|
||||
return value
|
||||
|
||||
|
||||
def declaration(**overrides: object) -> ContractDeclaration:
|
||||
value: dict[str, object] = {
|
||||
"operation_id": "listKnowledgeSpaces",
|
||||
"method": "GET",
|
||||
"path": "knowledge-spaces",
|
||||
"required_scope": "knowledge-spaces:read",
|
||||
"response_kind": "buffered",
|
||||
"max_response_bytes": 1_048_576,
|
||||
"request_headers": (),
|
||||
"response_headers": (),
|
||||
"response_media_types": ("application/json",),
|
||||
}
|
||||
value.update(overrides)
|
||||
return cast(ContractDeclaration, value)
|
||||
|
||||
|
||||
def _contract_declaration(operation: KnowledgeFSOperation) -> ContractDeclaration:
|
||||
return {
|
||||
"operation_id": operation.operation_id,
|
||||
"method": operation.method,
|
||||
"path": operation.path,
|
||||
"required_scope": operation.required_scope,
|
||||
"response_kind": operation.response_kind,
|
||||
"max_response_bytes": operation.max_response_bytes,
|
||||
"request_headers": operation.request_headers,
|
||||
"response_headers": operation.response_headers,
|
||||
"response_media_types": operation.response_media_types,
|
||||
}
|
||||
|
||||
|
||||
def write_fake_pnpm(path: Path, document: dict[str, object]) -> None:
|
||||
path.write_text(
|
||||
"#!/usr/bin/env python3\n"
|
||||
"import sys\n"
|
||||
"from pathlib import Path\n"
|
||||
"output = Path(sys.argv[sys.argv.index('--output') + 1])\n"
|
||||
f"output.write_text({json.dumps(document)!r})\n"
|
||||
)
|
||||
path.chmod(0o755)
|
||||
298
api/tests/unit_tests/services/test_knowledge_fs_proxy.py
Normal file
298
api/tests/unit_tests/services/test_knowledge_fs_proxy.py
Normal file
@ -0,0 +1,298 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
import pytest
|
||||
from pydantic import SecretStr
|
||||
|
||||
from core.helper import ssrf_proxy
|
||||
from core.rbac import RBACPermission
|
||||
from core.tools.errors import ToolSSRFError
|
||||
from services.knowledge_fs_proxy import (
|
||||
KNOWLEDGE_FS_CONSOLE_OPERATIONS,
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSMethod,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSTimeoutError,
|
||||
KnowledgeFSTransportError,
|
||||
authorize_knowledge_fs_request,
|
||||
get_knowledge_fs_operation,
|
||||
proxy_knowledge_fs_request,
|
||||
)
|
||||
from services.knowledge_fs_proxy import (
|
||||
_forward_knowledge_fs_request as forward_knowledge_fs_request,
|
||||
)
|
||||
|
||||
_JWT_SECRET = "production-secret-with-at-least-32-bytes"
|
||||
|
||||
|
||||
def _set_config(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
base_url: str | None = "http://knowledge-fs.test",
|
||||
timeout_seconds: float = 7.5,
|
||||
jwt_secret: str | None = _JWT_SECRET,
|
||||
) -> None:
|
||||
values = {
|
||||
"KNOWLEDGE_FS_BASE_URL": base_url,
|
||||
"KNOWLEDGE_FS_TIMEOUT_SECONDS": timeout_seconds,
|
||||
"KNOWLEDGE_FS_JWT_SECRET": SecretStr(jwt_secret) if jwt_secret is not None else None,
|
||||
}
|
||||
for name, value in values.items():
|
||||
monkeypatch.setattr(f"services.knowledge_fs_proxy.dify_config.{name}", value, raising=False)
|
||||
|
||||
|
||||
def test_console_registry_starts_with_list_and_create_operations() -> None:
|
||||
assert tuple(operation.operation_id for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS) == (
|
||||
"listKnowledgeSpaces",
|
||||
"createKnowledgeSpace",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "operation_id", "scope", "permission", "requires_dataset_editor"),
|
||||
[
|
||||
("GET", "listKnowledgeSpaces", "knowledge-spaces:read", RBACPermission.DATASET_READONLY, False),
|
||||
(
|
||||
"POST",
|
||||
"createKnowledgeSpace",
|
||||
"knowledge-spaces:write",
|
||||
RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_console_registry_preserves_contract_and_policy(
|
||||
method: KnowledgeFSMethod,
|
||||
operation_id: str,
|
||||
scope: str,
|
||||
permission: RBACPermission,
|
||||
requires_dataset_editor: bool,
|
||||
) -> None:
|
||||
operation = get_knowledge_fs_operation(method, "knowledge-spaces")
|
||||
|
||||
assert operation.operation_id == operation_id
|
||||
assert operation.required_scope == scope
|
||||
assert operation.rbac_permission == permission
|
||||
assert operation.requires_dataset_editor is requires_dataset_editor
|
||||
assert operation.max_response_bytes == 1_048_576
|
||||
assert operation.request_headers == ("x-trace-id",)
|
||||
assert operation.response_headers == ("x-trace-id",)
|
||||
assert operation.response_media_types == ("application/json",)
|
||||
|
||||
|
||||
def test_unconfigured_kfs_is_rejected_before_external_io(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_config(monkeypatch, base_url=None, jwt_secret=None)
|
||||
request = MagicMock()
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", request)
|
||||
|
||||
with pytest.raises(KnowledgeFSConfigurationError, match="incomplete"):
|
||||
forward_knowledge_fs_request(
|
||||
account_id="account-dev", method="GET", path="knowledge-spaces", tenant_id="tenant-dev"
|
||||
)
|
||||
|
||||
request.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["GET", "POST"])
|
||||
def test_list_and_create_forward_raw_request(monkeypatch: pytest.MonkeyPatch, method: KnowledgeFSMethod) -> None:
|
||||
_set_config(monkeypatch, base_url="http://knowledge-fs.test/gateway")
|
||||
response = httpx.Response(201, content=b'{"id":"space-1"}', headers={"Content-Type": "application/json"})
|
||||
request = MagicMock(return_value=response)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", request)
|
||||
body = b'{"name":"Product docs"}' if method == "POST" else None
|
||||
query = b"limit=20&cursor=first" if method == "GET" else None
|
||||
|
||||
result = forward_knowledge_fs_request(
|
||||
account_id="account-dev",
|
||||
method=method,
|
||||
path="knowledge-spaces",
|
||||
tenant_id="tenant-dev",
|
||||
query=query,
|
||||
body=body,
|
||||
)
|
||||
|
||||
assert result.response.content == response.content
|
||||
assert result.response_kind == "buffered"
|
||||
assert request.call_args.kwargs["method"] == method
|
||||
assert request.call_args.kwargs["url"] == "http://knowledge-fs.test/gateway/knowledge-spaces"
|
||||
assert request.call_args.kwargs["params"] == query
|
||||
assert request.call_args.kwargs["content"] == body
|
||||
assert request.call_args.kwargs["follow_redirects"] is False
|
||||
assert request.call_args.kwargs["max_retries"] == 0
|
||||
assert request.call_args.kwargs["stream_response"] is True
|
||||
|
||||
|
||||
def test_proxy_forwards_only_registry_declared_headers(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=True)
|
||||
upstream = MagicMock()
|
||||
forward = MagicMock(return_value=upstream)
|
||||
monkeypatch.setattr(
|
||||
"services.knowledge_fs_proxy.RBACService.CheckAccess.check",
|
||||
MagicMock(return_value=True),
|
||||
)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy._forward_knowledge_fs_request", forward)
|
||||
|
||||
result = proxy_knowledge_fs_request(
|
||||
account=account,
|
||||
method="POST",
|
||||
path="knowledge-spaces",
|
||||
tenant_id="tenant-1",
|
||||
request_headers={"Authorization": "browser-secret", "X-Trace-Id": "trace-1"},
|
||||
)
|
||||
|
||||
assert result is upstream
|
||||
assert forward.call_args.kwargs["request_headers"] == {"x-trace-id": "trace-1"}
|
||||
|
||||
|
||||
def test_authorization_rejects_workspace_rbac_denial(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=True)
|
||||
check_access = MagicMock(return_value=False)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access)
|
||||
|
||||
with pytest.raises(KnowledgeFSAccessDeniedError):
|
||||
authorize_knowledge_fs_request(
|
||||
account=account,
|
||||
tenant_id="tenant-1",
|
||||
operation=get_knowledge_fs_operation("GET", "knowledge-spaces"),
|
||||
)
|
||||
|
||||
check_access.assert_called_once_with(
|
||||
"tenant-1",
|
||||
"account-1",
|
||||
scene="dataset_readonly",
|
||||
resource_type="dataset",
|
||||
)
|
||||
|
||||
|
||||
def test_create_rejects_non_dataset_editor_before_rbac(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=False)
|
||||
check_access = MagicMock(return_value=True)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access)
|
||||
|
||||
with pytest.raises(KnowledgeFSAccessDeniedError, match="dataset edit access"):
|
||||
authorize_knowledge_fs_request(
|
||||
account=account,
|
||||
tenant_id="tenant-1",
|
||||
operation=get_knowledge_fs_operation("POST", "knowledge-spaces"),
|
||||
)
|
||||
|
||||
check_access.assert_not_called()
|
||||
|
||||
|
||||
def test_authorization_uses_the_declared_editor_policy(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=False)
|
||||
check_access = MagicMock(return_value=True)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access)
|
||||
operation = get_knowledge_fs_operation("POST", "knowledge-spaces")._replace(requires_dataset_editor=False)
|
||||
|
||||
authorize_knowledge_fs_request(account=account, tenant_id="tenant-1", operation=operation)
|
||||
|
||||
check_access.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "expected_scope"),
|
||||
[("GET", "knowledge-spaces:read"), ("POST", "knowledge-spaces:write")],
|
||||
)
|
||||
def test_auth_signs_current_principals_and_declared_scope(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: KnowledgeFSMethod,
|
||||
expected_scope: str,
|
||||
) -> None:
|
||||
_set_config(monkeypatch)
|
||||
response = httpx.Response(200, content=b'{"items":[]}', headers={"Content-Type": "application/json"})
|
||||
request = MagicMock(return_value=response)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", request)
|
||||
|
||||
forward_knowledge_fs_request(
|
||||
account_id="account-1",
|
||||
method=method,
|
||||
path="knowledge-spaces",
|
||||
tenant_id="tenant-1",
|
||||
)
|
||||
|
||||
authorization = request.call_args.kwargs["headers"]["Authorization"]
|
||||
claims = jwt.decode(
|
||||
authorization.removeprefix("Bearer "),
|
||||
_JWT_SECRET,
|
||||
algorithms=["HS256"],
|
||||
audience="knowledge-fs",
|
||||
issuer="dify",
|
||||
)
|
||||
assert claims["dify_account_id"] == "dify-account:account-1"
|
||||
assert claims["sub"] == "dify-workspace:tenant-1"
|
||||
assert claims["tenant_id"] == "tenant-1"
|
||||
assert claims["scopes"] == [expected_scope]
|
||||
assert claims["caller_kind"] == "interactive"
|
||||
assert claims["exp"] - claims["iat"] == 60
|
||||
|
||||
|
||||
def test_buffered_response_rejects_non_empty_body_without_content_type(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_config(monkeypatch)
|
||||
response = httpx.Response(200, content=b"<html>unexpected</html>")
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", MagicMock(return_value=response))
|
||||
|
||||
with pytest.raises(KnowledgeFSTransportError, match="unsupported media type"):
|
||||
forward_knowledge_fs_request(
|
||||
account_id="account-dev", method="GET", path="knowledge-spaces", tenant_id="tenant-dev"
|
||||
)
|
||||
|
||||
assert response.is_closed
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "expected_exception"),
|
||||
[
|
||||
(
|
||||
httpx.ReadTimeout("timed out", request=httpx.Request("GET", "http://knowledge-fs.test")),
|
||||
KnowledgeFSTimeoutError,
|
||||
),
|
||||
(
|
||||
httpx.ConnectError("unavailable", request=httpx.Request("GET", "http://knowledge-fs.test")),
|
||||
KnowledgeFSTransportError,
|
||||
),
|
||||
(ssrf_proxy.ResponseTooLargeError("too large"), KnowledgeFSTransportError),
|
||||
(ToolSSRFError("blocked"), KnowledgeFSConfigurationError),
|
||||
],
|
||||
)
|
||||
def test_transport_failures_are_normalized(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
error: Exception,
|
||||
expected_exception: type[Exception],
|
||||
) -> None:
|
||||
_set_config(monkeypatch)
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", MagicMock(side_effect=error))
|
||||
|
||||
with pytest.raises(expected_exception):
|
||||
forward_knowledge_fs_request(
|
||||
account_id="account-dev", method="GET", path="knowledge-spaces", tenant_id="tenant-dev"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path"),
|
||||
[
|
||||
("GET", "openapi.json"),
|
||||
("GET", "knowledge-spaces/space-1"),
|
||||
("PATCH", "knowledge-spaces"),
|
||||
("POST", "queries"),
|
||||
("POST", "knowledge-spaces/space-1/documents"),
|
||||
],
|
||||
)
|
||||
def test_unregistered_route_is_rejected_before_external_io(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
) -> None:
|
||||
_set_config(monkeypatch)
|
||||
request = MagicMock()
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", request)
|
||||
|
||||
with pytest.raises(KnowledgeFSRouteNotAllowedError):
|
||||
forward_knowledge_fs_request(account_id="account-dev", method=method, path=path, tenant_id="tenant-dev")
|
||||
|
||||
request.assert_not_called()
|
||||
@ -11,3 +11,8 @@ PLUGIN_REMOTE_INSTALL_PORT=5003
|
||||
PLUGIN_MAX_PACKAGE_SIZE=52428800
|
||||
PLUGIN_DAEMON_TIMEOUT=600.0
|
||||
INNER_API_KEY_FOR_PLUGIN=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1
|
||||
KNOWLEDGE_FS_ENABLED=${KNOWLEDGE_FS_ENABLED:-false}
|
||||
KNOWLEDGE_FS_BASE_URL=
|
||||
KNOWLEDGE_FS_JWT_SECRET=
|
||||
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=300
|
||||
KNOWLEDGE_FS_TIMEOUT_SECONDS=10
|
||||
|
||||
Loading…
Reference in New Issue
Block a user