Security: safely parse untrusted symbolic gain expressions

README.md

@@ -62,6 +62,27 @@ For a quick demonstration, press **Load Example** and then **Solve Diagram**.
 
 ---
 
+## Supported Gain Expressions
+
+Branch gains are parsed as **symbolic mathematics**, never as executable Python.
+Examples include:
+
+- `G1*G2` or `G_1*G_2` — multiply symbolic gains.
+- `1/(s+1)` — rational transfer expressions.
+- `G^2` — exponentiation (the caret is accepted).
+- `sin(s)`, `sqrt(s^2+1)`, `exp(-s)`, and `Rational(1,2)` — supported math functions.
+
+Allowed operations: addition, subtraction, multiplication, division, powers,
+parentheses, numbers, variable names, and selected mathematical functions.
+Other Python constructs (imports, property access, indexing, string literals,
+arbitrary function calls, or comprehensions) are rejected. To reduce resource
+abuse, gains have limits on expression length and complexity.
+
+If the solver reports an invalid gain, replace it with a mathematical
+expression using the supported syntax.
+
+---
+
 ## How the Solver Works
 
 The application follows the same engineering workflow used when solving a Signal Flow Graph manually:

app.py

@@ -2,6 +2,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from fastapi.staticfiles import StaticFiles
+from gain_parser import parse_gain
 import networkx as nx
 import sympy as sp
 import itertools
@@ -21,9 +22,6 @@ class DiagramData(BaseModel):
     output_node: str
     nodes: list[str] = []
 
-def _normalize_gain(expr: str) -> str:
-    return (expr or "").replace("^", "**")
-
 @app.post("/validate")
 def validate_diagram(data: DiagramData):
     errors: list[dict] = []
@@ -66,7 +64,7 @@ def validate_diagram(data: DiagramData):
             continue
 
         try:
-            sym_weight = sp.sympify(_normalize_gain(raw_gain))
+            sym_weight = parse_gain(raw_gain)
         except Exception as e:
             errors.append({"type": "invalid_gain", "message": f"Invalid gain on edge {frm} -> {to}: {e}", "edge": {"from": frm, "to": to}, "gain": raw_gain})
             continue
@@ -201,7 +199,7 @@ def solve_diagram(data: DiagramData):
         G.add_nodes_from(node_names)
         dummy_idx = 0
         for conn in data.connections:
-            sym_weight = sp.sympify(_normalize_gain(conn.gain))
+            sym_weight = parse_gain(conn.gain)
             frm = (conn.from_node or "").strip()
             to = (conn.to_node or "").strip()
 

gain_parser.py

@@ -0,0 +1,85 @@
+"""Parse symbolic signal-flow gains without evaluating arbitrary Python code."""
+
+import ast
+import math
+
+import sympy as sp
+
+_ALLOWED_FUNCTIONS = {
+    "sin": sp.sin, "cos": sp.cos, "tan": sp.tan,
+    "asin": sp.asin, "acos": sp.acos, "atan": sp.atan,
+    "sinh": sp.sinh, "cosh": sp.cosh, "tanh": sp.tanh,
+    "exp": sp.exp, "log": sp.log, "sqrt": sp.sqrt,
+    "Abs": sp.Abs, "Rational": sp.Rational,
+}
+_ALLOWED_CONSTANTS = {"pi": sp.pi, "E": sp.E, "I": sp.I}
+_MAX_EXPRESSION_LENGTH = 256
+_MAX_AST_NODES = 96
+_MAX_LITERAL_EXPONENT = 100
+
+
+def parse_gain(source: str) -> sp.Expr:
+    """Parse gains such as G1*G2, 1/(s+1), and sqrt(s^2+1).
+
+    The restricted AST allows numbers, symbols, arithmetic and a short list
+    of math functions. It never evaluates a gain as arbitrary Python code.
+    """
+    expression = (source or "").strip().replace("^", "**")
+    if not expression:
+        raise ValueError("Gain expression is empty")
+    if len(expression) > _MAX_EXPRESSION_LENGTH:
+        raise ValueError("Gain expression is too long")
+
+    try:
+        parsed = ast.parse(expression, mode="eval")
+    except SyntaxError as exc:
+        raise ValueError("Invalid gain expression syntax") from exc
+    if sum(1 for _ in ast.walk(parsed)) > _MAX_AST_NODES:
+        raise ValueError("Gain expression is too complex")
+
+    def convert(node: ast.AST) -> sp.Expr:
+        if isinstance(node, ast.Constant):
+            value = node.value
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError("Only numeric constants are allowed")
+            if isinstance(value, float):
+                if not math.isfinite(value):
+                    raise ValueError("Gain constants must be finite")
+                return sp.Float(str(value))
+            return sp.Integer(value)
+        if isinstance(node, ast.Name):
+            return _ALLOWED_CONSTANTS.get(node.id, sp.Symbol(node.id))
+        if isinstance(node, ast.UnaryOp):
+            operand = convert(node.operand)
+            if isinstance(node.op, ast.USub):
+                return -operand
+            if isinstance(node.op, ast.UAdd):
+                return operand
+        if isinstance(node, ast.BinOp):
+            left, right = convert(node.left), convert(node.right)
+            if isinstance(node.op, ast.Add):
+                return left + right
+            if isinstance(node.op, ast.Sub):
+                return left - right
+            if isinstance(node.op, ast.Mult):
+                return left * right
+            if isinstance(node.op, ast.Div):
+                return left / right
+            if isinstance(node.op, ast.Pow):
+                if right.is_number and not right.free_symbols:
+                    if right.is_real is not True or abs(right) > _MAX_LITERAL_EXPONENT:
+                        raise ValueError("Numeric exponent is outside the supported range")
+                return left ** right
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            function = _ALLOWED_FUNCTIONS.get(node.func.id)
+            if function is not None and not node.keywords:
+                try:
+                    return function(*(convert(arg) for arg in node.args))
+                except (TypeError, ValueError, ZeroDivisionError) as exc:
+                    raise ValueError("Invalid gain function arguments") from exc
+        raise ValueError("Unsupported gain expression construct")
+
+    try:
+        return convert(parsed.body)
+    except (TypeError, ZeroDivisionError) as exc:
+        raise ValueError("Invalid gain expression") from exc

