Security: Fix CWE-95 (Eval Injection) in _get_torch_dtype()

src/unitxt/inference.py

@@ -546,17 +546,37 @@ def _get_torch_dtype(self):
                 f"'{self.torch_dtype}' was given instead."
             )
 
-        try:
-            dtype = eval(self.torch_dtype)
-        except (AttributeError, TypeError) as e:
-            raise ValueError(
-                f"Incorrect value of 'torch_dtype' was given: '{self.torch_dtype}'."
-            ) from e
+        # Security fix: Use a lookup table instead of eval() to prevent code injection
+        # This addresses CWE-95 (Eval Injection) vulnerability
+        torch_dtypes = {
+            "torch.float16": torch.float16,
+            "torch.float32": torch.float32,
+            "torch.float64": torch.float64,
+            "torch.bfloat16": torch.bfloat16,
+            "torch.float": torch.float,
+            "torch.double": torch.double,
+            "torch.half": torch.half,
+            "torch.int8": torch.int8,
+            "torch.int16": torch.int16,
+            "torch.int32": torch.int32,
+            "torch.int64": torch.int64,
+            "torch.int": torch.int,
+            "torch.long": torch.long,
+            "torch.short": torch.short,
+            "torch.uint8": torch.uint8,
+            "torch.bool": torch.bool,
+            "torch.complex64": torch.complex64,
+            "torch.complex128": torch.complex128,
+            "torch.cfloat": torch.cfloat,
+            "torch.cdouble": torch.cdouble,
+        }
+
+        dtype = torch_dtypes.get(self.torch_dtype)
 
-        if not isinstance(dtype, torch.dtype):
+        if dtype is None:
             raise ValueError(
-                f"'torch_dtype' must be an instance of 'torch.dtype', however, "
-                f"'{dtype}' is an instance of '{type(dtype)}'."
+                f"Incorrect value of 'torch_dtype' was given: '{self.torch_dtype}'. "
+                f"Supported values are: {', '.join(sorted(torch_dtypes.keys()))}"
             )
 
         return dtype

tests/inference/test_inference_engine.py

@@ -680,3 +680,41 @@ def test_hf_auto_model_and_hf_pipeline_equivalency(self):
             self.assertEqual(
                 pipeline_inference_model_predictions, auto_inference_model_predictions
             )
+
+    def test_torch_dtype_security_fix_fast(self):
+        """Fast unit test for CWE-95 security fix that doesn't load models.
+
+        This test directly tests the _get_torch_dtype() method without
+        initializing the full inference engine, making it much faster.
+        """
+        import torch
+
+        # Create a minimal mock engine with just torch_dtype attribute
+        engine = HFAutoModelInferenceEngine.__new__(HFAutoModelInferenceEngine)
+
+        # Test valid dtypes
+        valid_dtypes = [
+            ("torch.float16", torch.float16),
+            ("torch.float32", torch.float32),
+            ("torch.bfloat16", torch.bfloat16),
+        ]
+
+        for dtype_str, expected_dtype in valid_dtypes:
+            engine.torch_dtype = dtype_str
+            result = engine._get_torch_dtype()
+            self.assertEqual(result, expected_dtype)
+
+        # Test malicious payload is rejected
+        malicious_payload = 'torch.typename.__globals__["__builtins__"]["__import__"]("os").system("id")'
+        engine.torch_dtype = malicious_payload
+
+        with self.assertRaises(ValueError) as context:
+            engine._get_torch_dtype()
+
+        self.assertIn("Incorrect value of 'torch_dtype'", str(context.exception))
+
+        # Test invalid dtypes are rejected
+        for invalid_dtype in ["torch.invalid_dtype", "torch.float128", "numpy.float32"]:
+            engine.torch_dtype = invalid_dtype
+            with self.assertRaises(ValueError):
+                engine._get_torch_dtype()