Bound and expire custodial upgrade challenges

src/msg/plugins/identity.py

@@ -255,6 +255,17 @@ async def custodial_upgrade_start(ctx,request,tx):
                 'request_id':request.request_id}
         verify(public,canonical(signed),decode(Signature,args['possession_proof']),
                purpose='custodial-upgrade-start')
+        # Keep challenge storage bounded even if a client repeatedly abandons
+        # upgrades.  Expired attempts no longer need to be recoverable, while a
+        # live attempt must be finished (or allowed to expire) before restarting.
+        # The no-op update serializes concurrent starts for the same subject.
+        tx.execute('UPDATE custodial_vault SET status=status WHERE subject=?',
+                   (subject.resource_id,),write=True)
+        tx.execute('DELETE FROM custodial_upgrades WHERE subject=? AND expires_at<=?',
+                   (subject.resource_id,wire(ctx.now)),write=True)
+        require(tx.one('''SELECT 1 FROM custodial_upgrades WHERE subject=?
+            AND status IN ('pending','pending_rewrap') AND expires_at>? LIMIT 1''',
+            (subject.resource_id,wire(ctx.now))) is None,'custodial_upgrade_in_progress')
         challenge_id=new_id('cupg')
         ephemeral=X25519PrivateKey.generate()
         expires=ctx.now+timedelta(seconds=300)

src/msg/workers/maintenance.py

@@ -34,7 +34,8 @@ async def purge_revisions(app, tx, resource, *, actor, request_id, reason):
 
 
 async def _cleanup(app,tx, *, scheduled):
-    counts={'expired_resources':0,'expired_transfers':0,'expired_query_sources':0}
+    counts={'expired_resources':0,'expired_transfers':0,'expired_query_sources':0,
+            'expired_custodial_upgrades':0}
     if not tx.setting('runtime_config',{}).get('cleanup_enabled',True):
         return counts
     now=app.clock()
@@ -111,6 +112,9 @@ async def _cleanup(app,tx, *, scheduled):
         tx.execute('DELETE FROM settings WHERE key=?',(key,),write=True)
         counts['expired_query_sources']+=1
     tx.execute('DELETE FROM email_challenges WHERE expires<=?',(wire(now),),write=True)
+    expired=tx.one('SELECT COUNT(*) FROM custodial_upgrades WHERE expires_at<=?',(wire(now),))[0]
+    tx.execute('DELETE FROM custodial_upgrades WHERE expires_at<=?',(wire(now),),write=True)
+    counts['expired_custodial_upgrades']=expired
     return counts
 
 

tests/test_custodial_upgrade.py

@@ -1,5 +1,6 @@
 """Custodial exit requires both new private keys and preserves recovery on failure."""
 import os
+from datetime import timedelta
 
 import pytest
 
@@ -8,7 +9,7 @@
 from msg.security.age_keys import generate_age_key
 from msg.security.crypto import Ed25519Signer,verify
 from msg.security.vault import client_upgrade_proof
-from test_service import call
+from test_service import NOW,call
 
 
 async def begin(app,subject,token,signer,recipient, *, handle='new-self',rid='upgrade-start'):
@@ -34,6 +35,30 @@ async def finish(app,subject,token,signer,age_identity,challenge, *, rid='upgrad
         'migration_ack':wire(acknowledgement)},subject=subject,token=token,rid=rid)
 
 
+@pytest.mark.asyncio
+async def test_only_one_live_custodial_upgrade_is_allowed(installed):
+    app,_=installed
+    created=await call(app,'identity.custodial_create',{'handle':'cust-bounded',
+        'nonce':b64(os.urandom(32))})
+    subject=created.data['subject_id']
+    token=(created.data['credential_id'],unb64(created.data['token']))
+    first=await begin(app,subject,token,Ed25519Signer.generate(),generate_age_key()[1],
+                      rid='bounded-first')
+    second=await begin(app,subject,token,Ed25519Signer.generate(),generate_age_key()[1],
+                       rid='bounded-second')
+    assert first.status=='ok'
+    assert second.status=='error' and second.error.code=='custodial_upgrade_in_progress'
+    async with app.metadata.transaction(write=False) as tx:
+        assert tx.one('SELECT COUNT(*) FROM custodial_upgrades WHERE subject=?',(subject,))[0]==1
+
+    app.clock=lambda: NOW+timedelta(seconds=301)
+    replacement=await begin(app,subject,token,Ed25519Signer.generate(),generate_age_key()[1],
+                            rid='bounded-replacement')
+    assert replacement.status=='ok'
+    async with app.metadata.transaction(write=False) as tx:
+        assert tx.one('SELECT COUNT(*) FROM custodial_upgrades WHERE subject=?',(subject,))[0]==1
+
+
 @pytest.mark.asyncio
 async def test_wrong_age_proof_fails_without_losing_old_token_or_vault(installed):
     app,_=installed

tests/test_operations_admin.py

@@ -48,6 +48,20 @@ async def test_policy_cleanup_is_explicit_records_tombstone_and_does_not_delete_
     assert reading.error.code=='resource_purged'
 
 
+@pytest.mark.asyncio
+async def test_cleanup_deletes_expired_custodial_upgrades(installed):
+    app,_=installed
+    async with app.metadata.transaction(write=True) as tx:
+        tx.execute('''INSERT INTO custodial_upgrades
+            (id,subject,credential_id,status,expires_at,challenge,body)
+            VALUES (?,?,?,?,?,?,?)''',('cupg_expired','u_test','cred_test','pending',
+            wire(NOW-timedelta(seconds=1)),'{}','{}'),write=True)
+    result=await run_maintenance(app,'cleanup_expired',scheduled=True)
+    assert result['expired_custodial_upgrades']==1
+    async with app.metadata.transaction(write=False) as tx:
+        assert tx.one('SELECT COUNT(*) FROM custodial_upgrades')[0]==0
+
+
 @pytest.mark.asyncio
 async def test_doctor_does_not_create_database_or_repair_bootstrap(installed,tmp_path):
     app,root=installed