dosubot[bot]
commented
1 month ago Hey @0xfatty, I'm here to help you with any bugs, questions, or contributions you have in mind. Let's get started!
To address the security vulnerability in the Custom Component feature of the Langflow framework that allows for arbitrary code execution when hosted online, you should implement input validation and sanitization. Additionally, you can leverage the existing structure of the Component class to ensure that the code provided is safe and does not contain malicious elements.
Here is an example of how you can modify the build method to include input validation and sanitization, while also utilizing the existing Component class structure:
import re
from langflow.custom import CustomComponent
from langflow.schema import Record
from langflow.custom.eval import eval_custom_component_code
from langflow.utils import validate
class Component(CustomComponent):
display_name = "Custom Component"
description = "Use as a template to create your own component."
documentation: str = "http://docs.langflow.org/components/custom"
icon = "custom_components"
def build_config(self):
return {"param": {"display_name": "Parameter"}}
def build(self, param: str) -> Record:
# Basic input validation to prevent arbitrary code execution
if not isinstance(param, str):
raise ValueError("Parameter must be a string")
# Example of sanitizing input to remove potentially dangerous characters
sanitized_param = re.sub(r'[^\w\s]', '', param)
# Validate and execute the custom component code safely
if not self.code:
raise ComponentCodeNullError(
status_code=400,
detail={"error": self.ERROR_CODE_NULL, "traceback": ""},
)
if not self._function_entrypoint_name:
raise ComponentFunctionEntrypointNameNullError(
status_code=400,
detail={
"error": self.ERROR_FUNCTION_ENTRYPOINT_NAME_NULL,
"traceback": "",
},
)
custom_function = validate.create_function(self.code, self._function_entrypoint_name)
result = custom_function(sanitized_param)
return Record(data=result)In this example:
- The
buildmethod checks if theparamis a string. - It sanitizes the input by removing any non-alphanumeric characters using a regular expression.
- It validates and executes the custom component code safely by leveraging the existing
Componentclass methods.
This approach helps mitigate the risk of arbitrary code execution by ensuring that the input is safe before processing it and by using the existing validation mechanisms provided by the Langflow framework [1][2].
#### Details
**About Dosu** This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research. Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.To continue the conversation, mention @dosu.
0xfatty
YamonBot
Quick note Langflow is a very interesting and useful framework for those who work with AI projects. Personally, it's a cool project to learn from. It seems the Custom Component feature was launched to support user-defined Python scripts that use the Langflow provided library. This is a great feature. However, upon reviewing its documentation in [1] and [2], it seems there is no mention of potential security issues. Hence, perhaps the following finding was likely not expected. [1] https://docs.langflow.org/components/custom [2] https://docs.langflow.org/guidelines/custom-component
Describe the bug The Custom Component feature allows users to provide their own Python scripts using the CustomComponent class provided by the Langflow library. This is excellent for local testing and experimentation. However, if the framework is hosted online, it creates a potential security issue where a bad actor can leverage the opportunity to provide arbitrary Python code and gain code execution ability against the hosting server.
Impacted API POST /api/v1/custom_component
Browser and Version
To Reproduce Steps to reproduce the behavior:
CustomComponent, withinComponentclass, provide the following Python functionCheck & Save, the/api/v1/custom_componentAPI is invoked to process the provided Python script, which then leads to OS command execution. The output will be Base64 encoded and sent to a malicious server.Screenshots
Additional context The vulnerability allows for arbitrary code execution by injecting malicious code through the Custom Component feature. This could lead to significant security risks, including data theft, unauthorized access, and potential disruption of services (especially when being hosted publicly)