Added balance pre-check on file pins (#799)

nesitor · Andres D. Molins · web-flow · commit fdbb7304e0e8 · 2025-06-02T12:45:18.000+02:00
* Feature: To avoid saving un-needed files on IPFS, we added a balance pre-check on STORE messages before doing the message processing.

* Split `verify_and_fetch` method in two and do the pre-balance-check before downloading it, only in IPFS case.

* Fix: Solve missing references to old method `verify_and_fetch` and replaced by the new one `verify_message`.

* Fix: Solve way of calculating the message_cost.

* Fix: Solved errors and code quality issues.

* Fix: Changes way of instantiate the model class.

* Fix: Allow small files to be pinned by "free"

* Fix: Added types compatibility

* Added missing test cases for IPFS and content handler.

* Fix: Solved test issues and applied good typing.

* Fix: Removed redundant `if`.

---------

Co-authored-by: Andres D. Molins &lt;nesitor@gmail.com&gt;
diff --git a/src/aleph/handlers/content/content_handler.py b/src/aleph/handlers/content/content_handler.py
@@ -48,6 +48,17 @@ async def process(self, session: DbSession, messages: List[MessageDb]) -> None:
         """
         pass
 
+    async def pre_check_balance(self, session: DbSession, message: MessageDb) -> None:
+        """
+        Checks whether the user has enough Aleph tokens before processing the message.
+
+        Raises InsufficientBalanceException if the balance of the user is too low.
+
+        :param session: DB session.
+        :param message: Pending Message being processed.
+        """
+        pass
+
     async def check_balance(
         self, session: DbSession, message: MessageDb
     ) -> List[AccountCostsDb] | None:
diff --git a/src/aleph/handlers/content/store.py b/src/aleph/handlers/content/store.py
@@ -8,6 +8,7 @@
 import asyncio
 import datetime as dt
 import logging
+from decimal import Decimal
 from typing import List, Optional, Set
 
 import aioipfs
@@ -31,6 +32,7 @@
 from aleph.db.models.account_costs import AccountCostsDb
 from aleph.exceptions import AlephStorageException, UnknownHashError
 from aleph.handlers.content.content_handler import ContentHandler
+from aleph.schemas.cost_estimation_messages import CostEstimationStoreContent
 from aleph.services.cost import calculate_storage_size, get_total_and_detailed_costs
 from aleph.storage import StorageService
 from aleph.toolkit.constants import MAX_UNAUTHENTICATED_UPLOAD_FILE_SIZE, MiB
@@ -206,6 +208,67 @@ async def fetch_related_content(
             size=size,
         )
 
+    async def pre_check_balance(self, session: DbSession, message: MessageDb):
+        content = _get_store_content(message)
+        assert isinstance(content, StoreContent)
+
+        if are_store_and_program_free(message):
+            return None
+
+        # This check is essential to ensure that files are not added to the system
+        # on the current node when the configuration disables storing of files.
+        config = get_config()
+        ipfs_enabled = config.ipfs.enabled.value
+
+        current_balance = get_total_balance(session=session, address=content.address)
+        current_cost = get_total_cost_for_address(
+            session=session, address=content.address
+        )
+
+        engine = content.item_type
+        # Initially only do that balance pre-check for ipfs files.
+        if engine == ItemType.ipfs and ipfs_enabled:
+            ipfs_byte_size = await self.storage_service.ipfs_service.get_ipfs_size(
+                content.item_hash
+            )
+            if ipfs_byte_size:
+                storage_mib = Decimal(ipfs_byte_size / MiB)
+
+                # Allow users to pin small files
+                if storage_mib and storage_mib <= (
+                    MAX_UNAUTHENTICATED_UPLOAD_FILE_SIZE / MiB
+                ):
+                    LOGGER.debug(
+                        f"Cost for {message.item_hash} supposed to be free as size is {storage_mib}"
+                    )
+                    return None
+
+                computable_content_data = {
+                    **content.model_dump(),
+                    "estimated_size_mib": int(storage_mib),
+                }
+                computable_content = CostEstimationStoreContent.model_validate(
+                    computable_content_data
+                )
+
+                message_cost, _ = get_total_and_detailed_costs(
+                    session, computable_content, message.item_hash
+                )
+            else:
+                message_cost = Decimal(0)
+        else:
+            message_cost = Decimal(0)
+
+        required_balance = current_cost + message_cost
+
+        if current_balance < required_balance:
+            raise InsufficientBalanceException(
+                balance=current_balance,
+                required_balance=required_balance,
+            )
+
+        return None
+
     async def check_balance(
         self, session: DbSession, message: MessageDb
     ) -> List[AccountCostsDb]:
diff --git a/src/aleph/handlers/message_handler.py b/src/aleph/handlers/message_handler.py
@@ -376,14 +376,12 @@ async def insert_costs(
             insert_stmt = make_costs_upsert_query(costs)
             session.execute(insert_stmt)
 
-    async def verify_and_fetch(
-        self, session: DbSession, pending_message: PendingMessageDb
-    ) -> MessageDb:
+    async def verify_message(self, pending_message: PendingMessageDb) -> MessageDb:
         await self.verify_signature(pending_message=pending_message)
         validated_message = await self.fetch_pending_message(
             pending_message=pending_message
         )
-        await self.fetch_related_content(session=session, message=validated_message)
+
         return validated_message
 
     async def process(
@@ -430,11 +428,16 @@ async def process(
                 error_code=ErrorCode.FORGOTTEN_DUPLICATE,
             )
 
-        message = await self.verify_and_fetch(
-            session=session, pending_message=pending_message
-        )
+        # First check the message content and verify it
+        message = await self.verify_message(pending_message=pending_message)
 
+        # Do a balance pre-check to avoid saving related data
         content_handler = self.get_content_handler(message.type)
+        await content_handler.pre_check_balance(session=session, message=message)
+
+        # Fetch related content like the IPFS associated file
+        await self.fetch_related_content(session=session, message=message)
+
         await content_handler.check_dependencies(session=session, message=message)
         await content_handler.check_permissions(session=session, message=message)
         costs = await content_handler.check_balance(session=session, message=message)
diff --git a/src/aleph/jobs/fetch_pending_messages.py b/src/aleph/jobs/fetch_pending_messages.py
@@ -55,8 +55,8 @@ def __init__(
     async def fetch_pending_message(self, pending_message: PendingMessageDb):
         with self.session_factory() as session:
             try:
-                message = await self.message_handler.verify_and_fetch(
-                    session=session, pending_message=pending_message
+                message = await self.message_handler.verify_message(
+                    pending_message=pending_message
                 )
                 session.execute(
                     make_pending_message_fetched_statement(
diff --git a/src/aleph/services/cost.py b/src/aleph/services/cost.py
@@ -152,10 +152,10 @@ def _get_product_price_type(
     settings: Settings,
     price_aggregate: Union[AggregateDb, dict],
 ) -> ProductPriceType:
-    if isinstance(content, StoreContent):
+    if isinstance(content, (StoreContent, CostEstimationStoreContent)):
         return ProductPriceType.STORAGE
 
-    if isinstance(content, ProgramContent):
+    if isinstance(content, (ProgramContent, CostEstimationProgramContent)):
         is_on_demand = not content.on.persistent
         return (
             ProductPriceType.PROGRAM
@@ -225,7 +225,7 @@ def _get_nb_compute_units(
 def _get_compute_unit_multiplier(content: CostComputableContent) -> int:
     compute_unit_multiplier = 1
     if (
-        isinstance(content, ProgramContent)
+        isinstance(content, (ProgramContent, CostEstimationProgramContent))
         and not content.on.persistent
         and content.environment.internet
     ):
@@ -289,7 +289,7 @@ def _get_execution_volumes_costs(
 ) -> List[AccountCostsDb]:
     volumes: List[RefVolume | SizedVolume] = []
 
-    if isinstance(content, InstanceContent):
+    if isinstance(content, (InstanceContent, CostEstimationInstanceContent)):
         volumes.append(
             SizedVolume(
                 CostType.EXECUTION_INSTANCE_VOLUME_ROOTFS,
@@ -298,7 +298,7 @@ def _get_execution_volumes_costs(
             )
         )
 
-    elif isinstance(content, ProgramContent):
+    elif isinstance(content, (ProgramContent, CostEstimationProgramContent)):
         if (
             isinstance(content, CostEstimationProgramContent)
             and content.code.estimated_size_mib
@@ -365,7 +365,7 @@ def _get_execution_volumes_costs(
         # or with same values for different volumes causing unique key constraint errors
         name_prefix = f"#{i}"
 
-        if isinstance(volume, ImmutableVolume):
+        if isinstance(volume, (ImmutableVolume, CostEstimationImmutableVolume)):
             name = (
                 f"{name_prefix}:{volume.mount or CostType.EXECUTION_VOLUME_INMUTABLE}"
             )
@@ -582,7 +582,7 @@ def get_detailed_costs(
     settings = settings or _get_settings(session)
     pricing = pricing or _get_product_price(session, content, settings)
 
-    if isinstance(content, StoreContent):
+    if isinstance(content, (StoreContent, CostEstimationStoreContent)):
         return _calculate_storage_costs(session, content, pricing, item_hash)
     else:
         return _calculate_executable_costs(session, content, pricing, item_hash)
diff --git a/src/aleph/services/ipfs/service.py b/src/aleph/services/ipfs/service.py
@@ -57,6 +57,73 @@ async def get_public_address(self):
             if "127.0.0.1" in address and "/tcp" in address and "/p2p" in address:
                 return address.replace("127.0.0.1", public_ip)
 
+    async def get_ipfs_size(
+        self, hash: str, timeout: int = 1, tries: int = 1
+    ) -> Optional[int]:
+        try_count = 0
+        result = None
+        while (result is None) and (try_count < tries):
+            try_count += 1
+            try:
+                dag_node = await asyncio.wait_for(
+                    self.ipfs_client.dag.get(hash), timeout=timeout
+                )
+                if isinstance(dag_node, dict):
+                    if "Data" in dag_node and isinstance(dag_node["Data"], dict):
+                        # This is the common structure for UnixFS nodes after aioipfs parsing
+                        if "filesize" in dag_node["Data"]:
+                            result = dag_node["Data"]["filesize"]
+                        elif (
+                            "Tsize" in dag_node["Data"]
+                        ):  # Less common, but good to check
+                            result = dag_node["Data"]["Tsize"]
+                        elif (
+                            "Tsize" in dag_node
+                        ):  # Sometimes it might be at the top level directly
+                            result = dag_node["Tsize"]
+                    elif "Links" in dag_node and isinstance(dag_node["Links"], list):
+                        total_size = 0
+                        for link in dag_node["Links"]:
+                            # In case it's a link list, get the Tsize property if exists
+                            if "Tsize" in link and isinstance(link["Tsize"], int):
+                                total_size += link["Tsize"]
+                            else:
+                                LOGGER.error(
+                                    f"Error: CID {hash} did not return a list structure. Type: {type(link)}"
+                                )
+                        result = total_size
+
+                    elif (
+                        "Size" in dag_node
+                    ):  # Occasionally, 'Size' might refer to total size, but often it's block
+                        # size
+                        result = dag_node["Size"]
+                else:
+                    # For raw blocks, dag_node might be bytes. block.stat is better for those.
+                    # For other codecs, the structure will vary.
+                    LOGGER.info(
+                        f"Warning: CID {hash} did not return a dictionary structure. Type: {type(dag_node)}"
+                    )
+                    block_stat = await asyncio.wait_for(
+                        self.ipfs_client.block.stat(hash), timeout=timeout
+                    )
+                    result = block_stat["Size"]
+            except aioipfs.APIError:
+                result = None
+                await asyncio.sleep(0.5)
+                continue
+            except asyncio.TimeoutError:
+                result = None
+                await asyncio.sleep(0.5)
+            except (
+                concurrent.futures.CancelledError,
+                aiohttp.client_exceptions.ClientConnectorError,
+            ):
+                try_count -= 1  # do not count as a try.
+                await asyncio.sleep(0.1)
+
+        return result
+
     async def get_ipfs_content(
         self, hash: str, timeout: int = 1, tries: int = 1
     ) -> Optional[bytes]:
diff --git a/tests/message_processing/test_process_stores.py b/tests/message_processing/test_process_stores.py
diff --git a/tests/services/test_ipfs_service.py b/tests/services/test_ipfs_service.py
diff --git a/tests/test_network.py b/tests/test_network.py
