|
|
|
@@ -1,6 +1,10 @@ |
|
|
|
import json |
|
|
|
from collections.abc import Generator |
|
|
|
from contextlib import AbstractContextManager |
|
|
|
|
|
|
|
import httpx |
|
|
|
import httpx_sse |
|
|
|
from httpx_sse import connect_sse |
|
|
|
|
|
|
|
from configs import dify_config |
|
|
|
from core.mcp.types import ErrorData, JSONRPCError |
|
|
|
@@ -55,20 +59,42 @@ def create_ssrf_proxy_mcp_http_client( |
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
def ssrf_proxy_sse_connect(url, **kwargs): |
|
|
|
def ssrf_proxy_sse_connect(url: str, **kwargs) -> AbstractContextManager[httpx_sse.EventSource]: |
|
|
|
"""Connect to SSE endpoint with SSRF proxy protection. |
|
|
|
|
|
|
|
This function creates an SSE connection using the configured proxy settings |
|
|
|
to prevent SSRF attacks when connecting to external endpoints. |
|
|
|
to prevent SSRF attacks when connecting to external endpoints. It returns |
|
|
|
a context manager that yields an EventSource object for SSE streaming. |
|
|
|
|
|
|
|
The function handles HTTP client creation and cleanup automatically, but |
|
|
|
also accepts a pre-configured client via kwargs. |
|
|
|
|
|
|
|
Args: |
|
|
|
url: The SSE endpoint URL |
|
|
|
**kwargs: Additional arguments passed to the SSE connection |
|
|
|
url (str): The SSE endpoint URL to connect to |
|
|
|
**kwargs: Additional arguments passed to the SSE connection, including: |
|
|
|
- client (httpx.Client, optional): Pre-configured HTTP client. |
|
|
|
If not provided, one will be created with SSRF protection. |
|
|
|
- method (str, optional): HTTP method to use, defaults to "GET" |
|
|
|
- headers (dict, optional): HTTP headers to include in the request |
|
|
|
- timeout (httpx.Timeout, optional): Timeout configuration for the connection |
|
|
|
|
|
|
|
Returns: |
|
|
|
EventSource object for SSE streaming |
|
|
|
AbstractContextManager[httpx_sse.EventSource]: A context manager that yields an EventSource |
|
|
|
object for SSE streaming. The EventSource provides access to server-sent events. |
|
|
|
|
|
|
|
Example: |
|
|
|
```python |
|
|
|
with ssrf_proxy_sse_connect(url, headers=headers) as event_source: |
|
|
|
for sse in event_source.iter_sse(): |
|
|
|
print(sse.event, sse.data) |
|
|
|
``` |
|
|
|
|
|
|
|
Note: |
|
|
|
If a client is not provided in kwargs, one will be automatically created |
|
|
|
with SSRF protection based on the application's configuration. If an |
|
|
|
exception occurs during connection, any automatically created client |
|
|
|
will be cleaned up automatically. |
|
|
|
""" |
|
|
|
from httpx_sse import connect_sse |
|
|
|
|
|
|
|
# Extract client if provided, otherwise create one |
|
|
|
client = kwargs.pop("client", None) |
|
|
|
@@ -101,7 +127,9 @@ def ssrf_proxy_sse_connect(url, **kwargs): |
|
|
|
raise |
|
|
|
|
|
|
|
|
|
|
|
def create_mcp_error_response(request_id: int | str | None, code: int, message: str, data=None): |
|
|
|
def create_mcp_error_response( |
|
|
|
request_id: int | str | None, code: int, message: str, data=None |
|
|
|
) -> Generator[bytes, None, None]: |
|
|
|
"""Create MCP error response""" |
|
|
|
error_data = ErrorData(code=code, message=message, data=data) |
|
|
|
json_response = JSONRPCError( |