Skip to content

Commit d141e4f

Browse files
authored
Set a custom toml loader for Dynaconf (#2087)
1 parent 8bd134e commit d141e4f

File tree

4 files changed

+202
-5
lines changed

4 files changed

+202
-5
lines changed

pr_agent/config_loader.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,15 @@
88
PR_AGENT_TOML_KEY = 'pr-agent'
99

1010
current_dir = dirname(abspath(__file__))
11+
12+
dynconf_kwargs = {'core_loaders': [], # DISABLE default loaders, otherwise will load toml files more than once.
13+
'loaders': ['pr_agent.custom_merge_loader', 'dynaconf.loaders.env_loader'], # Use a custom loader to merge sections, but overwrite their overlapping values. Also support ENV variables to take precedence.
14+
'root_path': join(current_dir, "settings"), #Used for Dynaconf.find_file() - So that root path points to settings folder, since we disabled all core loaders.
15+
'merge_enabled': True # In case more than one file is sent, merge them. Must be set to True, otherwise, a .toml file with section [XYZ] overwrites the entire section of a previous .toml file's [XYZ] and we want it to only overwrite the overlapping fields under such section
16+
}
1117
global_settings = Dynaconf(
1218
envvar_prefix=False,
13-
merge_enabled=True,
19+
load_dotenv=False, # Security: Don't load .env files
1420
settings_files=[join(current_dir, f) for f in [
1521
"settings/configuration.toml",
1622
"settings/ignore.toml",
@@ -33,7 +39,8 @@
3339
"settings/pr_help_docs_headings_prompts.toml",
3440
"settings/.secrets.toml",
3541
"settings_prod/.secrets.toml",
36-
]]
42+
]],
43+
**dynconf_kwargs
3744
)
3845

3946

