Threshold Custody: Add internal signing protocol

crates/custody/src/threshold.rs

@@ -1 +1,82 @@
+use anyhow::{anyhow, Result};
+use penumbra_keys::{keys::AddressIndex, Address, FullViewingKey};
+use penumbra_proto::custody::v1alpha1::{self as pb};
+use penumbra_transaction::AuthorizationData;
+use tonic::{async_trait, Request, Response, Status};
+
+use crate::AuthorizeRequest;
+
 mod config;
+mod sign;
+
+/// A custody backend using threshold signing.  
+///
+/// This backend is initialized with a full viewing key, but only a share
+/// of the spend key, which is not enough to sign on its own. Instead,
+/// other signers with the same type of configuration need to cooperate
+/// to help produce a signature.
+pub struct Threshold {}
+
+impl Threshold {
+    /// Try and create the necessary signatures to authorize the transaction plan.
+    async fn authorize(&self, _request: AuthorizeRequest) -> Result<AuthorizationData> {
+        todo!()
+    }
+
+    /// Return the full viewing key.
+    fn export_full_viewing_key(&self) -> FullViewingKey {
+        todo!()
+    }
+
+    /// Get the address associated with an index.
+    ///
+    /// This is just to match the API of the custody trait.
+    fn confirm_address(&self, _index: AddressIndex) -> Address {
+        todo!()
+    }
+}
+
+#[async_trait]
+impl pb::custody_protocol_service_server::CustodyProtocolService for Threshold {
+    async fn authorize(
+        &self,
+        request: Request<pb::AuthorizeRequest>,
+    ) -> Result<Response<pb::AuthorizeResponse>, Status> {
+        let request = request
+            .into_inner()
+            .try_into()
+            .map_err(|e| Status::invalid_argument(format!("{e}")))?;
+        let data = self.authorize(request).await.map_err(|e| {
+            Status::internal(format!("Failed to process authorization request: {e}"))
+        })?;
+        Ok(Response::new(pb::AuthorizeResponse {
+            data: Some(data.into()),
+        }))
+    }
+
+    async fn export_full_viewing_key(
+        &self,
+        _request: Request<pb::ExportFullViewingKeyRequest>,
+    ) -> Result<Response<pb::ExportFullViewingKeyResponse>, Status> {
+        let fvk = self.export_full_viewing_key();
+        Ok(Response::new(pb::ExportFullViewingKeyResponse {
+            full_viewing_key: Some(fvk.into()),
+        }))
+    }
+
+    async fn confirm_address(
+        &self,
+        request: Request<pb::ConfirmAddressRequest>,
+    ) -> Result<Response<pb::ConfirmAddressResponse>, Status> {
+        let index = request
+            .into_inner()
+            .address_index
+            .ok_or(anyhow!("ConfirmAddressRequest missing address_index"))
+            .and_then(|x| x.try_into())
+            .map_err(|e| Status::invalid_argument(format!("{e}")))?;
+        let address = self.confirm_address(index);
+        Ok(Response::new(pb::ConfirmAddressResponse {
+            address: Some(address.into()),
+        }))
+    }
+}

crates/custody/src/threshold/sign.rs

