ios_client

chat_core.py

@@ -1044,30 +1044,34 @@ class ChatClient:
     # ------------------------------------------------------------------
 
     async def register(self, username: str, password: str, email: str) -> tuple[bool, str]:
-        """Register user. Generates RSA + Ed25519 + prekeys."""
+        """Register user. Generates RSA + Ed25519 in memory (saved to disk
+        only after server confirms registration via confirm_registration)."""
         self.username = username
         self.email = email
         pwd_bytes = bytearray(password.encode("utf-8")) if password else None
 
         try:
-            # RSA keys for login
-            priv, pub, err = load_keys(email, password=bytes(pwd_bytes) if pwd_bytes else None)
+            pwd = bytes(pwd_bytes) if pwd_bytes else None
+            # Try loading existing keys (previous successful registration)
+            priv, pub, err = load_keys(email, password=pwd)
             if priv is None:
                 priv, pub = generate_rsa_keypair()
-                save_keys(email, priv, pub, password=bytes(pwd_bytes) if pwd_bytes else None)
             self.private_key = priv
             self.public_key = pub
 
-            # Ed25519 identity keys
-            ed_priv, ed_pub = _load_identity_keys(email, password=bytes(pwd_bytes) if pwd_bytes else None)
+            try:
+                ed_priv, ed_pub = _load_identity_keys(email, password=pwd)
+            except Exception:
+                ed_priv, ed_pub = None, None
             if ed_priv is None:
                 ed_priv, ed_pub = generate_identity_keypair()
-                _save_identity_keys(email, ed_priv, ed_pub, password=bytes(pwd_bytes) if pwd_bytes else None)
             self.identity_private = ed_priv
             self.identity_public = ed_pub
             self._cache_key = derive_self_encryption_key(ed_priv)
             self._local_key = derive_local_storage_key(ed_priv)
-            self._load_verification_stores()
+
+            # Store password for saving keys after confirm
+            self._reg_password = pwd
         finally:
             if pwd_bytes:
                 pwd_bytes[:] = b'\x00' * len(pwd_bytes)
@@ -1100,6 +1104,7 @@ class ChatClient:
                 **extra_fields,
             )
         if start["status"] != "ok":
+            self._reg_password = None
             return False, start["data"]["message"]
         code = start["data"].get("code")
         if code:
@@ -1109,6 +1114,12 @@ class ChatClient:
     async def confirm_registration(self, email: str, username: str, code: str) -> tuple[bool, str]:
         confirm = await self.send_and_recv("register_confirm", email=email, code=code)
         if confirm["status"] == "ok":
+            # Registration confirmed — NOW save keys to disk
+            pwd = getattr(self, "_reg_password", None)
+            save_keys(email, self.private_key, self.public_key, password=pwd)
+            _save_identity_keys(email, self.identity_private, self.identity_public, password=pwd)
+            self._reg_password = None
+            self._load_verification_stores()
             # Upload prekeys immediately after registration
             await self._generate_and_upload_prekeys()
             return True, f"Registered as '{username}' (ID: {confirm['data']['user_id']})"