pengjunjie/langflow
1
0
Fork 1
mirror of https://github.com/langflow-ai/langflow.git synced 2026-07-23 15:09:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: protect URL components from SSRF (#13643)

Browse Source 
* fix: protect URL components from SSRF

* [autofix.ci] apply automated fixes

* fix: address ci failures

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Eric Hare
2026-06-15 10:37:47 -07:00
committed by GitHub
parent 02e884552c
commit a074efa9b7
24 changed files with 769 additions and 179 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.secrets.baseline
View File

@ -3565,7 +3565,7 @@
"filename": "src/backend/tests/unit/components/languagemodels/test_xai.py",
"hashed_secret": "3bee216ebc256d692260fc3adc765050508fef5e",
"is_verified": false,
"line_number": 150,
"line_number": 159,
"is_secret": false
}
],
@ -9287,5 +9287,5 @@
}
]
},
"generated_at": "2026-06-10T08:36:23Z"
"generated_at": "2026-06-12T22:46:30Z"
}

8
src/backend/base/langflow/initial_setup/starter_projects/Structured Data Analysis Agent.json
View File

File diff suppressed because one or more lines are too long

5
src/backend/tests/integration/components/languagemodels/test_chatollama_integration.py
View File

@ -7,6 +7,11 @@ from lfx.schema.dataframe import DataFrame
from lfx.schema.message import Message
@pytest.fixture(autouse=True)
def allow_local_ollama(monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_ALLOWED_HOSTS", "localhost")
@pytest.mark.integration
class TestChatOllamaIntegration:
"""Integration tests for ChatOllama structured output flow."""

4
src/backend/tests/unit/components/embeddings/test_ollama_embeddings_component.py
View File

@ -919,3 +919,7 @@ class TestOllamaEmbeddingsComponent(ComponentTestBaseWithoutClient):
component = component_class()
output_names = [out.name for out in component.outputs]
assert "embeddings" in output_names
@pytest.fixture(autouse=True)
def disable_ssrf_protection(self, monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_PROTECTION_ENABLED", "false")

4
src/backend/tests/unit/components/languagemodels/test_chatollama_component.py
View File

@ -1260,3 +1260,7 @@ class TestChatOllamaComponent(ComponentTestBaseWithoutClient):
assert "format" in call_args, "format should be in call when enabled"
assert call_args["format"] == "json"
assert model == mock_instance
@pytest.fixture(autouse=True)
def disable_ssrf_protection(self, monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_PROTECTION_ENABLED", "false")

9
src/backend/tests/unit/components/languagemodels/test_deepseek.py
View File

@ -82,8 +82,8 @@ def test_deepseek_build_model(mock_chat_openai, temperature, max_tokens):
def test_deepseek_get_models(mocker):
component = DeepSeekModelComponent()
# Mock requests.get
mock_get = mocker.patch("requests.get")
# Mock SSRF-safe httpx helper
mock_get = mocker.patch("lfx.components.deepseek.deepseek.ssrf_safe_httpx_get")
mock_response = MagicMock()
mock_response.json.return_value = {"data": [{"id": "deepseek-chat"}, {"id": "deepseek-coder"}]}
mock_get.return_value = mock_response
@ -110,3 +110,8 @@ def test_deepseek_error_handling(mock_chat_openai):
with pytest.raises(Exception, match="Invalid API key"):
component.build_model()
@pytest.fixture(autouse=True)
def disable_ssrf_protection(monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_PROTECTION_ENABLED", "false")

22
src/backend/tests/unit/components/languagemodels/test_litellm_proxy.py
View File

@ -72,7 +72,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(),
)
mock_chat_openai = mocker.patch(
@ -98,7 +98,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(),
)
mock_chat_openai = mocker.patch(
@ -116,7 +116,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(),
)
mock_chat_openai = mocker.patch(
@ -133,7 +133,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
def test_validate_proxy_connection_success(self, component_class, default_kwargs, mocker):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(),
)
# Should not raise
@ -142,7 +142,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
def test_validate_proxy_connection_auth_failure(self, component_class, default_kwargs, mocker):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(status_code=401),
)
with pytest.raises(ValueError, match="Authentication failed"):
@ -152,7 +152,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
default_kwargs["model_name"] = "invalid-model-name"
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(models=[{"id": "gpt-4o"}]),
)
with pytest.raises(ValueError, match=r"invalid-model-name.*not found"):
@ -161,7 +161,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
def test_validate_proxy_connection_connect_error(self, component_class, default_kwargs, mocker):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
side_effect=httpx.ConnectError("Connection refused"),
)
with pytest.raises(ValueError, match="Could not connect"):
@ -170,7 +170,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
def test_validate_proxy_connection_timeout(self, component_class, default_kwargs, mocker):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
side_effect=httpx.TimeoutException("Timed out"),
)
with pytest.raises(ValueError, match="timed out"):
@ -179,7 +179,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
def test_validate_proxy_connection_empty_models_list(self, component_class, default_kwargs, mocker):
component = component_class(**default_kwargs)
mocker.patch(
"lfx.components.litellm.litellm_proxy.httpx.get",
"lfx.components.litellm.litellm_proxy.ssrf_safe_httpx_get",
return_value=_mock_models_response(models=[]),
)
# Empty models list should not raise (proxy may not report models)
@ -238,3 +238,7 @@ class TestLiteLLMProxyComponent(ComponentTestBaseWithoutClient):
with patch.dict("sys.modules", {"openai": None}):
message = component._get_exception_message(Exception("test"))
assert message is None
@pytest.fixture(autouse=True)
def disable_ssrf_protection(self, monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_PROTECTION_ENABLED", "false")

22
src/backend/tests/unit/components/languagemodels/test_xai.py
View File

@ -101,8 +101,15 @@ class TestXAIComponent(ComponentTestBaseWithoutClient):
component.base_url = "https://api.x.ai/v1"
component.seed = 1
http_client = MagicMock()
http_async_client = MagicMock()
mock_client_kwargs = mocker.patch(
"lfx.components.xai.xai.ssrf_protected_openai_clients_for_url",
return_value={"http_client": http_client, "http_async_client": http_async_client},
)
mock_chat_openai = mocker.patch("lfx.components.xai.xai.ChatOpenAI", return_value=MagicMock())
model = component.build_model()
mock_client_kwargs.assert_called_once_with("https://api.x.ai/v1")
mock_chat_openai.assert_called_once_with(
max_tokens=100,
model_kwargs={},
@ -111,12 +118,14 @@ class TestXAIComponent(ComponentTestBaseWithoutClient):
api_key="test-key",
temperature=0.7,
seed=1,
http_client=http_client,
http_async_client=http_async_client,
)
assert model == mock_chat_openai.return_value
def test_get_models(self):
component = XAIModelComponent()
with patch("requests.get") as mock_get:
with patch("lfx.components.xai.xai.ssrf_safe_httpx_get") as mock_get:
mock_response = MagicMock()
mock_response.json.return_value = {
"models": [
@ -191,8 +200,11 @@ class TestXAIComponent(ComponentTestBaseWithoutClient):
component = XAIModelComponent()
build_config = {"model_name": {"options": []}}
updated_config = component.update_build_config(build_config, "test-key", "api_key")
assert "model_name" in updated_config
with patch.object(component, "get_models", return_value=["grok-2-latest"]) as mock_get_models:
updated_config = component.update_build_config(build_config, "test-key", "api_key")
assert "model_name" in updated_config
updated_config = component.update_build_config(build_config, "grok-2-latest", "model_name")
assert "model_name" in updated_config
updated_config = component.update_build_config(build_config, "grok-2-latest", "model_name")
assert "model_name" in updated_config
assert mock_get_models.call_count == 2

183
src/backend/tests/unit/components/test_ssrf_guarded_url_components.py Normal file
View File

@ -0,0 +1,183 @@
from unittest.mock import patch
import pytest
from lfx.components.deepseek.deepseek import DEEPSEEK_MODELS, DeepSeekModelComponent
from lfx.components.glean.glean_search_api import GleanAPIWrapper
from lfx.components.homeassistant.home_assistant_control import HomeAssistantControl
from lfx.components.homeassistant.list_home_assistant_states import ListHomeAssistantStates
from lfx.components.huggingface.huggingface_inference_api import HuggingFaceInferenceAPIEmbeddingsComponent
from lfx.components.litellm.litellm_proxy import LiteLLMProxyComponent
from lfx.components.lmstudio.lmstudioembeddings import LMStudioEmbeddingsComponent
from lfx.components.lmstudio.lmstudiomodel import LMStudioModelComponent
from lfx.components.ollama.ollama import ChatOllamaComponent
from lfx.components.ollama.ollama_embeddings import OllamaEmbeddingsComponent
from lfx.components.xai.xai import XAI_DEFAULT_MODELS, XAIModelComponent
from lfx.utils.ssrf_protection import SSRFProtectionError
BLOCKED_URL = "http://169.254.169.254/latest/meta-data"
@pytest.fixture(autouse=True)
def enable_ssrf_protection(monkeypatch):
monkeypatch.setenv("LANGFLOW_SSRF_PROTECTION_ENABLED", "true")
monkeypatch.delenv("LANGFLOW_SSRF_ALLOWED_HOSTS", raising=False)
@pytest.mark.asyncio
async def test_lmstudio_model_update_blocks_metadata_url_before_httpx():
component = LMStudioModelComponent()
build_config = {"base_url": {"load_from_db": False, "value": BLOCKED_URL}, "model_name": {"options": []}}
with patch("httpx.AsyncClient.get") as mock_get, pytest.raises(ValueError, match="SSRF Protection"):
await component.update_build_config(build_config, None, "model_name")
mock_get.assert_not_called()
def test_lmstudio_model_build_blocks_metadata_url_before_openai_client():
component = LMStudioModelComponent(base_url=BLOCKED_URL, model_name="model", api_key="test")
with (
patch("lfx.components.lmstudio.lmstudiomodel.ChatOpenAI") as mock_chat_openai,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_model()
mock_chat_openai.assert_not_called()
def test_lmstudio_embeddings_build_blocks_metadata_url_before_sdk_client():
component = LMStudioEmbeddingsComponent(base_url=BLOCKED_URL, model="model", api_key="test")
with (
patch("lfx.components.lmstudio.lmstudioembeddings.NVIDIAEmbeddings", create=True) as mock_embeddings,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_embeddings()
mock_embeddings.assert_not_called()
def test_home_assistant_list_states_blocks_metadata_url_before_httpx():
component = ListHomeAssistantStates()
with patch("httpx.Client.get") as mock_get:
result = component._list_states("token", BLOCKED_URL)
assert "SSRF Protection" in result
mock_get.assert_not_called()
def test_home_assistant_control_blocks_metadata_url_before_httpx():
component = HomeAssistantControl()
with patch("httpx.Client.post") as mock_post:
result = component._control_device("token", BLOCKED_URL, "turn_on", "switch.test")
assert "SSRF Protection" in result
mock_post.assert_not_called()
def test_deepseek_model_fetch_blocks_metadata_url_before_httpx():
component = DeepSeekModelComponent(api_base=BLOCKED_URL, api_key="test")
with patch("httpx.Client.get") as mock_get:
models = component.get_models()
assert models == DEEPSEEK_MODELS
assert "SSRF Protection" in component.status
mock_get.assert_not_called()
def test_deepseek_build_blocks_metadata_url_before_openai_client():
component = DeepSeekModelComponent(api_base=BLOCKED_URL, api_key="test")
with patch("langchain_openai.ChatOpenAI") as mock_chat_openai, pytest.raises(ValueError, match="SSRF Protection"):
component.build_model()
mock_chat_openai.assert_not_called()
def test_xai_model_fetch_blocks_metadata_url_before_httpx():
component = XAIModelComponent(base_url=BLOCKED_URL, api_key="test")
with patch("httpx.Client.get") as mock_get:
models = component.get_models()
assert models == XAI_DEFAULT_MODELS
assert "SSRF Protection" in component.status
mock_get.assert_not_called()
def test_xai_build_blocks_metadata_url_before_openai_client():
component = XAIModelComponent(base_url=BLOCKED_URL, api_key="test")
with (
patch("lfx.components.xai.xai.ChatOpenAI") as mock_chat_openai,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_model()
mock_chat_openai.assert_not_called()
def test_glean_blocks_metadata_url_before_httpx_post():
wrapper = GleanAPIWrapper(glean_api_url=BLOCKED_URL, glean_access_token="test-access-token") # noqa: S106
with patch("httpx.Client.post") as mock_post, pytest.raises(SSRFProtectionError):
wrapper._search_api_results("query")
mock_post.assert_not_called()
def test_huggingface_build_blocks_metadata_url_before_sdk_client():
component = HuggingFaceInferenceAPIEmbeddingsComponent(
inference_endpoint=BLOCKED_URL,
model_name="model",
)
with (
patch.object(component, "create_huggingface_embeddings") as mock_create,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_embeddings()
mock_create.assert_not_called()
def test_ollama_embeddings_build_blocks_metadata_url_before_sdk_client():
component = OllamaEmbeddingsComponent(base_url=BLOCKED_URL, model_name="model")
with (
patch("lfx.components.ollama.ollama_embeddings.OllamaEmbeddings") as mock_embeddings,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_embeddings()
mock_embeddings.assert_not_called()
def test_ollama_build_blocks_metadata_url_before_sdk_client():
component = ChatOllamaComponent(base_url=BLOCKED_URL, model_name="model", mirostat="Disabled")
with (
patch("lfx.components.ollama.ollama.ChatOllama") as mock_chat_ollama,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_model()
mock_chat_ollama.assert_not_called()
def test_litellm_build_blocks_metadata_url_before_httpx_and_openai_client():
component = LiteLLMProxyComponent(api_base=BLOCKED_URL, api_key="test", model_name="model")
with (
patch("httpx.Client.get") as mock_get,
patch("lfx.components.litellm.litellm_proxy.ChatOpenAI") as mock_chat_openai,
pytest.raises(ValueError, match="SSRF Protection"),
):
component.build_model()
mock_get.assert_not_called()
mock_chat_openai.assert_not_called()

82
src/lfx/src/lfx/_assets/component_index.json
View File

File diff suppressed because one or more lines are too long

14
src/lfx/src/lfx/components/deepseek/deepseek.py
View File

@ -1,4 +1,4 @@
import requests
import httpx
from pydantic.v1 import SecretStr
from typing_extensions import override
@ -6,6 +6,8 @@ from lfx.base.models.model import LCModelComponent
from lfx.field_typing import LanguageModel
from lfx.field_typing.range_spec import RangeSpec
from lfx.inputs.inputs import BoolInput, DictInput, DropdownInput, IntInput, SecretStrInput, SliderInput, StrInput
from lfx.utils.ssrf_httpx import ssrf_protected_openai_clients_for_url, ssrf_safe_httpx_get
from lfx.utils.ssrf_protection import SSRFProtectionError
DEEPSEEK_MODELS = ["deepseek-chat"]
@ -83,11 +85,14 @@ class DeepSeekModelComponent(LCModelComponent):
headers = {"Authorization": f"Bearer {self.api_key}", "Accept": "application/json"}
try:
response = requests.get(url, headers=headers, timeout=10)
response = ssrf_safe_httpx_get(url, headers=headers, timeout=10)
response.raise_for_status()
model_list = response.json()
return [model["id"] for model in model_list.get("data", [])]
except requests.RequestException as e:
except SSRFProtectionError as e:
self.status = f"SSRF Protection: {e}"
return DEEPSEEK_MODELS
except httpx.HTTPError as e:
self.status = f"Error fetching models: {e}"
return DEEPSEEK_MODELS
@ -106,6 +111,8 @@ class DeepSeekModelComponent(LCModelComponent):
raise ImportError(msg) from e
api_key = SecretStr(self.api_key).get_secret_value() if self.api_key else None
ssrf_client_kwargs = ssrf_protected_openai_clients_for_url(self.api_base)
output = ChatOpenAI(
model=self.model_name,
temperature=self.temperature if self.temperature is not None else 0.1,
@ -115,6 +122,7 @@ class DeepSeekModelComponent(LCModelComponent):
api_key=api_key,
streaming=self.stream if hasattr(self, "stream") else False,
seed=self.seed,
**ssrf_client_kwargs,
)
if self.json_mode:

4
src/lfx/src/lfx/components/glean/glean_search_api.py
View File

@ -2,7 +2,6 @@ import json
from typing import Any
from urllib.parse import urljoin
import httpx
from langchain_core.tools import StructuredTool, ToolException
from pydantic import BaseModel
from pydantic.v1 import Field
@ -13,6 +12,7 @@ from lfx.inputs.inputs import IntInput, MultilineInput, NestedDictInput, SecretS
from lfx.io import Output
from lfx.schema.data import Data
from lfx.schema.dataframe import DataFrame
from lfx.utils.ssrf_httpx import ssrf_safe_httpx_post
class GleanSearchAPISchema(BaseModel):
@ -82,7 +82,7 @@ class GleanAPIWrapper(BaseModel):
def _search_api_results(self, query: str, **kwargs: Any) -> list[dict[str, Any]]:
request_details = self._prepare_request(query, **kwargs)
response = httpx.post(
response = ssrf_safe_httpx_post(
request_details["url"],
json=request_details["payload"],
headers=request_details["headers"],

10
src/lfx/src/lfx/components/homeassistant/home_assistant_control.py
View File

@ -1,7 +1,7 @@
import json
from typing import Any
import requests
import httpx
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field
@ -9,6 +9,8 @@ from lfx.base.langchain_utilities.model import LCToolComponent
from lfx.field_typing import Tool
from lfx.inputs.inputs import SecretStrInput, StrInput
from lfx.schema.data import Data
from lfx.utils.ssrf_httpx import ssrf_safe_httpx_post
from lfx.utils.ssrf_protection import SSRFProtectionError
class HomeAssistantControl(LCToolComponent):
@ -132,11 +134,13 @@ class HomeAssistantControl(LCToolComponent):
}
payload = {"entity_id": entity_id}
response = requests.post(url, headers=headers, json=payload, timeout=10)
response = ssrf_safe_httpx_post(url, headers=headers, json=payload, timeout=10)
response.raise_for_status()
return response.json() # HA response JSON on success
except requests.exceptions.RequestException as e:
except SSRFProtectionError as e:
return f"Error: SSRF Protection: {e}"
except httpx.HTTPError as e:
return f"Error: Failed to call service. {e}"
except Exception as e: # noqa: BLE001
return f"An unexpected error occurred: {e}"

10
src/lfx/src/lfx/components/homeassistant/list_home_assistant_states.py
View File

@ -1,7 +1,7 @@
import json
from typing import Any
import requests
import httpx
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field
@ -9,6 +9,8 @@ from lfx.base.langchain_utilities.model import LCToolComponent
from lfx.field_typing import Tool
from lfx.inputs.inputs import SecretStrInput, StrInput
from lfx.schema.data import Data
from lfx.utils.ssrf_httpx import ssrf_safe_httpx_get
from lfx.utils.ssrf_protection import SSRFProtectionError
class ListHomeAssistantStates(LCToolComponent):
@ -103,14 +105,16 @@ class ListHomeAssistantStates(LCToolComponent):
"Content-Type": "application/json",
}
url = f"{base_url}/api/states"
response = requests.get(url, headers=headers, timeout=10)
response = ssrf_safe_httpx_get(url, headers=headers, timeout=10)
response.raise_for_status()
all_states = response.json()
if filter_domain:
return [st for st in all_states if st.get("entity_id", "").startswith(f"{filter_domain}.")]
except requests.exceptions.RequestException as e:
except SSRFProtectionError as e:
return f"Error: SSRF Protection: {e}"
except httpx.HTTPError as e:
return f"Error: Failed to fetch states. {e}"
except (ValueError, TypeError) as e:
return f"Error processing response: {e}"

10
src/lfx/src/lfx/components/huggingface/huggingface_inference_api.py
View File

@ -1,6 +1,6 @@
from urllib.parse import urlparse
import requests
import httpx
from langchain_community.embeddings.huggingface import HuggingFaceInferenceAPIEmbeddings
# Next update: use langchain_huggingface
@ -10,6 +10,7 @@ from tenacity import retry, stop_after_attempt, wait_fixed
from lfx.base.embeddings.model import LCEmbeddingsModel
from lfx.field_typing import Embeddings
from lfx.io import MessageTextInput, Output, SecretStrInput
from lfx.utils.ssrf_httpx import ssrf_safe_httpx_get, validate_url_for_ssrf_or_raise
class HuggingFaceInferenceAPIEmbeddingsComponent(LCEmbeddingsModel):
@ -57,15 +58,15 @@ class HuggingFaceInferenceAPIEmbeddingsComponent(LCEmbeddingsModel):
raise ValueError(msg)
try:
response = requests.get(f"{inference_endpoint}/health", timeout=5)
except requests.RequestException as e:
response = ssrf_safe_httpx_get(f"{inference_endpoint}/health", timeout=5)
except httpx.HTTPError as e:
msg = (
f"Inference endpoint '{inference_endpoint}' is not responding. "
"Please ensure the URL is correct and the service is running."
)
raise ValueError(msg) from e
if response.status_code != requests.codes.ok:
if response.status_code != httpx.codes.OK:
msg = f"Hugging Face health check failed: {response.status_code}"
raise ValueError(msg)
# returning True to solve linting error
@ -84,6 +85,7 @@ class HuggingFaceInferenceAPIEmbeddingsComponent(LCEmbeddingsModel):
def build_embeddings(self) -> Embeddings:
api_url = self.get_api_url()
validate_url_for_ssrf_or_raise(api_url)
is_local_url = (
api_url.startswith(("http://localhost", "http://127.0.0.1", "http://0.0.0.0", "http://docker"))

9
src/lfx/src/lfx/components/litellm/litellm_proxy.py
View File

@ -6,6 +6,8 @@ from lfx.field_typing import LanguageModel
from lfx.field_typing.range_spec import RangeSpec
from lfx.inputs.inputs import IntInput, SecretStrInput, SliderInput, StrInput
from lfx.utils.secrets import secret_value_to_str
from lfx.utils.ssrf_httpx import ssrf_protected_openai_clients_for_url, ssrf_safe_httpx_get
from lfx.utils.ssrf_protection import SSRFProtectionError
class LiteLLMProxyComponent(LCModelComponent):
@ -72,6 +74,7 @@ class LiteLLMProxyComponent(LCModelComponent):
def build_model(self) -> LanguageModel:
"""Build the LiteLLM proxy model."""
api_key = secret_value_to_str(self.api_key) or ""
ssrf_client_kwargs = ssrf_protected_openai_clients_for_url(self.api_base)
self._validate_proxy_connection(api_key)
@ -84,6 +87,7 @@ class LiteLLMProxyComponent(LCModelComponent):
timeout=self.timeout,
max_retries=self.max_retries,
streaming=self.stream,
**ssrf_client_kwargs,
)
def _validate_proxy_connection(self, api_key: str) -> None:
@ -92,11 +96,14 @@ class LiteLLMProxyComponent(LCModelComponent):
models_url = f"{base_url}/models"
try:
response = httpx.get(
response = ssrf_safe_httpx_get(
models_url,
headers={"Authorization": f"Bearer {api_key}"},
timeout=10,
)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except httpx.ConnectError as e:
msg = (
f"Could not connect to LiteLLM Proxy at {base_url}. Verify the URL is correct and the proxy is running."

18
src/lfx/src/lfx/components/lmstudio/lmstudioembeddings.py
View File

@ -1,12 +1,12 @@
from typing import Any
from urllib.parse import urljoin
import httpx
from lfx.base.embeddings.model import LCEmbeddingsModel
from lfx.field_typing import Embeddings
from lfx.inputs.inputs import DropdownInput, SecretStrInput
from lfx.io import FloatInput, MessageTextInput
from lfx.utils.ssrf_httpx import ssrf_safe_async_get, validate_url_for_ssrf_or_raise
from lfx.utils.ssrf_protection import SSRFProtectionError
class LMStudioEmbeddingsComponent(LCEmbeddingsModel):
@ -31,12 +31,14 @@ class LMStudioEmbeddingsComponent(LCEmbeddingsModel):
async def get_model(base_url_value: str) -> list[str]:
try:
url = urljoin(base_url_value, "/v1/models")
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
data = response.json()
response = await ssrf_safe_async_get(url)
response.raise_for_status()
data = response.json()
return [model["id"] for model in data.get("data", [])]
return [model["id"] for model in data.get("data", [])]
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except Exception as e:
msg = "Could not retrieve models. Please, make sure the LM Studio server is running."
raise ValueError(msg) from e
@ -71,6 +73,8 @@ class LMStudioEmbeddingsComponent(LCEmbeddingsModel):
]
def build_embeddings(self) -> Embeddings:
validate_url_for_ssrf_or_raise(self.base_url)
try:
from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
except ImportError as e:

25
src/lfx/src/lfx/components/lmstudio/lmstudiomodel.py
View File

@ -8,6 +8,8 @@ from lfx.base.models.model import LCModelComponent
from lfx.field_typing import LanguageModel
from lfx.field_typing.range_spec import RangeSpec
from lfx.inputs.inputs import DictInput, DropdownInput, FloatInput, IntInput, SecretStrInput, StrInput
from lfx.utils.ssrf_httpx import ssrf_protected_openai_clients_for_url, ssrf_safe_async_get
from lfx.utils.ssrf_protection import SSRFProtectionError
class LMStudioModelComponent(LCModelComponent):
@ -24,9 +26,11 @@ class LMStudioModelComponent(LCModelComponent):
if base_url_load_from_db:
base_url_value = await self.get_variables(base_url_value, field_name)
try:
async with httpx.AsyncClient() as client:
response = await client.get(urljoin(base_url_value, "/v1/models"), timeout=2.0)
response.raise_for_status()
response = await ssrf_safe_async_get(urljoin(base_url_value, "/v1/models"), timeout=2.0)
response.raise_for_status()
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except httpx.HTTPError:
msg = "Could not access the default LM Studio URL. Please, specify the 'Base URL' field."
self.log(msg)
@ -39,12 +43,14 @@ class LMStudioModelComponent(LCModelComponent):
async def get_model(base_url_value: str) -> list[str]:
try:
url = urljoin(base_url_value, "/v1/models")
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
data = response.json()
response = await ssrf_safe_async_get(url)
response.raise_for_status()
data = response.json()
return [model["id"] for model in data.get("data", [])]
return [model["id"] for model in data.get("data", [])]
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except Exception as e:
msg = "Could not retrieve models. Please, make sure the LM Studio server is running."
raise ValueError(msg) from e
@ -103,6 +109,8 @@ class LMStudioModelComponent(LCModelComponent):
base_url = self.base_url or "http://localhost:1234/v1"
seed = self.seed
ssrf_client_kwargs = ssrf_protected_openai_clients_for_url(base_url)
return ChatOpenAI(
max_tokens=max_tokens or None,
model_kwargs=model_kwargs,
@ -111,6 +119,7 @@ class LMStudioModelComponent(LCModelComponent):
api_key=lmstudio_api_key,
temperature=temperature if temperature is not None else 0.1,
seed=seed,
**ssrf_client_kwargs,
)
def _get_exception_message(self, e: Exception):

102
src/lfx/src/lfx/components/ollama/ollama.py
View File

@ -28,6 +28,12 @@ from lfx.log.logger import logger
from lfx.schema.data import Data
from lfx.schema.dataframe import DataFrame
from lfx.schema.table import EditMode
from lfx.utils.ssrf_httpx import (
ssrf_protected_httpx_client_kwargs_for_url,
ssrf_safe_async_get,
ssrf_safe_async_post,
)
from lfx.utils.ssrf_protection import SSRFProtectionError
from lfx.utils.util import transform_localhost_url
HTTP_STATUS_OK = 200
@ -263,6 +269,8 @@ class ChatOllamaComponent(LCModelComponent):
"Learn more at https://docs.ollama.com/openai#openai-compatibility"
)
sync_client_kwargs, async_client_kwargs = ssrf_protected_httpx_client_kwargs_for_url(transformed_base_url)
try:
output_format = self._parse_format_field(self.format) if self.enable_structured_output else None
except Exception as e:
@ -297,6 +305,10 @@ class ChatOllamaComponent(LCModelComponent):
headers = self.headers
if headers is not None:
llm_params["client_kwargs"] = {"headers": headers}
if sync_client_kwargs:
llm_params["sync_client_kwargs"] = sync_client_kwargs
if async_client_kwargs:
llm_params["async_client_kwargs"] = async_client_kwargs
# Remove parameters with None values
llm_params = {k: v for k, v in llm_params.items() if v is not None}
@ -314,19 +326,21 @@ class ChatOllamaComponent(LCModelComponent):
async def is_valid_ollama_url(self, url: str) -> bool:
try:
async with httpx.AsyncClient() as client:
url = transform_localhost_url(url)
if not url:
return False
# Strip /v1 suffix if present, as Ollama API endpoints are at root level
url = url.rstrip("/").removesuffix("/v1")
if not url.endswith("/"):
url = url + "/"
return (
await client.get(url=urljoin(url, "api/tags"), headers=self.headers)
).status_code == HTTP_STATUS_OK
url = transform_localhost_url(url)
if not url:
return False
# Strip /v1 suffix if present, as Ollama API endpoints are at root level
url = url.rstrip("/").removesuffix("/v1")
if not url.endswith("/"):
url = url + "/"
response = await ssrf_safe_async_get(urljoin(url, "api/tags"), headers=self.headers)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except httpx.RequestError:
return False
else:
return response.status_code == HTTP_STATUS_OK
async def update_build_config(self, build_config: dict, field_value: Any, field_name: str | None = None):
if field_name == "enable_structured_output": # bind enable_structured_output boolean to format show value
@ -409,44 +423,46 @@ class ChatOllamaComponent(LCModelComponent):
# Ollama REST API to return model capabilities
show_url = urljoin(base_url, "api/show")
async with httpx.AsyncClient() as client:
headers = self.headers
# Fetch available models
tags_response = await client.get(url=tags_url, headers=headers)
tags_response.raise_for_status()
models = tags_response.json()
if asyncio.iscoroutine(models):
models = await models
await logger.adebug(f"Available models: {models}")
headers = self.headers
# Fetch available models
tags_response = await ssrf_safe_async_get(tags_url, headers=headers)
tags_response.raise_for_status()
models = tags_response.json()
if asyncio.iscoroutine(models):
models = await models
await logger.adebug(f"Available models: {models}")
# Filter models that are NOT embedding models
model_ids = []
for model in models[self.JSON_MODELS_KEY]:
model_name = model[self.JSON_NAME_KEY]
await logger.adebug(f"Checking model: {model_name}")
# Filter models that are NOT embedding models
model_ids = []
for model in models[self.JSON_MODELS_KEY]:
model_name = model[self.JSON_NAME_KEY]
await logger.adebug(f"Checking model: {model_name}")
payload = {"model": model_name}
show_response = await client.post(url=show_url, json=payload, headers=headers)
show_response.raise_for_status()
json_data = show_response.json()
if asyncio.iscoroutine(json_data):
json_data = await json_data
payload = {"model": model_name}
show_response = await ssrf_safe_async_post(show_url, json=payload, headers=headers)
show_response.raise_for_status()
json_data = show_response.json()
if asyncio.iscoroutine(json_data):
json_data = await json_data
capabilities = json_data.get(self.JSON_CAPABILITIES_KEY)
await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
capabilities = json_data.get(self.JSON_CAPABILITIES_KEY)
await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
# If capabilities not provided, assume it's a completion model (backwards compatibility
# with older Ollama versions that don't return capabilities from /api/show)
if capabilities is None:
if not tool_model_enabled:
model_ids.append(model_name)
# If tool_model_enabled is True but no capabilities info, skip the model
# since we can't verify tool support
elif self.DESIRED_CAPABILITY in capabilities and (
not tool_model_enabled or self.TOOL_CALLING_CAPABILITY in capabilities
):
# If capabilities not provided, assume it's a completion model (backwards compatibility
# with older Ollama versions that don't return capabilities from /api/show)
if capabilities is None:
if not tool_model_enabled:
model_ids.append(model_name)
# If tool_model_enabled is True but no capabilities info, skip the model
# since we can't verify tool support
elif self.DESIRED_CAPABILITY in capabilities and (
not tool_model_enabled or self.TOOL_CALLING_CAPABILITY in capabilities
):
model_ids.append(model_name)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except (httpx.RequestError, ValueError) as e:
msg = "Could not get model names from Ollama."
raise ValueError(msg) from e

86
src/lfx/src/lfx/components/ollama/ollama_embeddings.py
View File

@ -9,6 +9,12 @@ from lfx.base.models.model import LCModelComponent
from lfx.field_typing import Embeddings
from lfx.io import DropdownInput, Output, SecretStrInput, StrInput
from lfx.log.logger import logger
from lfx.utils.ssrf_httpx import (
ssrf_protected_httpx_client_kwargs_for_url,
ssrf_safe_async_get,
ssrf_safe_async_post,
)
from lfx.utils.ssrf_protection import SSRFProtectionError
from lfx.utils.util import transform_localhost_url
HTTP_STATUS_OK = 200
@ -81,11 +87,17 @@ class OllamaEmbeddingsComponent(LCModelComponent):
"Learn more at https://docs.ollama.com/openai#openai-compatibility"
)
sync_client_kwargs, async_client_kwargs = ssrf_protected_httpx_client_kwargs_for_url(transformed_base_url)
llm_params = {
"model": self.model_name,
"base_url": transformed_base_url,
}
if sync_client_kwargs:
llm_params["sync_client_kwargs"] = sync_client_kwargs
if async_client_kwargs:
llm_params["async_client_kwargs"] = async_client_kwargs
if self.headers:
llm_params["client_kwargs"] = {"headers": self.headers}
@ -133,35 +145,37 @@ class OllamaEmbeddingsComponent(LCModelComponent):
# Ollama REST API to return model capabilities
show_url = urljoin(base_url, "api/show")
async with httpx.AsyncClient() as client:
headers = self.headers
# Fetch available models
tags_response = await client.get(url=tags_url, headers=headers)
tags_response.raise_for_status()
models = tags_response.json()
if asyncio.iscoroutine(models):
models = await models
await logger.adebug(f"Available models: {models}")
headers = self.headers
# Fetch available models
tags_response = await ssrf_safe_async_get(tags_url, headers=headers)
tags_response.raise_for_status()
models = tags_response.json()
if asyncio.iscoroutine(models):
models = await models
await logger.adebug(f"Available models: {models}")
# Filter models that are embedding models
model_ids = []
for model in models[self.JSON_MODELS_KEY]:
model_name = model[self.JSON_NAME_KEY]
await logger.adebug(f"Checking model: {model_name}")
# Filter models that are embedding models
model_ids = []
for model in models[self.JSON_MODELS_KEY]:
model_name = model[self.JSON_NAME_KEY]
await logger.adebug(f"Checking model: {model_name}")
payload = {"model": model_name}
show_response = await client.post(url=show_url, json=payload, headers=headers)
show_response.raise_for_status()
json_data = show_response.json()
if asyncio.iscoroutine(json_data):
json_data = await json_data
payload = {"model": model_name}
show_response = await ssrf_safe_async_post(show_url, json=payload, headers=headers)
show_response.raise_for_status()
json_data = show_response.json()
if asyncio.iscoroutine(json_data):
json_data = await json_data
capabilities = json_data.get(self.JSON_CAPABILITIES_KEY, [])
await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
capabilities = json_data.get(self.JSON_CAPABILITIES_KEY, [])
await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
if self.EMBEDDING_CAPABILITY in capabilities:
model_ids.append(model_name)
if self.EMBEDDING_CAPABILITY in capabilities:
model_ids.append(model_name)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except (httpx.RequestError, ValueError) as e:
msg = "Could not get model names from Ollama."
raise ValueError(msg) from e
@ -170,16 +184,18 @@ class OllamaEmbeddingsComponent(LCModelComponent):
async def is_valid_ollama_url(self, url: str) -> bool:
try:
async with httpx.AsyncClient() as client:
url = transform_localhost_url(url)
if not url:
return False
# Strip /v1 suffix if present, as Ollama API endpoints are at root level
url = url.rstrip("/").removesuffix("/v1")
if not url.endswith("/"):
url = url + "/"
return (
await client.get(url=urljoin(url, "api/tags"), headers=self.headers)
).status_code == HTTP_STATUS_OK
url = transform_localhost_url(url)
if not url:
return False
# Strip /v1 suffix if present, as Ollama API endpoints are at root level
url = url.rstrip("/").removesuffix("/v1")
if not url.endswith("/"):
url = url + "/"
response = await ssrf_safe_async_get(urljoin(url, "api/tags"), headers=self.headers)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
except httpx.RequestError:
return False
else:
return response.status_code == HTTP_STATUS_OK

13
src/lfx/src/lfx/components/xai/xai.py
View File

@ -1,4 +1,4 @@
import requests
import httpx
from langchain_openai import ChatOpenAI
from pydantic.v1 import SecretStr
from typing_extensions import override
@ -15,6 +15,8 @@ from lfx.inputs.inputs import (
SecretStrInput,
SliderInput,
)
from lfx.utils.ssrf_httpx import ssrf_protected_openai_clients_for_url, ssrf_safe_httpx_get
from lfx.utils.ssrf_protection import SSRFProtectionError
XAI_DEFAULT_MODELS = ["grok-2-latest"]
@ -97,7 +99,7 @@ class XAIModelComponent(LCModelComponent):
headers = {"Authorization": f"Bearer {self.api_key}", "Accept": "application/json"}
try:
response = requests.get(url, headers=headers, timeout=10)
response = ssrf_safe_httpx_get(url, headers=headers, timeout=10)
response.raise_for_status()
data = response.json()
@ -108,7 +110,10 @@ class XAIModelComponent(LCModelComponent):
models.update(model.get("aliases", []))
return sorted(models) if models else XAI_DEFAULT_MODELS
except requests.RequestException as e:
except SSRFProtectionError as e:
self.status = f"SSRF Protection: {e}"
return XAI_DEFAULT_MODELS
except httpx.HTTPError as e:
self.status = f"Error fetching models: {e}"
return XAI_DEFAULT_MODELS
@ -130,6 +135,7 @@ class XAIModelComponent(LCModelComponent):
json_mode = self.json_mode
seed = self.seed
ssrf_client_kwargs = ssrf_protected_openai_clients_for_url(base_url)
api_key = SecretStr(api_key).get_secret_value() if api_key else None
output = ChatOpenAI(
@ -140,6 +146,7 @@ class XAIModelComponent(LCModelComponent):
api_key=api_key,
temperature=temperature if temperature is not None else 0.1,
seed=seed,
**ssrf_client_kwargs,
)
if json_mode:

118
src/lfx/src/lfx/utils/ssrf_httpx.py Normal file
View File

@ -0,0 +1,118 @@
"""SSRF-protected helpers for ``httpx`` call sites."""
from __future__ import annotations
from typing import Any
from urllib.parse import urlparse
import httpx
from lfx.utils.ssrf_protection import (
SSRFProtectionError,
is_ssrf_protection_enabled,
validate_and_resolve_url,
validate_url_for_ssrf,
)
from lfx.utils.ssrf_transport import (
SSRFProtectedSyncTransport,
SSRFProtectedTransport,
create_ssrf_protected_client,
create_ssrf_protected_sync_client,
)
def validate_url_for_ssrf_or_raise(url: str) -> None:
"""Validate a user URL and raise a UI-facing error when it is blocked."""
try:
validate_url_for_ssrf(url, warn_only=False)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
def _raise_if_following_redirects(request_kwargs: dict[str, Any]) -> None:
if request_kwargs.get("follow_redirects"):
msg = "SSRF-protected httpx helpers do not support automatic redirect following."
raise SSRFProtectionError(msg)
def _async_client_for_url(url: str, validated_ips: list[str]) -> httpx.AsyncClient:
if is_ssrf_protection_enabled() and validated_ips:
hostname = urlparse(url).hostname
if hostname:
return create_ssrf_protected_client(hostname=hostname, validated_ips=validated_ips)
return httpx.AsyncClient()
def _sync_client_for_url(url: str, validated_ips: list[str]) -> httpx.Client:
if is_ssrf_protection_enabled() and validated_ips:
hostname = urlparse(url).hostname
if hostname:
return create_ssrf_protected_sync_client(hostname=hostname, validated_ips=validated_ips)
return httpx.Client()
def ssrf_protected_httpx_client_kwargs_for_url(url: str) -> tuple[dict[str, Any], dict[str, Any]]:
"""Return sync/async httpx kwargs that enforce SSRF protection for SDK clients."""
try:
validated_url, validated_ips = validate_and_resolve_url(url)
except SSRFProtectionError as e:
msg = f"SSRF Protection: {e}"
raise ValueError(msg) from e
if not is_ssrf_protection_enabled():
return {}, {}
sync_kwargs: dict[str, Any] = {"follow_redirects": False}
async_kwargs: dict[str, Any] = {"follow_redirects": False}
hostname = urlparse(validated_url).hostname
if hostname and validated_ips:
ip_list = list(validated_ips)
sync_kwargs["transport"] = SSRFProtectedSyncTransport(pinned_ips={hostname: ip_list})
async_kwargs["transport"] = SSRFProtectedTransport(pinned_ips={hostname: ip_list})
return sync_kwargs, async_kwargs
def ssrf_protected_openai_clients_for_url(url: str) -> dict[str, httpx.Client | httpx.AsyncClient]:
"""Return pinned sync and async clients for OpenAI-compatible LangChain components."""
sync_kwargs, async_kwargs = ssrf_protected_httpx_client_kwargs_for_url(url)
if not sync_kwargs and not async_kwargs:
return {}
return {
"http_client": httpx.Client(**sync_kwargs),
"http_async_client": httpx.AsyncClient(**async_kwargs),
}
async def ssrf_safe_async_get(url: str, **request_kwargs: Any) -> httpx.Response:
"""Perform an async GET with SSRF validation and DNS pinning."""
_raise_if_following_redirects(request_kwargs)
validated_url, validated_ips = validate_and_resolve_url(url)
async with _async_client_for_url(validated_url, validated_ips) as client:
return await client.get(url=validated_url, **request_kwargs)
async def ssrf_safe_async_post(url: str, **request_kwargs: Any) -> httpx.Response:
"""Perform an async POST with SSRF validation and DNS pinning."""
_raise_if_following_redirects(request_kwargs)
validated_url, validated_ips = validate_and_resolve_url(url)
async with _async_client_for_url(validated_url, validated_ips) as client:
return await client.post(url=validated_url, **request_kwargs)
def ssrf_safe_httpx_get(url: str, **request_kwargs: Any) -> httpx.Response:
"""Perform a synchronous GET with SSRF validation and DNS pinning."""
_raise_if_following_redirects(request_kwargs)
validated_url, validated_ips = validate_and_resolve_url(url)
with _sync_client_for_url(validated_url, validated_ips) as client:
return client.get(url=validated_url, **request_kwargs)
def ssrf_safe_httpx_post(url: str, **request_kwargs: Any) -> httpx.Response:
"""Perform a synchronous POST with SSRF validation and DNS pinning."""
_raise_if_following_redirects(request_kwargs)
validated_url, validated_ips = validate_and_resolve_url(url)
with _sync_client_for_url(validated_url, validated_ips) as client:
return client.post(url=validated_url, **request_kwargs)

131
src/lfx/src/lfx/utils/ssrf_transport.py
View File

@ -135,6 +135,82 @@ class DNSPinningNetworkBackend(httpcore.AsyncNetworkBackend):
await self._backend.sleep(seconds)
class DNSPinningSyncNetworkBackend(httpcore.NetworkBackend):
"""Synchronous network backend that pins DNS resolution to validated IP addresses."""
def __init__(self, pinned_ips: dict[str, list[str]], backend: httpcore.NetworkBackend | None = None):
self.pinned_ips = pinned_ips
if backend is None:
backend = httpcore.SyncBackend()
self._backend = backend
logger.debug(f"Created sync DNS pinning network backend with pinned IPs: {pinned_ips}")
def connect_tcp(
self,
host: str,
port: int,
timeout: float | None = None,
local_address: str | None = None,
socket_options: Iterable[httpcore.SOCKET_OPTION] | None = None,
) -> httpcore.NetworkStream:
"""Connect to TCP socket, using pinned IPs for hosts that were validated."""
if host in self.pinned_ips:
pinned_ips = self.pinned_ips[host]
if not pinned_ips:
msg = f"DNS pinning: Host {host} is marked for pinning but has no pinned IPs"
logger.error(msg)
raise RuntimeError(msg)
logger.debug(f"DNS pinning: Connecting to pinned IPs {pinned_ips} for hostname {host}")
last_error = None
for pinned_ip in pinned_ips:
try:
logger.debug(f"DNS pinning: Attempting connection to {pinned_ip}")
return self._backend.connect_tcp(
host=pinned_ip,
port=port,
timeout=timeout,
local_address=local_address,
socket_options=socket_options,
)
except (OSError, TimeoutError) as e:
last_error = e
logger.debug(f"DNS pinning: Failed to connect to {pinned_ip}: {e}")
continue
if last_error is None:
msg = f"DNS pinning: All pinned IPs failed for {host} but no error was captured"
raise RuntimeError(msg)
raise last_error
return self._backend.connect_tcp(
host=host,
port=port,
timeout=timeout,
local_address=local_address,
socket_options=socket_options,
)
def connect_unix_socket(
self,
path: str,
timeout: float | None = None,
socket_options: Iterable[httpcore.SOCKET_OPTION] | None = None,
) -> httpcore.NetworkStream:
"""Connect to Unix socket (pass through to underlying backend)."""
return self._backend.connect_unix_socket(
path=path,
timeout=timeout,
socket_options=socket_options,
)
def sleep(self, seconds: float) -> None:
"""Sleep for specified duration (pass through to underlying backend)."""
self._backend.sleep(seconds)
class SSRFProtectedTransport(httpx.AsyncHTTPTransport):
"""HTTP transport that pins DNS resolution to validated IPs.
@ -222,6 +298,52 @@ class SSRFProtectedTransport(httpx.AsyncHTTPTransport):
logger.debug(f"Created SSRF protected transport with pinned IPs: {pinned_ips}")
class SSRFProtectedSyncTransport(httpx.HTTPTransport):
"""Synchronous HTTP transport that pins DNS resolution to validated IPs."""
def __init__(
self,
pinned_ips: dict[str, list[str]],
verify: bool | str | ssl.SSLContext = True, # noqa: FBT001, FBT002
cert: tuple[str, str] | tuple[str, str, str] | str | None = None,
trust_env: bool = True, # noqa: FBT001, FBT002
http1: bool = True, # noqa: FBT001, FBT002
http2: bool = False, # noqa: FBT001, FBT002
limits: httpx.Limits = httpx.Limits(), # noqa: B008
proxy: httpx._types.ProxyTypes | None = None,
uds: str | None = None,
local_address: str | None = None,
retries: int = 0,
socket_options: Iterable[httpcore.SOCKET_OPTION] | None = None,
):
network_backend = DNSPinningSyncNetworkBackend(pinned_ips=pinned_ips)
ssl_context = create_ssl_context(verify=verify, cert=cert, trust_env=trust_env)
if proxy is not None:
proxy = Proxy(url=proxy) if isinstance(proxy, (str, URL)) else proxy
if proxy is None:
self._pool = httpcore.ConnectionPool(
ssl_context=ssl_context,
max_connections=limits.max_connections,
max_keepalive_connections=limits.max_keepalive_connections,
keepalive_expiry=limits.keepalive_expiry,
http1=http1,
http2=http2,
uds=uds,
local_address=local_address,
retries=retries,
socket_options=socket_options,
network_backend=network_backend,
)
else:
msg = "DNS pinning with proxies is not currently supported"
raise NotImplementedError(msg)
self.pinned_ips = pinned_ips
logger.debug(f"Created synchronous SSRF protected transport with pinned IPs: {pinned_ips}")
def create_ssrf_protected_client(
hostname: str, validated_ips: list[str] | tuple[str, ...], **client_kwargs
) -> httpx.AsyncClient:
@ -240,3 +362,12 @@ def create_ssrf_protected_client(
ip_list = list(validated_ips) if isinstance(validated_ips, tuple) else validated_ips
transport = SSRFProtectedTransport(pinned_ips={hostname: ip_list})
return httpx.AsyncClient(transport=transport, **client_kwargs)
def create_ssrf_protected_sync_client(
hostname: str, validated_ips: list[str] | tuple[str, ...], **client_kwargs
) -> httpx.Client:
"""Create a synchronous httpx client with DNS pinning for SSRF protection."""
ip_list = list(validated_ips) if isinstance(validated_ips, tuple) else validated_ips
transport = SSRFProtectedSyncTransport(pinned_ips={hostname: ip_list})
return httpx.Client(transport=transport, **client_kwargs)

55
src/lfx/tests/unit/utils/test_ssrf_httpx.py Normal file
View File

@ -0,0 +1,55 @@
import os
import socket
from unittest.mock import patch
import httpcore
import pytest
from lfx.utils.ssrf_httpx import ssrf_safe_httpx_get
from lfx.utils.ssrf_protection import SSRFProtectionError
class TestSSRFSafeHTTPX:
def test_direct_internal_ip_is_blocked(self):
with (
patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
patch("httpx.Client.get") as mock_get,
pytest.raises(SSRFProtectionError),
):
ssrf_safe_httpx_get("http://169.254.169.254/latest/meta-data/", timeout=5)
mock_get.assert_not_called()
def test_sync_dns_pinning_prevents_rebinding_attack(self):
call_count = 0
connected_to_ip = None
def mock_getaddrinfo(_hostname, _port, *_args, **_kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 0))]
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0))]
def mock_connect_tcp(_self, host, port, **_kwargs):
nonlocal connected_to_ip
assert port == 8080
connected_to_ip = host
return httpcore.MockStream(
[
b"HTTP/1.1 200 OK\r\n",
b"Content-Type: application/json\r\n",
b"Content-Length: 15\r\n",
b"\r\n",
b'{"status":"ok"}',
]
)
with (
patch("socket.getaddrinfo", side_effect=mock_getaddrinfo),
patch.dict(os.environ, {"LANGFLOW_SSRF_PROTECTION_ENABLED": "true"}),
patch.object(httpcore.SyncBackend, "connect_tcp", mock_connect_tcp),
):
response = ssrf_safe_httpx_get("http://rebinding.test:8080/models", timeout=5)
assert response.status_code == 200
assert call_count == 1
assert connected_to_ip == "8.8.8.8"