Merge pull request #1015 from google/py-fix-1014

python bugfix: limit the number of bytes we read in case of an input with just many whitespaces

Author: reyammer

python/CHANGELOG.md

@@ -13,6 +13,7 @@ semver guidelines for more details about this.
 ## [Unreleased]
 
 - Mark python 3.13 as supported.
+- Bugfix: limit the number of bytes we read in case of an input with just many whitespaces. ([#1015](https://github.com/google/magika/pull/1015))
 
 
 ## [0.6.1] - 2025-03-19

python/src/magika/magika.py

@@ -875,7 +875,8 @@ def _get_result_or_features_from_bytes(
                 # If the n-th token is padding, then it means that,
                 # post-stripping, we do not have enough meaningful
                 # bytes.
-                result = self._get_result_from_few_bytes(content)
+                bytes_to_read = min(len(content), self._model_config.block_size)
+                result = self._get_result_from_few_bytes(content[0:bytes_to_read])
                 return result, None
 
             else:
@@ -938,8 +939,9 @@ def _get_result_or_features_from_stream(
                 # If the n-th token is padding, then it means that,
                 # post-stripping, we do not have enough meaningful
                 # bytes.
+                bytes_to_read = min(bytes_stream_size, self._model_config.block_size)
                 stream.seek(0)
-                content = stream.read()
+                content = stream.read(bytes_to_read)
                 result = self._get_result_from_few_bytes(content)
                 return result, None
 

python/tests/test_magika_python_module.py

@@ -222,6 +222,63 @@ def test_magika_module_with_python_and_non_python_content() -> None:
     assert res.prediction.output.label == ContentTypeLabel.TXT
 
 
+def test_magika_module_with_whitespaces() -> None:
+    m = Magika()
+
+    ws_nums = sorted(
+        {
+            1,
+            m._model_config.min_file_size_for_dl - 1,
+            m._model_config.min_file_size_for_dl,
+            m._model_config.min_file_size_for_dl + 1,
+            m._model_config.beg_size - 1,
+            m._model_config.beg_size,
+            m._model_config.beg_size + 1,
+            m._model_config.end_size - 1,
+            m._model_config.end_size,
+            m._model_config.end_size + 1,
+            m._model_config.beg_size + m._model_config.end_size - 1,
+            m._model_config.beg_size + m._model_config.end_size,
+            m._model_config.beg_size + m._model_config.end_size + 1,
+            m._model_config.beg_size + m._model_config.end_size + 1,
+            m._model_config.block_size - 1,
+            m._model_config.block_size,
+            m._model_config.block_size + 1,
+            2 * m._model_config.block_size - 1,
+            2 * m._model_config.block_size,
+            2 * m._model_config.block_size + 1,
+            4 * m._model_config.block_size - 1,
+            4 * m._model_config.block_size,
+            4 * m._model_config.block_size + 1,
+        }
+    )
+
+    for ws_num in ws_nums:
+        print(f"Calling indentify_bytes with {ws_num} whitespaces")
+        content = b" " * ws_num
+        res = m.identify_bytes(content)
+        assert (
+            res.ok
+            and res.dl.label == ContentTypeLabel.UNDEFINED
+            and res.output.label == ContentTypeLabel.TXT
+        )
+        res = m.identify_stream(io.BytesIO(content))
+        assert (
+            res.ok
+            and res.dl.label == ContentTypeLabel.UNDEFINED
+            and res.output.label == ContentTypeLabel.TXT
+        )
+        with tempfile.TemporaryDirectory() as td:
+            tf_path = Path(td) / "test.bin"
+            tf_path.write_bytes(content)
+            res = m.identify_path(tf_path)
+            assert (
+                res.ok
+                and res.dl.label == ContentTypeLabel.UNDEFINED
+                and res.output.label == ContentTypeLabel.TXT
+            )
+
+
 def test_magika_module_with_different_prediction_modes() -> None:
     model_dir = utils.get_default_model_dir()
     m = Magika(model_dir=model_dir, prediction_mode=PredictionMode.BEST_GUESS)