tests/test_gain_parser.py

@@ -0,0 +1,72 @@
+"""Security and compatibility tests for the signal-flow gain parser."""
+
+import unittest
+
+import sympy as sp
+
+from app import Connection, DiagramData, solve_diagram, validate_diagram
+from gain_parser import parse_gain
+
+
+class GainParserTests(unittest.TestCase):
+    def test_arithmetic_and_control_system_notation(self):
+        cases = [
+            ("G1*G2", "G1*G2"),
+            ("1/(s+1)", "1/(s+1)"),
+            ("G_1^2/(1+H)", "G_1**2/(1+H)"),
+            ("-0.5*H", "-0.5*H"),
+            ("1e-3*s", "0.001*s"),
+            ("Rational(1,2)*K", "K/2"),
+            ("sin(s)+sqrt(s^2+1)", "sin(s)+sqrt(s**2+1)"),
+        ]
+        for expression, expected in cases:
+            with self.subTest(expression=expression):
+                actual = parse_gain(expression)
+                self.assertEqual(sp.simplify(actual - sp.sympify(expected)), 0)
+
+    def test_rejects_python_execution_and_attribute_access(self):
+        invalid = [
+            '__import__("os").system("echo unsafe")',
+            'open("secret", "r")',
+            "(0).__class__",
+            "s[0]",
+            "[G1 for G1 in (1,2)]",
+            "lambda s: s",
+            "Symbol('x')",
+            "unknown_function(s)",
+            "2**1000000",
+        ]
+        for expression in invalid:
+            with self.subTest(expression=expression):
+                with self.assertRaises(ValueError):
+                    parse_gain(expression)
+
+    def test_rejects_non_finite_or_oversized_expressions(self):
+        for expression in ["1e309", "G" * 257, ""]:
+            with self.subTest(expression=expression):
+                with self.assertRaises(ValueError):
+                    parse_gain(expression)
+
+    def test_api_rejects_unsafe_gain_consistently(self):
+        data = DiagramData(
+            nodes=["S", "T"],
+            input_node="S",
+            output_node="T",
+            connections=[
+                Connection(
+                    id="unsafe",
+                    from_node="S",
+                    to_node="T",
+                    gain='__import__("os").system("echo unsafe")',
+                )
+            ],
+        )
+        validated = validate_diagram(data)
+        solved = solve_diagram(data)
+        self.assertEqual(validated["status"], "error")
+        self.assertIn("invalid_gain", {e["type"] for e in validated["errors"]})
+        self.assertEqual(solved["status"], "error")
+
+
+if __name__ == "__main__":
+    unittest.main()