feat: restrict storage

client_manager.py

@@ -11,7 +11,10 @@ from .crud import (
     get_config_for_all_active_relays,
     get_event,
     get_events,
+    get_prunable_events,
+    get_storage_for_public_key,
     mark_events_deleted,
+    prune_old_events,
 )
 from .models import ClientConfig, NostrEvent, NostrEventType, NostrFilter, RelayConfig
 
@@ -156,6 +159,12 @@ class NostrClientConnection:
             await self._send_msg(resp_nip20)
             return None
 
+        valid, message = await self._validate_storage(e)
+        if not valid:
+            resp_nip20 += [valid, message]
+            await self._send_msg(resp_nip20)
+            return None
+
         try:
             if e.is_replaceable_event():
                 await delete_events(
@@ -237,6 +246,29 @@ class NostrClientConnection:
 
         return True, ""
 
+    async def _validate_storage(self, e: NostrEvent) -> Tuple[bool, str]:
+        if self.client_config.free_storage_value == 0:
+            if not self.client_config.is_paid_relay:
+                return False, "Cannot write event, relay is read-only"
+            # todo: handeld paid paid plan
+            return True, "Temp OK"
+            
+
+        stored_bytes = await get_storage_for_public_key(self.relay_id, e.pubkey)
+        if self.client_config.is_paid_relay:
+            # todo: handeld paid paid plan
+            return True, "Temp OK"
+            
+        if (stored_bytes + e.size_bytes) <= self.client_config.free_storage_bytes_value:
+             return True, ""
+        
+        if self.client_config.full_storage_action == "block":
+            return False, f"Cannot write event, no more storage available for public key: '{e.pubkey}'"
+        
+        await prune_old_events(self.relay_id, e.pubkey, e.size_bytes)
+
+        return True, ""
+
     def _exceeded_max_events_per_second(self) -> bool:
         if self.client_config.max_events_per_second == 0:
             return False

crud.py

@@ -1,7 +1,5 @@
 import json
-from typing import Any, List, Optional
-
-from lnbits.helpers import urlsafe_short_hash
+from typing import Any, List, Optional, Tuple
 
 from . import db
 from .models import NostrEvent, NostrFilter, NostrRelay, RelayConfig
@@ -121,13 +119,24 @@ async def get_event(relay_id: str, id: str) -> Optional[NostrEvent]:
 
 async def get_storage_for_public_key(relay_id: str, pubkey: str) -> int:
     """Returns the storage space in bytes for all the events of a public key. Deleted events are also counted"""
-    
+
     row = await db.fetchone("SELECT SUM(size) FROM nostrrelay.events WHERE relay_id = ? AND pubkey = ?", (relay_id, pubkey,))
     if not row:
         return 0
 
-    return row["sum"]
+    return round(row["sum"])
 
+async def get_prunable_events(relay_id: str, pubkey: str) -> List[Tuple[str, int]]:
+    """ Return the oldest 10 000 events. Only the `id` and the size are returned, so the data size should be small"""
+    query = """
+            SELECT id, size FROM nostrrelay.events
+            WHERE relay_id = ? AND pubkey = ?
+            ORDER BY created_at ASC LIMIT 10000
+        """
+
+    rows = await db.fetchall(query, (relay_id, pubkey))
+
+    return [(r["id"], r["size"]) for r in rows]
 
 
 async def mark_events_deleted(relay_id: str,  filter: NostrFilter):
@@ -146,6 +155,20 @@ async def delete_events(relay_id: str,  filter: NostrFilter):
     await db.execute(query, tuple(values))
     #todo: delete tags
 
+async def prune_old_events(relay_id: str, pubkey: str, space_to_regain: int):
+    prunable_events = await get_prunable_events(relay_id, pubkey)
+    prunable_event_ids = []
+    size = 0
+
+    for pe in prunable_events:
+        prunable_event_ids.append(pe[0])
+        size += pe[1]
+
+        if size > space_to_regain:
+            break
+
+    await delete_events(relay_id, NostrFilter(ids=prunable_event_ids))
+
 
 async def delete_all_events(relay_id: str):
     query = "DELETE from nostrrelay.events WHERE relay_id = ?"

models.py

@@ -24,7 +24,8 @@ class ClientConfig(BaseModel):
     created_at_seconds_future = Field(0, alias="createdAtSecondsFuture")
 
     
-    free_storage_value = Field("1", alias="freeStorageValue")
+    is_paid_relay = Field(False, alias="isPaidRelay")
+    free_storage_value = Field(1, alias="freeStorageValue")
     free_storage_unit = Field("MB", alias="freeStorageUnit")
     full_storage_action = Field("prune", alias="fullStorageAction")
     
@@ -48,11 +49,17 @@ class ClientConfig(BaseModel):
     def created_at_in_future(self) -> int:
         return self.created_at_days_future * 86400 + self.created_at_hours_future * 3600 + self.created_at_minutes_future * 60 + self.created_at_seconds_future
 
+    @property
+    def free_storage_bytes_value(self):
+        value = self.free_storage_value * 1024
+        if self.free_storage_unit == "MB":
+            value *= 1024
+        return value
+
     class Config:
         allow_population_by_field_name = True
 
 class RelayConfig(ClientConfig):
-    is_paid_relay = Field(False, alias="isPaidRelay")
     wallet = Field("")
     cost_to_join = Field(0, alias="costToJoin")
     free_storage = Field(0, alias="freeStorage")