fix: use execute to consume keys used to publish

Author: PromiseFru

grpc_services/v3/utils.py

@@ -12,7 +12,7 @@
     X25519PublicKey,
 )
 from sqlalchemy import delete, insert, select
-from sqlalchemy.orm import Session, contains_eager
+from sqlalchemy.orm import Session
 
 from db import get_session
 from lib_relaysms_payload_specs.generated import relaysms_spec_payload as rrs
@@ -36,35 +36,33 @@ def get_keys_for_decryption(
     Fetch token and all necessary keys for decryption.
     Returns (Token, TokenHash, ss_kid, es_kid, es_kid_pk, ec_kid_pk).
     """
-    stmt = (
-        select(Token)
-        .join(Token.token_hash)
-        .join(TokenHash.server_keys)
-        .join(TokenHash.client_keys)
-        .where(Token.token_id == token_id_bytes)
-        .where(ServerEphemeralKey.key_index == key_id)
-        .where(ClientEphemeralKey.key_index == key_id)
-        .options(
-            contains_eager(Token.token_hash).contains_eager(TokenHash.server_keys),
-            contains_eager(Token.token_hash).contains_eager(TokenHash.client_keys),
-        )
-    )
-
-    token = session.scalars(stmt).first()
+    token = session.scalar(select(Token).where(Token.token_id == token_id_bytes))
     if not token:
-        raise ValueError("Token or associated keys not found")
+        raise ValueError("token not found")
 
-    token_hash_obj = token.token_hash
+    token_hash_obj = session.scalar(
+        select(TokenHash).where(TokenHash.id == token.token_hash_id)
+    )
     if not token_hash_obj:
-        raise ValueError("Token hash not found")
+        raise ValueError("token hash not found")
 
-    if not token_hash_obj.server_keys:
-        raise ValueError(f"Server ephemeral key not found for kid {key_id}")
-    se_key = token_hash_obj.server_keys[0]
+    se_key = session.scalar(
+        select(ServerEphemeralKey).where(
+            ServerEphemeralKey.token_hash_id == token_hash_obj.id,
+            ServerEphemeralKey.key_index == key_id,
+        )
+    )
+    if not se_key:
+        raise ValueError(f"server ephemeral key not found: kid={key_id}")
 
-    if not token_hash_obj.client_keys:
-        raise ValueError(f"Client ephemeral key not found for kid {key_id}")
-    ce_key = token_hash_obj.client_keys[0]
+    ce_key = session.scalar(
+        select(ClientEphemeralKey).where(
+            ClientEphemeralKey.token_hash_id == token_hash_obj.id,
+            ClientEphemeralKey.key_index == key_id,
+        )
+    )
+    if not ce_key:
+        raise ValueError(f"client ephemeral key not found: kid={key_id}")
 
     ss_kid = get_private_key(key_id, session).private_bytes_raw()
 

rest_services/v1/routes.py

@@ -188,9 +188,16 @@ def create_publications(
                     status_code=400, detail="Failed to deserialize payload"
                 ) from exc
 
+            t_id = payload.get_t_id()
+            if t_id is None:
+                logger.error("payload is missing token ID")
+                raise HTTPException(
+                    status_code=400, detail="Payload is missing token ID"
+                )
+
             with get_session() as db:
                 publish_content(
-                    token_id=struct.pack("<I", payload.get_t_id()),
+                    token_id=struct.pack("<I", t_id),
                     key_id=payload.get_kid(),
                     len_att=payload.get_len_att(),
                     content_ciphertext=payload.get_content(),
@@ -210,6 +217,7 @@ def create_publications(
                     db=db,
                 )
                 if joined is None:
+                    logger.info("Segment stored. Waiting for remaining segments.")
                     return PublishContentResponse(
                         message="Segment stored. Waiting for remaining segments."
                     )
@@ -229,4 +237,5 @@ def create_publications(
                 detail=f"Payload type {payload_type!r} is not supported.",
             )
 
+    logger.info("Content published successfully.")
     return PublishContentResponse(message="Content published successfully")

rest_services/v1/services.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: GPL-3.0-only
 """v1 Services for the REST API."""
 
+from sqlalchemy import delete
 from sqlalchemy.orm import Session
 
 from grpc_services.v3.utils import (
@@ -10,11 +11,13 @@
 )
 from lib_relaysms_payload_specs.generated import relaysms_spec_payload as rrs
 from logutils import get_logger
+from models.client_ephemeral_key import ClientEphemeralKey
 from models.payload_segment import create_if_not_exists as create_segment
 from models.payload_segment import get_all_data
 from models.payload_session import create as create_session
 from models.payload_session import delete as delete_session
 from models.payload_session import get_by_sender_and_session
+from models.server_ephemeral_key import ServerEphemeralKey
 from models.server_identity_key import mark_key_used as mark_ss_kid_used
 from models.token import update_token_data
 from models.token_hash import TokenHash
@@ -28,7 +31,6 @@ def _get_adapter_params(
     token_data: dict, content: rrs.V1ContentsContainer, *, extras: dict | None = None
 ) -> dict:
     cat_id = content.get_cat_id()
-    print(">>>>>>>>> ATTACHMENT:", content.get_attachment())
     extra_params = extras or {}
 
     match cat_id:
@@ -54,8 +56,18 @@ def _get_adapter_params(
 
 
 def _consume_used_keys(token_hash: TokenHash, key_index: int, session: Session) -> None:
-    session.delete(token_hash.server_keys[0])
-    session.delete(token_hash.client_keys[0])
+    session.execute(
+        delete(ServerEphemeralKey).where(
+            ServerEphemeralKey.token_hash_id == token_hash.id,
+            ServerEphemeralKey.key_index == key_index,
+        )
+    )
+    session.execute(
+        delete(ClientEphemeralKey).where(
+            ClientEphemeralKey.token_hash_id == token_hash.id,
+            ClientEphemeralKey.key_index == key_index,
+        )
+    )
     mark_ss_kid_used(key_index, session)