pr_agent/custom_merge_loader.py

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
from pathlib import Path
2+
import tomllib #tomllib should be used instead of Py toml for Python 3.11+
3+
4+
from jinja2.exceptions import SecurityError
5+
6+
from pr_agent.log import get_logger
7+
8+
def load(obj, env=None, silent=True, key=None, filename=None):
9+
"""
10+
Load and merge TOML configuration files into a Dynaconf settings object using a secure, in-house loader.
11+
This loader:
12+
- Replaces list and dict fields instead of appending/updating (non-default Dynaconf behavior).
13+
- Enforces several security checks (e.g., disallows includes/preloads and enforces .toml files).
14+
- Supports optional single-key loading.
15+
Args:
16+
obj: The Dynaconf settings instance to update.
17+
env: The current environment name (upper case). Defaults to 'DEVELOPMENT'. Note: currently unused.
18+
silent (bool): If True, suppress exceptions and log warnings/errors instead.
19+
key (str | None): Load only this top-level key (section) if provided; otherwise, load all keys from the files.
20+
filename (str | None): Custom filename for tests (not used when settings_files are provided).
21+
Returns:
22+
None
23+
"""
24+
25+
MAX_TOML_SIZE_IN_BYTES = 100 * 1024 * 1024 # Prevent out of mem. exceptions by limiting to 100 MBs which is sufficient for upto 1M lines
26+
27+
# Get the list of files to load
28+
# TODO: hasattr(obj, 'settings_files') for some reason returns False. Need to use 'settings_file'
29+
settings_files = obj.settings_files if hasattr(obj, 'settings_files') else (
30+
obj.settings_file) if hasattr(obj, 'settings_file') else []
31+
if not settings_files or not isinstance(settings_files, list):
32+
get_logger().warning("No settings files specified, or missing keys "
33+
"(tried looking for 'settings_files' or 'settings_file'), or not a list. Skipping loading.",
34+
artifact={'toml_obj_attributes_names': dir(obj)})
35+
return
36+
37+
# Storage for all loaded data
38+
accumulated_data = {}
39+
40+
# Security: Check for forbidden configuration options
41+
if hasattr(obj, 'includes') and obj.includes:
42+
if not silent:
43+
raise SecurityError("Configuration includes forbidden option: 'includes'. Skipping loading.")
44+
get_logger().error("Configuration includes forbidden option: 'includes'. Skipping loading.")
45+
return
46+
if hasattr(obj, 'preload') and obj.preload:
47+
if not silent:
48+
raise SecurityError("Configuration includes forbidden option: 'preload'. Skipping loading.")
49+
get_logger().error("Configuration includes forbidden option: 'preload'. Skipping loading.")
50+
return
51+
52+
for settings_file in settings_files:
53+
try:
54+
# Load the TOML file
55+
file_path = Path(settings_file)
56+
# Security: Only allow .toml files
57+
if file_path.suffix.lower() != '.toml':
58+
get_logger().warning(f"Only .toml files are allowed. Skipping: {settings_file}")
59+
continue
60+
61+
if not file_path.exists():
62+
get_logger().warning(f"Settings file not found: {settings_file}. Skipping it.")
63+
continue
64+
65+
if file_path.stat().st_size > MAX_TOML_SIZE_IN_BYTES:
66+
get_logger().warning(f"Settings file too large (> {MAX_TOML_SIZE_IN_BYTES} bytes): {settings_file}. Skipping it.")
67+
continue
68+
69+
with open(file_path, 'rb') as f:
70+
file_data = tomllib.load(f)
71+
72+
# Handle sections (like [config], [default], etc.)
73+
if not isinstance(file_data, dict):
74+
get_logger().warning(f"TOML root is not a table in '{settings_file}'. Skipping.")
75+
continue
76+
77+
# Security: Check file contents for forbidden directives
78+
validate_file_security(file_data, settings_file)
79+
80+
for section_name, section_data in file_data.items():
81+
if not isinstance(section_data, dict):
82+
get_logger().warning(f"Section '{section_name}' in '{settings_file}' is not a table. Skipping.")
83+
continue
84+
for field, field_value in section_data.items():
85+
if section_name not in accumulated_data:
86+
accumulated_data[section_name] = {}
87+
accumulated_data[section_name][field] = field_value
88+
89+
except Exception as e:
90+
if not silent:
91+
raise e
92+
get_logger().exception(f"Exception loading settings file: {settings_file}. Skipping.")
93+
94+
# Update the settings object
95+
for k, v in accumulated_data.items():
96+
if key is None or key == k:
97+
obj.set(k, v)
98+
99+
def validate_file_security(file_data, filename):
100+
"""
101+
Validate that the config file does not contain security-sensitive directives.
102+
103+
Args:
104+
file_data: Parsed TOML data representing the configuration contents.
105+
filename: The name or path of the file being validated (used for error messages).
106+
107+
Raises:
108+
SecurityError: If forbidden directives are found within the configuration, or if data too nested.
109+
"""
110+
MAX_DEPTH = 50
111+
112+
# Check for forbidden keys
113+
# Comprehensive list of forbidden keys with explanations
114+
forbidden_keys_to_reasons = {
115+
# Include mechanisms - allow loading arbitrary files
116+
'dynaconf_include': 'allows including other config files dynamically',
117+
'dynaconf_includes': 'allows including other config files dynamically',
118+
'includes': 'allows including other config files dynamically',
119+
120+
# Preload mechanisms - allow loading files before main config
121+
'preload': 'allows preloading files with potential code execution',
122+
'preload_for_dynaconf': 'allows preloading files with potential code execution',
123+
'preloads': 'allows preloading files with potential code execution',
124+
125+
# Merge controls - could be used to manipulate config behavior
126+
'dynaconf_merge': 'allows manipulating merge behavior',
127+
'dynaconf_merge_enabled': 'allows manipulating merge behavior',
128+
'merge_enabled': 'allows manipulating merge behavior',
129+
130+
# Loader controls - allow changing how configs are loaded
131+
'loaders_for_dynaconf': 'allows overriding loaders to execute arbitrary code',
132+
'loaders': 'allows overriding loaders to execute arbitrary code',
133+
'core_loaders': 'allows overriding core loaders',
134+
'core_loaders_for_dynaconf': 'allows overriding core loaders',
135+
136+
# Settings module - allows loading Python modules
137+
'settings_module': 'allows loading Python modules with code execution',
138+
'settings_file_for_dynaconf': 'could override settings file location',
139+
'settings_files_for_dynaconf': 'could override settings file location',
140+
141+
# Environment variable prefix manipulation
142+
'envvar_prefix': 'allows changing environment variable prefix',
143+
'envvar_prefix_for_dynaconf': 'allows changing environment variable prefix',
144+
}
145+
146+
# Check at the top level and in all sections
147+
def check_dict(data, path="", max_depth=MAX_DEPTH):
148+
if max_depth <= 0:
149+
raise SecurityError(
150+
f"Maximum nesting depth exceeded at {path}. "
151+
f"Possible attempt to cause stack overflow."
152+
)
153+
154+
for key, value in data.items():
155+
full_path = f"{path}.{key}" if path else key
156+
157+
if key.lower() in forbidden_keys_to_reasons:
158+
raise SecurityError(
159+
f"Security error in {filename}: "
160+
f"Forbidden directive '{key}' found at {full_path}. Reason: {forbidden_keys_to_reasons[key.lower()]}"
161+
)
162+
163+
# Recursively check nested dicts
164+
if isinstance(value, dict):
165+
check_dict(value, path=full_path, max_depth=(max_depth - 1))
166+
167+
check_dict(file_data, max_depth=MAX_DEPTH)