@@ -282,3 +282,176 @@ impl TypeUrl for FollowerRound2 {
 impl DomainType for FollowerRound2 {
     type Proto = pb::FollowerRound2;
 }
+
+/// Calculate the number of required signatures for a plan.
+///
+/// A plan can require more than one signature, hence the need for this method.
+fn required_signatures(plan: &TransactionPlan) -> usize {
+    plan.spend_plans().count() + plan.delegator_vote_plans().count()
+}
+
+pub struct CoordinatorState1 {
+    plan: TransactionPlan,
+    my_round1_reply: FollowerRound1,
+    my_round1_state: FollowerState,
+}
+
+pub struct CoordinatorState2 {
+    plan: TransactionPlan,
+    my_round2_reply: FollowerRound2,
+    effect_hash: EffectHash,
+    signing_packages: Vec<frost::SigningPackage>,
+}
+
+pub struct FollowerState {
+    plan: TransactionPlan,
+    nonces: Vec<frost::round1::SigningNonces>,
+}
+
+pub fn coordinator_round1(
+    rng: &mut impl CryptoRngCore,
+    config: &Config,
+    plan: TransactionPlan,
+) -> Result<(CoordinatorRound1, CoordinatorState1)> {
+    let required = required_signatures(&plan);
+    let message = CoordinatorRound1 { plan: plan.clone() };
+    let (my_round1_reply, my_round1_state) = follower_round1(rng, config, message.clone())?;
+    let state = CoordinatorState1 {
+        plan,
+        my_round1_reply,
+        my_round1_state,
+    };
+    Ok((message, state))
+}
+
+pub fn coordinator_round2(
+    config: &Config,
+    state: CoordinatorState1,
+    follower_messages: &[FollowerRound1],
+) -> Result<(CoordinatorRound2, CoordinatorState2)> {
+    let mut all_commitments = vec![BTreeMap::new(); required_signatures(&state.plan)];
+    for message in follower_messages
+        .iter()
+        .cloned()
+        .chain(iter::once(state.my_round1_reply))
+    {
+        let (pk, commitments) = message.checked_commitments()?;
+        if !config.verification_keys.contains(&pk) {
+            anyhow::bail!("unknown verification key: {:?}", pk);
+        }
+        // The public key acts as the identifier
+        let identifier = frost::Identifier::derive(pk.as_bytes().as_slice())?;
+        for (tree_i, com_i) in all_commitments.iter_mut().zip(commitments.into_iter()) {
+            tree_i.insert(identifier, com_i);
+        }
+    }
+    let reply = CoordinatorRound2 { all_commitments };
+
+    let my_round2_reply = follower_round2(config, state.my_round1_state, reply.clone())?;
+    let effect_hash = state.plan.effect_hash(&config.fvk);
+    let signing_packages = {
+        reply
+            .all_commitments
+            .iter()
+            .map(|tree| frost::SigningPackage::new(tree.clone(), effect_hash.as_ref()))
+            .collect()
+    };
+    let state = CoordinatorState2 {
+        plan: state.plan,
+        my_round2_reply,
+        effect_hash,
+        signing_packages,
+    };
+    Ok((reply, state))
+}
+
+pub fn coordinator_round3(
+    config: &Config,
+    state: CoordinatorState2,
+    follower_messages: &[FollowerRound2],
+) -> Result<AuthorizationData> {
+    let mut share_maps: Vec<HashMap<frost::Identifier, frost::round2::SignatureShare>> =
+        vec![HashMap::new(); required_signatures(&state.plan)];
+    for message in follower_messages
+        .iter()
+        .cloned()
+        .chain(iter::once(state.my_round2_reply))
+    {
+        let (pk, shares) = message.checked_shares()?;
+        if !config.verification_keys.contains(&pk) {
+            anyhow::bail!("unknown verification key: {:?}", pk);
+        }
+        let identifier = frost::Identifier::derive(pk.as_bytes().as_slice())?;
+        for (map_i, share_i) in share_maps.iter_mut().zip(shares.into_iter()) {
+            map_i.insert(identifier, share_i);
+        }
+    }
+    let mut spend_auths = state
+        .plan
+        .spend_plans()
+        .map(|x| x.randomizer)
+        .chain(state.plan.delegator_vote_plans().map(|x| x.randomizer))
+        .zip(share_maps.iter())
+        .zip(state.signing_packages.iter())
+        .map(|((randomizer, share_map), signing_package)| {
+            frost::aggregate_randomized(
+                signing_package,
+                &share_map,
+                &config.public_key_package,
+                randomizer,
+            )
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    let delegator_vote_auths = spend_auths.split_off(state.plan.spend_plans().count());
+    Ok(AuthorizationData {
+        effect_hash: state.effect_hash,
+        spend_auths,
+        delegator_vote_auths,
+    })
+}
+
+pub fn follower_round1(
+    rng: &mut impl CryptoRngCore,
+    config: &Config,
+    coordinator: CoordinatorRound1,
+) -> Result<(FollowerRound1, FollowerState)> {
+    let required = required_signatures(&coordinator.plan);
+    let (nonces, commitments) = (0..required)
+        .map(|_| frost::round1::commit(&config.key_package.secret_share(), rng))
+        .unzip();
+    let reply = FollowerRound1::make(&config.signing_key, commitments);
+    let state = FollowerState {
+        plan: coordinator.plan,
+        nonces,
+    };
+    Ok((reply, state))
+}
+
+pub fn follower_round2(
+    config: &Config,
+    state: FollowerState,
+    coordinator: CoordinatorRound2,
+) -> Result<FollowerRound2> {
+    let effect_hash = state.plan.effect_hash(&config.fvk);
+    let signing_packages = coordinator
+        .all_commitments
+        .into_iter()
+        .map(|tree| frost::SigningPackage::new(tree, effect_hash.as_ref()));
+    let shares = state
+        .plan
+        .spend_plans()
+        .map(|x| x.randomizer)
+        .chain(state.plan.delegator_vote_plans().map(|x| x.randomizer))
+        .zip(signing_packages)
+        .zip(state.nonces.into_iter())
+        .map(|((randomizer, signing_package), signer_nonces)| {
+            frost::round2::sign_randomized(
+                &signing_package,
+                &signer_nonces,
+                &config.key_package,
+                randomizer,
+            )
+        })
+        .collect::<Result<_, _>>()?;
+    Ok(FollowerRound2::make(&config.signing_key, shares))
+}