Repository: DCC-Lab/RayTracing

Issue:
Rays can load any "pickled" file
This can lead to unwanted behaviors. We should check if the content of the file is an iterable of `Ray` objects.

`load` directly extends (or override) the inner list or rays in `Rays`. This can cause unexpected problems without proper guidance to what caused the issue.

Comments:

Patch:
diff --git a/raytracing/rays.py b/raytracing/rays.py
index 9129491..a424fad 100644
--- a/raytracing/rays.py
+++ b/raytracing/rays.py
@@ -207,6 +207,10 @@ class Rays:
     def load(self, filePath, append=False):
         with open(filePath, 'rb') as infile:
             loadedRays = pickle.Unpickler(infile).load()
+            if not isinstance(loadedRays, collections.Iterable):
+                raise IOError(f"{filePath} does not contain an iterable of Ray objects.")
+            if not all([isinstance(ray, Ray) for ray in loadedRays]):
+                raise IOError(f"{filePath} must contain only Ray objects.")
             if append and self._rays is not None:
                 self._rays.extend(loadedRays)
             else:


Test patch:
diff --git a/raytracing/tests/testsRays.py b/raytracing/tests/testsRays.py
index 33ac2e3..833ade3 100644
--- a/raytracing/tests/testsRays.py
+++ b/raytracing/tests/testsRays.py
@@ -69,7 +69,6 @@ class TestRays(unittest.TestCase):
         r = Rays(listOfRays)
         self.assertEqual(len(r), len(listOfRays))
 
-
     def testRaysGetRay(self):
         raysList = [Ray(), Ray(1), Ray(-1)]
         rays = Rays(raysList)
@@ -244,7 +243,6 @@ class TestRays(unittest.TestCase):
         self.assertIsNone(r._anglesHistogramParameters)
         self.assertIsNone(r._xValuesAnglesHistogram)
 
-
     def testAppendInvalidInput(self):
         rays = Rays()
         with self.assertRaises(TypeError):
@@ -258,7 +256,7 @@ class TestRaysSaveAndLoad(unittest.TestCase):
         self.fileName = 'testFile.pkl'
         with open(self.fileName, 'wb') as file:
             pickle.Pickler(file).dump(self.testRays.rays)
-        time.sleep(1)  # Make sure everything is ok
+        time.sleep(0.5)  # Make sure everything is ok
 
     def tearDown(self) -> None:
         if os.path.exists(self.fileName):
@@ -294,6 +292,33 @@ class TestRaysSaveAndLoad(unittest.TestCase):
         self.assertLoadNotFailed(rays)  # We don't append, we override
         self.assertListEqual(rays.rays, self.testRays.rays)
 
+    def testLoadWrongFileContent(self):
+        wrongObj = 7734
+        fileName = 'wrongObj.pkl'
+        with open(fileName, 'wb') as file:
+            pickle.Pickler(file).dump(wrongObj)
+        time.sleep(0.5)  # Make sure everything is ok
+
+        try:
+            with self.assertRaises(IOError):
+                Rays().load(fileName)
+        except AssertionError as exception:
+            self.fail(str(exception))
+        finally:
+            os.remove(fileName)
+
+        wrongIterType = [Ray(), Ray(1), [1, 1]]
+        with open(fileName, 'wb') as file:
+            pickle.Pickler(file).dump(wrongIterType)
+        time.sleep(0.5)
+        try:
+            with self.assertRaises(IOError):
+                Rays().load(fileName)
+        except AssertionError as exception:
+            self.fail(str(exception))
+        finally:
+            os.remove(fileName)
+
     def assertSaveNotFailed(self, rays: Rays, name: str, deleteNow: bool = True):
         try:
             rays.save(name)