pr_agent/git_providers/utils.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,17 @@ def apply_repo_settings(pr_url):
3838
os.write(fd, repo_settings)
3939

4040
try:
41+
dynconf_kwargs = {'core_loaders': [], # DISABLE default loaders, otherwise will load toml files more than once.
42+
'loaders': ['pr_agent.custom_merge_loader'],
43+
# Use a custom loader to merge sections, but overwrite their overlapping values. Don't involve ENV variables.
44+
'merge_enabled': True # Merge multiple files; ensures [XYZ] sections only overwrite overlapping keys, not whole sections.
45+
}
46+
4147
new_settings = Dynaconf(settings_files=[repo_settings_file],
4248
# Disable all dynamic loading features
4349
load_dotenv=False, # Don't load .env files
44-
merge_enabled=False, # Don't allow merging from other sources
50+
envvar_prefix=False, # Drop DYNACONF for env. variables
51+
**dynconf_kwargs
4552
)
4653
except TypeError as e:
4754
# Fallback for older Dynaconf versions that don't support these parameters

pr_agent/tools/pr_config.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,24 @@ async def run(self):
3030
return ""
3131

3232
def _prepare_pr_configs(self) -> str:
33-
conf_file = get_settings().find_file("configuration.toml")
34-
conf_settings = Dynaconf(settings_files=[conf_file])
33+
try:
34+
conf_file = get_settings().find_file("configuration.toml")
35+
dynconf_kwargs = {'core_loaders': [], # DISABLE default loaders, otherwise will load toml files more than once.
36+
'loaders': ['pr_agent.custom_merge_loader'],
37+
# Use a custom loader to merge sections, but overwrite their overlapping values. Do not use ENV variables.
38+
'merge_enabled': True
39+
# Merge multiple TOML files; prevent full section overwrite—only overlapping keys in sections overwrite prior ones.
40+
}
41+
conf_settings = Dynaconf(settings_files=[conf_file],
42+
# Security: Disable all dynamic loading features
43+
load_dotenv=False, # Don't load .env files
44+
envvar_prefix=False,
45+
**dynconf_kwargs
46+
)
47+
except Exception as e:
48+
get_logger().error("Caught exception during Dynaconf loading. Returning empty dict",
49+
artifact={"exception": e})
50+
conf_settings = {}
3551
configuration_headers = [header.lower() for header in conf_settings.keys()]
3652
relevant_configs = {
3753
header: configs for header, configs in get_settings().to_dict().items()

0 commit comments

Comments
 (0)