TaprootWitness for musig2 signer

crates/secret-service-client/Cargo.toml

@@ -10,6 +10,7 @@ musig2 = { path = "../musig2" }
 quinn.workspace = true
 rkyv.workspace = true
 secret-service-proto = { version = "0.1.0", path = "../secret-service-proto" }
+strata-bridge-primitives.workspace = true
 terrors.workspace = true
 tokio.workspace = true
 tracing.workspace = true

crates/secret-service-client/src/lib.rs

@@ -31,6 +31,7 @@ use secret_service_proto::{
         WireMessage,
     },
 };
+use strata_bridge_primitives::scripts::taproot::TaprootWitness;
 use terrors::OneOf;
 use tokio::time::timeout;
 
@@ -428,12 +429,15 @@ struct Musig2Client {
 impl Musig2Signer<Client, Musig2FirstRound> for Musig2Client {
     fn new_session(
         &self,
-        ctx: KeyAggContext,
-        signer_idx: usize,
+        pubkeys: Vec<PublicKey>,
+        witness: TaprootWitness,
     ) -> impl Future<Output = Result<Result<Musig2FirstRound, SignerIdxOutOfBounds>, ClientError>> + Send
     {
         async move {
-            let msg = ClientMessage::Musig2NewSession { ctx, signer_idx };
+            let msg = ClientMessage::Musig2NewSession {
+                pubkeys: pubkeys.into_iter().map(|pk| pk.serialize()).collect(),
+                witness: witness.into(),
+            };
             let res = make_v1_req(&self.conn, msg, self.config.timeout).await?;
             let ServerMessage::Musig2NewSession(maybe_session_id) = res else {
                 return Err(ClientError::ProtocolError(res));
@@ -449,6 +453,18 @@ impl Musig2Signer<Client, Musig2FirstRound> for Musig2Client {
             })
         }
     }
+
+    fn pubkey(&self) -> impl Future<Output = <Client as Origin>::Container<PublicKey>> + Send {
+        async move {
+            let msg = ClientMessage::Musig2Pubkey;
+            let res = make_v1_req(&self.conn, msg, self.config.timeout).await?;
+            let ServerMessage::Musig2Pubkey { pubkey } = res else {
+                return Err(ClientError::ProtocolError(res));
+            };
+
+            PublicKey::from_slice(&pubkey).map_err(|_| ClientError::ProtocolError(res))
+        }
+    }
 }
 
 struct WotsClient {

crates/secret-service-proto/Cargo.toml

@@ -8,3 +8,4 @@ bitcoin.workspace = true
 musig2 = { path = "../musig2" }
 quinn.workspace = true
 rkyv.workspace = true
+strata-bridge-primitives.workspace = true

crates/secret-service-proto/src/v1/traits.rs

@@ -8,6 +8,7 @@ use musig2::{
 };
 use quinn::{ConnectionError, ReadExactError, WriteError};
 use rkyv::{rancor, Archive, Deserialize, Serialize};
+use strata_bridge_primitives::scripts::taproot::TaprootWitness;
 
 use super::wire::ServerMessage;
 
@@ -63,9 +64,10 @@ pub struct SignerIdxOutOfBounds {
 pub trait Musig2Signer<O: Origin, FirstRound>: Send + Sync {
     fn new_session(
         &self,
-        ctx: KeyAggContext,
-        signer_idx: usize,
+        pubkeys: Vec<PublicKey>,
+        witness: TaprootWitness,
     ) -> impl Future<Output = O::Container<Result<FirstRound, SignerIdxOutOfBounds>>> + Send;
+    fn pubkey(&self) -> impl Future<Output = O::Container<PublicKey>> + Send;
 }
 
 pub trait Musig2SignerFirstRound<O: Origin, SecondRound>: Send + Sync {

crates/secret-service-proto/src/v1/wire.rs

@@ -1,8 +1,14 @@
+use bitcoin::{
+    hashes::Hash,
+    taproot::{ControlBlock, TaprootError},
+    ScriptBuf, TapNodeHash,
+};
 use musig2::{
     errors::{RoundContributionError, RoundFinalizeError},
     KeyAggContext,
 };
 use rkyv::{with::Map, Archive, Deserialize, Serialize};
+use strata_bridge_primitives::scripts::taproot::TaprootWitness;
 
 use super::traits::{Musig2SessionId, SignerIdxOutOfBounds};
 
@@ -26,6 +32,9 @@ pub enum ServerMessage {
     },
 
     Musig2NewSession(Result<Musig2SessionId, SignerIdxOutOfBounds>),
+    Musig2Pubkey {
+        pubkey: [u8; 33],
+    },
 
     Musig2FirstRoundOurNonce {
         our_nonce: [u8; 66],
@@ -110,9 +119,10 @@ pub enum ClientMessage {
     P2PPubkey,
 
     Musig2NewSession {
-        ctx: KeyAggContext,
-        signer_idx: usize,
+        pubkeys: Vec<[u8; 33]>,
+        witness: SerializableTaprootWitness,
     },
+    Musig2Pubkey,
 
     Musig2FirstRoundOurNonce {
         session_id: usize,
@@ -163,3 +173,62 @@ pub enum ClientMessage {
         deposit_idx: u64,
     },
 }
+
+#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
+pub enum SerializableTaprootWitness {
+    Key,
+    Script {
+        script_buf: Vec<u8>,
+        control_block: Vec<u8>,
+    },
+    Tweaked {
+        tweak: [u8; 32],
+    },
+}
+
+impl From<TaprootWitness> for SerializableTaprootWitness {
+    fn from(witness: TaprootWitness) -> Self {
+        match witness {
+            TaprootWitness::Key => SerializableTaprootWitness::Key,
+            TaprootWitness::Script {
+                script_buf,
+                control_block,
+            } => SerializableTaprootWitness::Script {
+                script_buf: script_buf.into_bytes(),
+                control_block: control_block.serialize(),
+            },
+            TaprootWitness::Tweaked { tweak } => SerializableTaprootWitness::Tweaked {
+                tweak: tweak.to_raw_hash().to_byte_array(),
+            },
+        }
+    }
+}
+
+pub enum TaprootWitnessError {
+    InvalidWitnessType,
+    InvalidScriptControlBlock(TaprootError),
+}
+
+impl TryFrom<SerializableTaprootWitness> for TaprootWitness {
+    type Error = TaprootWitnessError;
+    fn try_from(value: SerializableTaprootWitness) -> Result<Self, Self::Error> {
+        match value {
+            SerializableTaprootWitness::Key => Ok(TaprootWitness::Key),
+            SerializableTaprootWitness::Script {
+                script_buf,
+                control_block,
+            } => {
+                let script_buf = ScriptBuf::from_bytes(script_buf);
+                let control_block = ControlBlock::decode(&control_block)
+                    .map_err(TaprootWitnessError::InvalidScriptControlBlock)?;
+                Ok(TaprootWitness::Script {
+                    script_buf,
+                    control_block,
+                })
+            }
+            SerializableTaprootWitness::Tweaked { tweak } => Ok(TaprootWitness::Tweaked {
+                tweak: TapNodeHash::from_byte_array(tweak),
+            }),
+        }
+    }
+}

crates/secret-service-server/Cargo.toml

@@ -12,6 +12,7 @@ parking_lot.workspace = true
 quinn.workspace = true
 rkyv.workspace = true
 secret-service-proto = { version = "0.1.0", path = "../secret-service-proto" }
+strata-bridge-primitives.workspace = true
 terrors.workspace = true
 tokio.workspace = true
 tracing.workspace = true

crates/secret-service-server/src/lib.rs

@@ -34,6 +34,7 @@ use secret_service_proto::{
     },
     wire::{ArchivedVersionedClientMessage, LengthUint, VersionedServerMessage, WireMessage},
 };
+use strata_bridge_primitives::scripts::taproot::TaprootWitness;
 use terrors::OneOf;
 use tokio::{
     sync::Mutex,
@@ -238,15 +239,25 @@ where
                 }
             }
 
-            ArchivedClientMessage::Musig2NewSession { ctx, signer_idx } => 'block: {
+            ArchivedClientMessage::Musig2NewSession { pubkeys, witness } => 'block: {
                 let signer = service.musig2_signer();
-                let Ok(ctx) = deserialize::<KeyAggContext, rancor::Error>(ctx) else {
+                let Ok(ser_witness) = deserialize::<_, rancor::Error>(witness) else {
                     break 'block ServerMessage::InvalidClientMessage;
                 };
-                let first_round = match signer
-                    .new_session(ctx, signer_idx.to_native() as usize)
-                    .await
-                {
+                let Ok(witness) = TaprootWitness::try_from(ser_witness)
+                    .map_err(|_| ServerMessage::InvalidClientMessage)
+                else {
+                    break 'block ServerMessage::InvalidClientMessage;
+                };
+                let Ok(pubkeys) = pubkeys
+                    .into_iter()
+                    .map(|data| PublicKey::from_slice(data))
+                    .collect::<Result<Vec<_>, _>>()
+                else {
+                    break 'block ServerMessage::InvalidClientMessage;
+                };
+
+                let first_round = match signer.new_session(pubkeys, witness).await {
                     Ok(fr) => fr,
                     Err(e) => break 'block ServerMessage::Musig2NewSession(Err(e)),
                 };
@@ -266,6 +277,10 @@ where
 
                 ServerMessage::Musig2NewSession(Ok(write_perm.session_id()))
             }
+            ArchivedClientMessage::Musig2Pubkey => ServerMessage::Musig2Pubkey {
+                pubkey: service.musig2_signer().pubkey().await.serialize(),
+            },
+
             ArchivedClientMessage::Musig2FirstRoundOurNonce { session_id } => {
                 let r = musig2_sm
                     .lock()

crates/secret-service/Cargo.toml

@@ -18,6 +18,7 @@ secret-service-server = { version = "0.1.0", path = "../secret-service-server" }
 serde.workspace = true
 sha2.workspace = true
 sled = "0.34.7"
+strata-bridge-primitives.workspace = true
 strata-key-derivation = { git = "https://github.com/alpenlabs/strata", version = "0.1.0" }
 terrors.workspace = true
 

crates/secret-service/src/disk/musig2.rs

@@ -1,8 +1,9 @@
 use std::future::Future;
 
+use bitcoin::key::Keypair;
 use musig2::{
     errors::{RoundContributionError, RoundFinalizeError},
-    secp256k1::{PublicKey, SecretKey},
+    secp256k1::{PublicKey, SecretKey, SECP256K1},
     FirstRound, KeyAggContext, LiftedSignature, SecNonceSpices, SecondRound,
 };
 use rand::{thread_rng, Rng};
@@ -13,44 +14,74 @@ use secret_service_proto::v1::traits::{
 };
 use secret_service_server::RoundPersister;
 use sled::Tree;
+use strata_bridge_primitives::scripts::taproot::TaprootWitness;
 use terrors::OneOf;
 
 pub struct Ms2Signer {
-    key: SecretKey,
+    kp: Keypair,
 }
 
 impl Ms2Signer {
     pub fn new(key: SecretKey) -> Self {
-        Self { key }
+        Self {
+            kp: Keypair::from_secret_key(SECP256K1, &key),
+        }
     }
 }
 
 impl Musig2Signer<Server, ServerFirstRound> for Ms2Signer {
     fn new_session(
         &self,
-        ctx: KeyAggContext,
-        signer_idx: usize,
+        mut pubkeys: Vec<PublicKey>,
+        witness: TaprootWitness,
     ) -> impl Future<Output = Result<ServerFirstRound, SignerIdxOutOfBounds>> + Send {
         async move {
             let nonce_seed = thread_rng().gen::<[u8; 32]>();
-            let ordered_public_keys = ctx.pubkeys().iter().cloned().map(|p| p.into()).collect();
+            if !pubkeys.contains(&self.kp.public_key()) {
+                pubkeys.push(self.kp.public_key());
+            }
+            pubkeys.sort();
+            let signer_index = pubkeys
+                .iter()
+                .position(|pk| pk == &self.kp.public_key())
+                .unwrap();
+            let mut ctx = KeyAggContext::new(pubkeys.clone()).unwrap();
+
+            match witness {
+                TaprootWitness::Key => {
+                    ctx = ctx
+                        .with_unspendable_taproot_tweak()
+                        .expect("must be able to tweak the key agg context")
+                }
+                TaprootWitness::Tweaked { tweak } => {
+                    ctx = ctx
+                        .with_taproot_tweak(tweak.as_ref())
+                        .expect("must be able to tweak the key agg context")
+                }
+                _ => {}
+            }
+
             let first_round = FirstRound::new(
                 ctx,
                 nonce_seed,
-                signer_idx,
-                SecNonceSpices::new().with_seckey(self.key.clone()),
+                signer_index,
+                SecNonceSpices::new().with_seckey(self.kp.secret_key()),
             )
             .map_err(|e| SignerIdxOutOfBounds {
                 index: e.index,
                 n_signers: e.n_signers,
             })?;
             Ok(ServerFirstRound {
                 first_round,
-                ordered_public_keys,
-                seckey: self.key.clone(),
+                ordered_public_keys: pubkeys,
+                seckey: self.kp.secret_key(),
             })
         }
     }
+
+    fn pubkey(&self) -> impl Future<Output = <Server as Origin>::Container<PublicKey>> + Send {
+        async move { self.kp.public_key() }
+    }
 }
 
 pub struct SledRoundPersist {