Limit memory reads for mmaped files (prometheus#329)

bloodearnest · brian-brazil · commit 38e9f48a8860 · 2018-10-18T12:53:35.000+01:00
* Add bounds check to mmap reads

Signed-off-by: Simon Davy &lt;simon.davy@canonical.com&gt;
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -562,6 +562,7 @@ class _MmapedDict(object):
 
     def __init__(self, filename, read_mode=False):
         self._f = open(filename, 'rb' if read_mode else 'a+b')
+        self._fname = filename
         if os.fstat(self._f.fileno()).st_size == 0:
             self._f.truncate(_INITIAL_MMAP_SIZE)
         self._capacity = os.fstat(self._f.fileno()).st_size
@@ -607,6 +608,10 @@ def _read_all_values(self):
 
         while pos < used:
             encoded_len = _unpack_integer(data, pos)[0]
+            # check we are not reading beyond bounds
+            if encoded_len + pos > used:
+                msg = 'Read beyond file size detected, %s is corrupted.'
+                raise RuntimeError(msg % self._fname)
             pos += 4
             encoded = unpack_from(('%ss' % encoded_len).encode(), data, pos)[0]
             padded_len = encoded_len + (8 - (encoded_len + 4) % 8)
diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py
@@ -296,6 +296,13 @@ def test_multi_expansion(self):
             [('abc', 42.0), (key, 123.0), ('def', 17.0)],
             list(self.d.read_all_values()))
 
+    def test_corruption_detected(self):
+        self.d.write_value('abc', 42.0)
+        # corrupt the written data
+        self.d._m[8:16] = b'somejunk'
+        with self.assertRaises(RuntimeError):
+            list(self.d.read_all_values())
+
     def tearDown(self):
         os.unlink(self.tempfile)
 
