feat: support contract upgrades in batches

Cargo.toml

@@ -28,6 +28,7 @@ members = [
     "crates/tee-authority",
     "crates/tee-context",
     "crates/tee-launcher",
+    "crates/test-large-contract",
     "crates/test-migration-contract",
     "crates/test-parallel-contract",
     "crates/test-port-allocator",

crates/contract/README.md

@@ -308,7 +308,11 @@ These functions require the caller to be a participant or candidate.
 | `vote_pk(key_event_id: KeyEventId, public_key: PublicKey)`                          | For Initializing state only. Votes for the public key for the given generation attempt; if enough votes are collected, transitions to the next domain to generate a key for, or if all domains are completed, transitions into Running. | `Result<(), Error>`       | TBD             | TBD                |
 | `vote_reshared(key_event_id: KeyEventId)`                                           | For Resharing state only. Votes for the success of the given resharing attempt; if enough votes are collected, transitions to the next domain to reshare for, or if all domains are completed, transitions into Running.                | `Result<(), Error>`       | TBD             | TBD                |
 | `vote_cancel_keygen(next_domain_id: u64)`                                           | For Initializing state only. Votes to cancel the key generation (identified by the next_domain_id) and revert to the Running state.                                                                                                     | `Result<(), Error>`       | TBD             | TBD                |
-| `propose_update(args: ProposeUpdateArgs)`                                           | Proposes an update to the contract, requiring an attached deposit.                                                                                                                                                                      | `Result<UpdateId, Error>` | TBD             | TBD                |
+| `propose_update(args: ProposeUpdateArgs)`                                           | Proposes a configuration update, requiring an attached deposit. Contract-code updates are proposed via the chunked upload flow below, not this method.                                                                                   | `Result<UpdateId, Error>` | TBD             | TBD                |
+| `start_contract_upload(args: StartContractUploadArgs)`                              | Begins a chunked contract-code upload by declaring the total size. Each voter may have at most one open upload at a time.                                                                                                                | `Result<(), Error>`       | TBD             | TBD                |
+| `upload_contract_chunk(args: UploadContractChunkArgs)`                              | Appends a chunk of contract code to the caller's in-progress upload, requiring a deposit covering the chunk's storage cost.                                                                                                              | `Result<(), Error>`       | TBD             | TBD                |
+| `finalize_contract_upload()`                                                        | Finalizes a completed chunked upload and registers it as a contract-code proposal, returning its `UpdateId`.                                                                                                                             | `Result<UpdateId, Error>` | TBD             | TBD                |
+| `clear_staged_contract()`                                                           | Abandons the caller's in-progress chunked upload and refunds the accumulated chunk-storage deposits.                                                                                                                                     | `Result<(), Error>`       | TBD             | TBD                |
 | `vote_update(id: UpdateId)`                                                         | Votes on a proposed update. If the threshold is met, the update is executed.                                                                                                                                                            | `Result<bool, Error>`     | TBD             | TBD                |
 | `submit_participant_info(proposed_participant_attestation: Attestation, tls_public_key: Ed25519PublicKey)` | Submits the tee participant info for a potential candidate. c.f. TEE section | `Result<(), Error>` | TBD | TBD |
 

crates/contract/src/lib.rs

@@ -48,7 +48,10 @@ use crate::{
     state::ContractNotInitialized,
     storage_keys::StorageKey,
     tee::tee_state::{TeeQuoteStatus, TeeState},
-    update::{ProposeUpdateArgs, ProposedUpdates, Update, UpdateId},
+    update::{
+        ProposedConfigUpdateArgs, ProposedUpdates, StagedContractUpload, StartContractUploadArgs,
+        Update, UpdateId, UploadContractChunkArgs,
+    },
 };
 use config::Config;
 use crypto_shared::{
@@ -151,6 +154,16 @@ pub struct MpcContract {
     pending_ckd_requests: LookupMap<CKDRequest, Vec<YieldIndex>>,
     pending_verify_foreign_tx_requests: LookupMap<VerifyForeignTransactionRequest, Vec<YieldIndex>>,
     proposed_updates: ProposedUpdates,
+    /// Per-account metadata for in-progress chunked contract uploads. Each voter
+    /// may have at most one open upload at a time. See [`start_contract_upload`],
+    /// [`upload_contract_chunk`], [`finalize_contract_upload`],
+    /// [`clear_staged_contract`].
+    staged_uploads: IterableMap<AccountId, StagedContractUpload>,
+    /// Chunk bytes for in-progress uploads, keyed by `(uploader, chunk_index)`.
+    /// Lives in a separate `LookupMap` (rather than inside `StagedContractUpload`)
+    /// so appending a chunk only writes the new chunk's bytes — not the entire
+    /// accumulated blob.
+    staged_chunks: LookupMap<(AccountId, usize), Vec<u8>>,
     node_foreign_chain_support: SupportedForeignChainsByNode,
     config: Config,
     tee_state: TeeState,
@@ -1195,16 +1208,15 @@ impl MpcContract {
             .vote_abort_key_event_instance(key_event_id)
     }
 
-    /// Propose update to either code or config, but not both of them at the same time.
+    /// Propose update to the config.
     #[payable]
     #[handle_result]
-    pub fn propose_update(
+    pub fn propose_config_update(
         &mut self,
-        #[serializer(borsh)] args: ProposeUpdateArgs,
+        #[serializer(borsh)] args: ProposedConfigUpdateArgs,
     ) -> Result<UpdateId, Error> {
-        // Only voters can propose updates:
         let proposer = self.voter_or_panic();
-        let update: Update = args.try_into()?;
+        let update: Update = args.into();
 
         let attached = env::attached_deposit();
         let required = ProposedUpdates::required_deposit(&update);
@@ -1234,6 +1246,169 @@ impl MpcContract {
         Ok(id)
     }
 
+    /// Begin a chunked contract-code upload.
+    #[payable]
+    #[handle_result]
+    pub fn start_contract_upload(
+        &mut self,
+        #[serializer(borsh)] args: StartContractUploadArgs,
+    ) -> Result<(), Error> {
+        let caller = self.voter_or_panic();
+
+        if self.staged_uploads.contains_key(&caller) {
+            return Err(InvalidParameters::MalformedPayload {
+                reason:
+                    "caller already has a staged upload; call clear_staged_contract to abandon it first"
+                        .into(),
+            }
+            .into());
+        }
+
+        // Track the deposit attached to `start` so it is refunded on
+        // `clear_staged_contract`, matching `StagedContractUpload::deposited`'s
+        // documented contract (the sum of `start` + `upload_chunk` deposits).
+        let mut staged = StagedContractUpload::new(args.total_size);
+        staged.deposited = env::attached_deposit();
+        self.staged_uploads.insert(caller, staged);
+
+        Ok(())
+    }
+
+    /// Append a chunk of contract code to the caller's in-progress upload.
+    #[payable]
+    #[handle_result]
+    pub fn upload_contract_chunk(
+        &mut self,
+        #[serializer(borsh)] args: UploadContractChunkArgs,
+    ) -> Result<(), Error> {
+        let caller = self.voter_or_panic();
+
+        let attached = env::attached_deposit();
+        let chunk_len = args.data.len();
+        let required = StagedContractUpload::required_deposit_for_bytes(chunk_len);
+        if attached < required {
+            return Err(InvalidParameters::InsufficientDeposit {
+                attached: attached.as_yoctonear(),
+                required: required.as_yoctonear(),
+            }
+            .into());
+        }
+
+        let staged =
+            self.staged_uploads
+                .get_mut(&caller)
+                .ok_or(InvalidParameters::MalformedPayload {
+                    reason: "no staged upload found; call start_contract_upload first".into(),
+                })?;
+
+        let chunk_index = staged.record_chunk(chunk_len)?;
+        staged.deposited = staged.deposited.saturating_add(attached);
+
+        // Store chunk bytes in a separate map so the metadata write is small.
+        self.staged_chunks.insert((caller, chunk_index), args.data);
+
+        Ok(())
+    }
+
+    /// Finalize the caller's chunked upload and register it as a contract-code
+    /// proposal.
+    #[handle_result]
+    pub fn finalize_contract_upload(&mut self) -> Result<UpdateId, Error> {
+        let caller = self.voter_or_panic();
+
+        let staged =
+            self.staged_uploads
+                .remove(&caller)
+                .ok_or(InvalidParameters::MalformedPayload {
+                    reason: "no staged upload found; call start_contract_upload first".into(),
+                })?;
+
+        if !staged.is_complete() {
+            // Reinstate the staged entry so the caller can resume or clear. This
+            // matches the contract invariant that a non-complete upload is either
+            // visible (via the staged_uploads map) or has been explicitly cleared.
+            let total_size = staged.total_size;
+            let received_bytes = staged.received_bytes;
+            self.staged_uploads.insert(caller, staged);
+            return Err(InvalidParameters::MalformedPayload {
+                reason: format!(
+                    "upload incomplete: received_bytes={received_bytes}, total_size={total_size}"
+                ),
+            }
+            .into());
+        }
+
+        let num_chunks = staged.num_chunks;
+        let total_size = staged.total_size;
+        let deposited = staged.deposited;
+
+        // Read (don't remove) the chunks first so the deposit can be validated before
+        // any storage is mutated: if the deposit is short we reinstate the staged entry
+        // and bail with the chunks still in place.
+        let mut code = Vec::with_capacity(total_size.get());
+        for i in 0..num_chunks {
+            let chunk = self
+                .staged_chunks
+                .get(&(caller.clone(), i))
+                .expect("chunk recorded in staged metadata must be present in staged_chunks");
+            code.extend_from_slice(chunk);
+        }
+
+        let update = Update::Contract(code);
+        // The proposal entry reserves storage for the code *plus* per-proposal
+        // vote-tracking overhead. The per-chunk deposit only backs the raw bytes, so
+        // require the accumulated deposit to cover the full proposal cost — matching
+        // the invariant `propose_update` enforces for config proposals.
+        let required = ProposedUpdates::required_deposit(&update);
+        if deposited < required {
+            self.staged_uploads.insert(caller, staged);
+            return Err(InvalidParameters::InsufficientDeposit {
+                attached: deposited.as_yoctonear(),
+                required: required.as_yoctonear(),
+            }
+            .into());
+        }
+
+        // Deposit is sufficient; free the staged chunk storage and register the proposal.
+        for i in 0..num_chunks {
+            self.staged_chunks.remove(&(caller.clone(), i));
+        }
+        let id = self.proposed_updates.propose(update);
+
+        // Refund any deposit beyond the proposal's storage cost.
+        if let Some(diff) = deposited.checked_sub(required)
+            && diff > NearToken::from_yoctonear(0)
+        {
+            Promise::new(caller.clone()).transfer(diff).detach();
+        }
+
+        Ok(id)
+    }
+
+    /// Abandon the caller's in-progress chunked upload and refund the accumulated
+    /// chunk-storage deposits.
+    #[handle_result]
+    pub fn clear_staged_contract(&mut self) -> Result<(), Error> {
+        let caller = self.voter_or_panic();
+
+        let staged =
+            self.staged_uploads
+                .remove(&caller)
+                .ok_or(InvalidParameters::MalformedPayload {
+                    reason: "no staged upload found".into(),
+                })?;
+
+        for i in 0..staged.num_chunks {
+            self.staged_chunks.remove(&(caller.clone(), i));
+        }
+
+        if staged.deposited > NearToken::from_yoctonear(0) {
+            Promise::new(caller).transfer(staged.deposited).detach();
+        }
+
+        Ok(())
+    }
+
     /// Vote for a proposed update given the [`UpdateId`] of the update.
     ///
     /// Returns `Ok(true)` if the amount of voters surpassed the threshold and the update was
@@ -1630,8 +1805,31 @@ impl MpcContract {
             }
         };
 
+        // An in-progress chunked upload owned by an account that is no longer a
+        // participant can never be finalized or cleared by its owner (both endpoints
+        // require `voter_or_panic`), so its chunk bytes and metadata would linger
+        // forever, locking the contract's balance behind storage staking. Drop them
+        // here and refund each owner's accumulated deposit.
+        let orphaned: Vec<(AccountId, usize, NearToken)> = self
+            .staged_uploads
+            .iter()
+            .filter(|(account, _)| !participants.is_participant_given_account_id(account))
+            .map(|(account, staged)| (account.clone(), staged.num_chunks, staged.deposited))
+            .collect();
+
         self.proposed_updates
             .remove_non_participant_votes(participants);
+
+        for (account, num_chunks, deposited) in orphaned {
+            for i in 0..num_chunks {
+                self.staged_chunks.remove(&(account.clone(), i));
+            }
+            self.staged_uploads.remove(&account);
+            if deposited > NearToken::from_yoctonear(0) {
+                Promise::new(account).transfer(deposited).detach();
+            }
+        }
+
         Ok(())
     }
 
@@ -1757,6 +1955,8 @@ impl MpcContract {
                 StorageKey::PendingVerifyForeignTxRequestsV2,
             ),
             proposed_updates: ProposedUpdates::default(),
+            staged_uploads: IterableMap::new(StorageKey::StagedContractUploads),
+            staged_chunks: LookupMap::new(StorageKey::StagedContractChunks),
             config: init_config.map(Into::into).unwrap_or_default(),
             tee_state,
             accept_requests: true,
@@ -1826,6 +2026,8 @@ impl MpcContract {
                 StorageKey::PendingVerifyForeignTxRequestsV2,
             ),
             proposed_updates: Default::default(),
+            staged_uploads: IterableMap::new(StorageKey::StagedContractUploads),
+            staged_chunks: LookupMap::new(StorageKey::StagedContractChunks),
             tee_state,
             accept_requests: true,
             node_migrations: NodeMigrations::default(),
@@ -4036,6 +4238,8 @@ mod tests {
                 ),
                 accept_requests: true,
                 proposed_updates: Default::default(),
+                staged_uploads: IterableMap::new(StorageKey::StagedContractUploads),
+                staged_chunks: LookupMap::new(StorageKey::StagedContractChunks),
                 node_foreign_chain_support: Default::default(),
                 config: Default::default(),
                 tee_state: Default::default(),

crates/contract/src/snapshots/mpc_contract__tests__mpc_contract_borsh_schema_has_not_changed.snap

@@ -674,6 +674,14 @@ BorshSchemaContainer {
                         "proposed_updates",
                         "ProposedUpdates",
                     ),
+                    (
+                        "staged_uploads",
+                        "IterableMap",
+                    ),
+                    (
+                        "staged_chunks",
+                        "LookupMap",
+                    ),
                     (
                         "node_foreign_chain_support",
                         "SupportedForeignChainsByNode",

crates/contract/src/storage_keys.rs

@@ -30,4 +30,8 @@ pub enum StorageKey {
     AllowedForeignChainProvidersV1,
     ForeignChainProviderVotesByVoterV1,
     ForeignChainProviderVotesByProposalV1,
+    /// Per-account `StagedContractUpload` metadata.
+    StagedContractUploads,
+    /// Chunk bytes belonging to in-progress uploads, keyed by `(AccountId, chunk_index)`.
+    StagedContractChunks,
 }