fix: make eval schedule string safe (#126)

Author: hanxiao

discoart/helper.py

@@ -20,7 +20,6 @@
 import yaml
 from clip.simple_tokenizer import SimpleTokenizer, whitespace_clean, basic_clean
 from packaging.version import Version
-
 from spellchecker import SpellChecker
 from tqdm.auto import tqdm
 
@@ -690,9 +689,21 @@ def _version_check(package: str = None, github_repo: str = None):
 _MAX_DIFFUSION_STEPS = 1000
 
 
+def _is_valid_schedule_str(val) -> bool:
+    r = re.match(r'(False\b|True\b|[\(\)\[\]0-9\, \.\*\+\-])+', val)
+    if r and r.group(0) == val:
+        return True
+    return False
+
+
 def _eval_scheduling_str(val) -> List[float]:
     if isinstance(val, str):
-        val = eval(val)
+        if _is_valid_schedule_str(val):
+            val = eval(val)
+        else:
+            raise ValueError(
+                f'invalid scheduling string: {val}, it contains unsafe code'
+            )
 
     if isinstance(val, (int, float, bool)):
         val = [val] * _MAX_DIFFUSION_STEPS

tests/test_config.py

@@ -1,5 +1,7 @@
+import pytest
+
 from discoart.config import default_args, save_config, load_config, export_python
-from discoart.helper import _eval_scheduling_str
+from discoart.helper import _eval_scheduling_str, _is_valid_schedule_str
 
 
 def test_export_load_config(tmpfile):
@@ -34,3 +36,28 @@ def test_eval_schedule_string():
     assert _eval_scheduling_str('True') == [True] * 1000
     assert _eval_scheduling_str('False') == [False] * 1000
     assert _eval_scheduling_str(True) == [True] * 1000
+
+
+@pytest.mark.parametrize(
+    'val, expected',
+    [
+        ('[100]*600+[200]*400', True),
+        ('[100]*600+[2.3]*400', True),
+        ('[100]*600+[2.3]*400', True),
+        ('1', True),
+        ('True', True),
+        ('Truetrue', False),
+        ('False', True),
+        ('true', False),
+        ('sdd ds', False),
+        ('[True, False]*1000', True),
+        ('[True]*500+[False]*400', True),
+        ('[0.5]*400+[0.2]*300+[True]*200', True),
+        ('[hello]*1000', False),
+        ('del a', False),
+        ('([1]+[2])*50', True),
+        ('[False,True,1,0.23,23,]*1000', True),
+    ],
+)
+def test_chec_schedule_str(val, expected):
+    assert _is_valid_schedule_str(val) == expected