mirror of https://github.com/langgenius/dify.git
feat: add MCP tools
This commit is contained in:
parent
9d496ed3dc
commit
b7be5c8c82
|
|
@ -106,6 +106,9 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
|
|||
verify_option = kwargs.pop("ssl_verify", dify_config.HTTP_REQUEST_NODE_SSL_VERIFY)
|
||||
client = _get_ssrf_client(verify_option)
|
||||
|
||||
# Extract follow_redirects for client.send() - it's not a build_request parameter
|
||||
follow_redirects = kwargs.pop("follow_redirects", True)
|
||||
|
||||
# Preserve user-provided Host header
|
||||
# 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.
|
||||
|
|
@ -120,9 +123,9 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
|
|||
# the request API to explicitly set headers before sending
|
||||
headers = {k: v for k, v in headers.items() if k.lower() != "host"}
|
||||
if user_provided_host is not None:
|
||||
headers["host"] = user_provided_host
|
||||
kwargs["headers"] = headers
|
||||
response = client.request(method=method, url=url, **kwargs)
|
||||
request.headers["Host"] = user_provided_host
|
||||
|
||||
response = client.send(request, follow_redirects=follow_redirects)
|
||||
|
||||
# Check for SSRF protection by Squid proxy
|
||||
if response.status_code in (401, 403):
|
||||
|
|
|
|||
|
|
@ -107,3 +107,24 @@ def test_host_header_preservation_with_user_header(mock_get_client):
|
|||
response = make_request("GET", "http://example.com", headers={"Host": custom_host})
|
||||
|
||||
assert response.status_code == 200
|
||||
# Verify build_request was called
|
||||
mock_client.build_request.assert_called_once()
|
||||
# Verify the Host header was set on the request object
|
||||
assert mock_request.headers.get("Host") == custom_host
|
||||
mock_client.send.assert_called_once_with(mock_request, follow_redirects=True)
|
||||
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client")
|
||||
@pytest.mark.parametrize("host_key", ["host", "HOST"])
|
||||
def test_host_header_preservation_case_insensitive(mock_get_client, host_key):
|
||||
"""Test that Host header is preserved regardless of case."""
|
||||
mock_client = MagicMock()
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_client.build_request.return_value = mock_request
|
||||
mock_get_client.return_value = mock_client
|
||||
response = make_request("GET", "http://example.com", headers={host_key: "api.example.com"})
|
||||
assert mock_request.headers.get("Host") == "api.example.com"
|
||||
|
|
|
|||
Loading…
Reference in New Issue