Change some types of bad file handling for datalayer

chia/data_layer/data_layer_errors.py

@@ -9,6 +9,10 @@ class IntegrityError(Exception):
     pass
 
 
+class FileDownloadError(Exception):
+    pass
+
+
 def build_message_with_hashes(message: str, bytes_objects: Iterable[bytes]) -> str:
     return "\n".join([message, *[f"    {b.hex()}" for b in bytes_objects]])
 

chia/data_layer/download_data.py

@@ -11,6 +11,7 @@ from typing import List, Optional
 import aiohttp
 from typing_extensions import Literal
 
+from chia.data_layer.data_layer_errors import FileDownloadError
 from chia.data_layer.data_layer_util import NodeType, Root, SerializedNode, ServerInfo, Status
 from chia.data_layer.data_store import DataStore
 from chia.types.blockchain_format.sized_bytes import bytes32
@@ -65,7 +66,8 @@ async def insert_into_data_store_from_file(
                 cur_chunk = reader.read(size_to_read)
                 if cur_chunk is None or cur_chunk == b"":
                     if size_to_read < 4:
-                        raise Exception("Incomplete read of length.")
+                        # Treat as download error instead of integrity error
+                        raise FileDownloadError("Incomplete read of length.")
                     break
                 chunk += cur_chunk
             if chunk == b"":
@@ -77,7 +79,8 @@ async def insert_into_data_store_from_file(
                 size_to_read = size - len(serialize_nodes_bytes)
                 cur_chunk = reader.read(size_to_read)
                 if cur_chunk is None or cur_chunk == b"":
-                    raise Exception("Incomplete read of blob.")
+                    # Treat as download error instead of integrity error
+                    raise FileDownloadError("Incomplete read of blob.")
                 serialize_nodes_bytes += cur_chunk
             serialized_node = SerializedNode.from_bytes(serialize_nodes_bytes)
 
@@ -155,7 +158,8 @@ async def insert_from_delta_file(
         if downloader is None:
             # use http downloader
             if not await http_download(client_foldername, filename, proxy_url, server_info, timeout, log):
-                break
+                await data_store.server_misses_file(tree_id, server_info, timestamp)
+                return False
         else:
             log.info(f"Using downloader {downloader} for store {tree_id.hex()}.")
             async with aiohttp.ClientSession() as session:
@@ -163,7 +167,8 @@ async def insert_from_delta_file(
                     res_json = await response.json()
                     if not res_json["downloaded"]:
                         log.error(f"Failed to download delta file {filename} from {downloader}: {res_json}")
-                        break
+                        await data_store.server_misses_file(tree_id, server_info, timestamp)
+                        return False
 
         log.info(f"Successfully downloaded delta file {filename}.")
         try:
@@ -188,10 +193,13 @@ async def insert_from_delta_file(
             await data_store.received_correct_file(tree_id, server_info)
         except asyncio.CancelledError:
             raise
-        except Exception:
+        except Exception as e:
             target_filename = client_foldername.joinpath(filename)
             os.remove(target_filename)
-            await data_store.received_incorrect_file(tree_id, server_info, timestamp)
+            if isinstance(e, FileDownloadError):
+                await data_store.server_misses_file(tree_id, server_info, timestamp)
+            else:
+                await data_store.received_incorrect_file(tree_id, server_info, timestamp)
             await data_store.rollback_to_generation(tree_id, existing_generation - 1)
             raise
 

tests/core/data_layer/test_data_store.py

@@ -3,14 +3,15 @@ from __future__ import annotations
 import itertools
 import logging
 import re
+import sqlite3
 import statistics
 from pathlib import Path
 from random import Random
-from typing import Any, Awaitable, Callable, Dict, List, Set, Tuple
+from typing import Any, Awaitable, Callable, Dict, List, Set, Tuple, Type
 
 import pytest
 
-from chia.data_layer.data_layer_errors import NodeHashError, TreeGenerationIncrementingError
+from chia.data_layer.data_layer_errors import FileDownloadError, NodeHashError, TreeGenerationIncrementingError
 from chia.data_layer.data_layer_util import (
     DiffData,
     InternalNode,
@@ -1236,6 +1237,39 @@ async def test_data_server_files(data_store: DataStore, tree_id: bytes32, test_d
         generation += 1
 
 
+@pytest.mark.parametrize(
+    argnames=["test_bytes", "exception", "exception_string"],
+    argvalues=[
+        [b"\x00\x00\x00", FileDownloadError, "Incomplete read of length"],
+        [b"\x00\x00\x00\x01", FileDownloadError, "Incomplete read of blob"],
+        [
+            b"\x00\x00\x00\x0b\x01\x00\x00\x00\x012\x00\x00\x00\x012",
+            sqlite3.IntegrityError,
+            "FOREIGN KEY constraint failed",
+        ],
+        [b"\x00\x00\x00\x01\0x01", AssertionError, ""],
+    ],
+    ids=["bad file length", "bad blob length", "integrity error", "bad node serialization"],
+)
+@pytest.mark.asyncio
+async def test_unreadable_files(
+    data_store: DataStore,
+    tree_id: bytes32,
+    tmp_path: Path,
+    test_bytes: bytes,
+    exception: Type[Exception],
+    exception_string: str,
+) -> None:
+    fake_file = tmp_path.joinpath("fake.dat")
+    with open(fake_file, "wb") as binary_file:
+        binary_file.write(test_bytes)
+
+    fake_root_hash = bytes32.from_hexstr("7e2ce962061b5b4d17abba02468ebb39e1382ad21ffb12ed18b95e32f5d07dc2")
+
+    with pytest.raises(exception, match=exception_string):
+        await insert_into_data_store_from_file(data_store, tree_id, fake_root_hash, fake_file)
+
+
 @pytest.mark.asyncio
 async def test_pending_roots(data_store: DataStore, tree_id: bytes32) -> None:
     key = b"\x01\x02"