silent_payments_psbt/

roles.rs

//! BIP 375 PSBT role typestate machine.
//!
//! Four distinct types enforce role ordering at compile time:
//! [`SpPsbtConstructor`] -> [`SpPsbtUpdater`] -> [`SpPsbtSigner`] -> [`SpPsbtExtractor`]
//!
//! Each transition consumes `self` and returns the next type, preventing
//! skipping or reordering roles. The compiler enforces correctness.

use bitcoin::psbt::Psbt;
use bitcoin::secp256k1::PublicKey;

use silent_payments_core::keys::{ScanPublicKey, SpendPublicKey};

use crate::dleq::DleqProof;
use crate::error::{DleqError, PsbtError};
use crate::fields;

// ---------------------------------------------------------------------------
// Constructor
// ---------------------------------------------------------------------------

/// BIP 375 Constructor role: adds SP output info and labels to the PSBT.
///
/// Created from an unsigned PSBT via [`SpPsbtConstructor::new`]. The constructor
/// allows adding Silent Payment recipient info (scan + spend keys) and optional
/// labels to PSBT outputs. Transitions to [`SpPsbtUpdater`] after at least one
/// SP output has been added.
pub struct SpPsbtConstructor {
    inner: Psbt,
}

impl SpPsbtConstructor {
    /// Wrap a PSBT for Silent Payment construction.
    ///
    /// No validation is performed at this stage -- the PSBT may have zero outputs
    /// initially (outputs can be added by the wallet before calling `add_sp_output`).
    pub fn new(psbt: Psbt) -> Result<Self, PsbtError> {
        Ok(Self { inner: psbt })
    }

    /// Add Silent Payment recipient info to a PSBT output.
    ///
    /// Stores the scan and spend public keys as `PSBT_OUT_SP_V0_INFO` in the
    /// output's unknown map.
    ///
    /// # Errors
    ///
    /// Returns [`PsbtError::OutputIndexOutOfBounds`] if `output_index` exceeds
    /// the number of PSBT outputs.
    pub fn add_sp_output(
        mut self,
        output_index: usize,
        scan_pubkey: &ScanPublicKey,
        spend_pubkey: &SpendPublicKey,
    ) -> Result<Self, PsbtError> {
        if output_index >= self.inner.outputs.len() {
            return Err(PsbtError::OutputIndexOutOfBounds {
                index: output_index,
                count: self.inner.outputs.len(),
            });
        }

        fields::set_sp_output_info(
            &mut self.inner.outputs[output_index],
            scan_pubkey.as_inner(),
            spend_pubkey.as_inner(),
        );

        Ok(self)
    }

    /// Add an optional label to a PSBT output.
    ///
    /// Stores the label as `PSBT_OUT_SP_V0_LABEL` in the output's unknown map.
    ///
    /// # Errors
    ///
    /// Returns [`PsbtError::OutputIndexOutOfBounds`] if `output_index` exceeds
    /// the number of PSBT outputs.
    pub fn add_sp_label(mut self, output_index: usize, label: u32) -> Result<Self, PsbtError> {
        if output_index >= self.inner.outputs.len() {
            return Err(PsbtError::OutputIndexOutOfBounds {
                index: output_index,
                count: self.inner.outputs.len(),
            });
        }

        fields::set_sp_output_label(&mut self.inner.outputs[output_index], label);

        Ok(self)
    }

    /// Transition to the Updater role.
    ///
    /// Validates that at least one PSBT output has SP_V0_INFO fields.
    ///
    /// # Errors
    ///
    /// Returns [`PsbtError::NoSpOutputs`] if no output has SP info.
    pub fn into_updater(self) -> Result<SpPsbtUpdater, PsbtError> {
        let has_sp_output = self
            .inner
            .outputs
            .iter()
            .any(|output| fields::get_sp_output_info(output).is_some());

        if !has_sp_output {
            return Err(PsbtError::NoSpOutputs);
        }

        Ok(SpPsbtUpdater { inner: self.inner })
    }
}

// ---------------------------------------------------------------------------
// Updater
// ---------------------------------------------------------------------------

/// BIP 375 Updater role: allows wallet-specific BIP32 operations on the PSBT.
///
/// This is a passthrough role. The wallet can add BIP32 derivation paths,
/// witness UTXOs, and other standard PSBT data via [`inner_mut`](SpPsbtUpdater::inner_mut).
/// Transitions to [`SpPsbtSigner`] without additional validation.
pub struct SpPsbtUpdater {
    inner: Psbt,
}

impl SpPsbtUpdater {
    /// Read access to the inner PSBT for wallet inspection.
    pub fn inner(&self) -> &Psbt {
        &self.inner
    }

    /// Write access to the inner PSBT for wallet modifications.
    ///
    /// Wallets use this to add BIP32 derivation paths, witness UTXOs,
    /// and other standard PSBT data before signing.
    pub fn inner_mut(&mut self) -> &mut Psbt {
        &mut self.inner
    }

    /// Transition to the Signer role.
    ///
    /// No validation is required -- BIP32 paths are optional.
    pub fn into_signer(self) -> SpPsbtSigner {
        SpPsbtSigner { inner: self.inner }
    }
}

// ---------------------------------------------------------------------------
// Signer
// ---------------------------------------------------------------------------

/// BIP 375 Signer role: adds ECDH shares with mandatory DLEQ proof validation.
///
/// Every call to [`add_ecdh_share`](SpPsbtSigner::add_ecdh_share) or
/// [`add_global_ecdh_share`](SpPsbtSigner::add_global_ecdh_share) verifies
/// the accompanying DLEQ proof before storing the share. Mixing global and
/// per-input shares for the same scan key is rejected.
pub struct SpPsbtSigner {
    inner: Psbt,
}

impl SpPsbtSigner {
    /// Add a per-input ECDH share with mandatory DLEQ proof verification.
    ///
    /// The DLEQ proof is verified before storing the share. If a global ECDH share
    /// already exists for this scan key, returns [`PsbtError::DuplicateShareType`].
    ///
    /// # Arguments
    ///
    /// * `input_index` - The PSBT input index to store the share on
    /// * `scan_pubkey` - The recipient's scan public key
    /// * `ecdh_share` - The signer's ECDH share (a_i * B_scan)
    /// * `dleq_proof` - The DLEQ proof for this share
    /// * `signer_pubkey` - The signer's public key (A = a * G)
    /// * `secp` - Secp256k1 verification context
    pub fn add_ecdh_share(
        mut self,
        input_index: usize,
        scan_pubkey: &ScanPublicKey,
        ecdh_share: &PublicKey,
        dleq_proof: &DleqProof,
        signer_pubkey: &PublicKey,
        secp: &bitcoin::secp256k1::Secp256k1<impl bitcoin::secp256k1::Verification>,
    ) -> Result<Self, PsbtError> {
        // Validate input index
        if input_index >= self.inner.inputs.len() {
            return Err(PsbtError::OutputIndexOutOfBounds {
                index: input_index,
                count: self.inner.inputs.len(),
            });
        }

        // Check for existing global share for this scan key (reject mixing)
        if fields::get_global_ecdh_share(&self.inner, scan_pubkey.as_inner()).is_some() {
            return Err(PsbtError::DuplicateShareType {
                scan_key: hex_encode_pubkey(scan_pubkey.as_inner()),
            });
        }

        // Verify DLEQ proof (mandatory -- no unchecked paths)
        let valid = dleq_proof.verify(signer_pubkey, scan_pubkey.as_inner(), ecdh_share, secp)?;
        if !valid {
            return Err(PsbtError::InvalidProof(DleqError::VerificationFailed));
        }

        // Store per-input ECDH share and DLEQ proof
        fields::set_input_ecdh_share(
            &mut self.inner.inputs[input_index],
            scan_pubkey.as_inner(),
            ecdh_share,
        );
        fields::set_input_dleq_proof(
            &mut self.inner.inputs[input_index],
            scan_pubkey.as_inner(),
            dleq_proof.as_bytes(),
        );

        Ok(self)
    }

    /// Add a global ECDH share with mandatory DLEQ proof verification.
    ///
    /// Used for single-signer workflows where the signer contributes a single
    /// aggregated share. The DLEQ proof is verified before storing.
    ///
    /// If per-input ECDH shares already exist for this scan key across any input,
    /// returns [`PsbtError::DuplicateShareType`].
    pub fn add_global_ecdh_share(
        mut self,
        scan_pubkey: &ScanPublicKey,
        ecdh_share: &PublicKey,
        dleq_proof: &DleqProof,
        signer_pubkey: &PublicKey,
        secp: &bitcoin::secp256k1::Secp256k1<impl bitcoin::secp256k1::Verification>,
    ) -> Result<Self, PsbtError> {
        // Check for existing per-input shares for this scan key (reject mixing)
        let has_per_input = self
            .inner
            .inputs
            .iter()
            .any(|input| fields::get_input_ecdh_share(input, scan_pubkey.as_inner()).is_some());

        if has_per_input {
            return Err(PsbtError::DuplicateShareType {
                scan_key: hex_encode_pubkey(scan_pubkey.as_inner()),
            });
        }

        // Verify DLEQ proof (mandatory -- no unchecked paths)
        let valid = dleq_proof.verify(signer_pubkey, scan_pubkey.as_inner(), ecdh_share, secp)?;
        if !valid {
            return Err(PsbtError::InvalidProof(DleqError::VerificationFailed));
        }

        // Store global ECDH share and DLEQ proof
        fields::set_global_ecdh_share(&mut self.inner, scan_pubkey.as_inner(), ecdh_share);
        fields::set_global_dleq_proof(
            &mut self.inner,
            scan_pubkey.as_inner(),
            dleq_proof.as_bytes(),
        );

        Ok(self)
    }

    /// Transition to the Extractor role.
    ///
    /// Validates that every unique scan key referenced in SP output info fields
    /// has either a global ECDH share or per-input shares on all inputs.
    ///
    /// # Errors
    ///
    /// Returns [`PsbtError::IncompleteSigning`] if any scan key lacks shares.
    pub fn into_extractor(self) -> Result<SpPsbtExtractor, PsbtError> {
        // Collect all unique scan pubkeys from SP output info fields
        let mut scan_keys: Vec<PublicKey> = Vec::new();
        for output in &self.inner.outputs {
            if let Some((scan_key, _spend_key)) = fields::get_sp_output_info(output) {
                if !scan_keys.contains(&scan_key) {
                    scan_keys.push(scan_key);
                }
            }
        }

        // For each scan key: check if global share OR all inputs have per-input shares
        for scan_key in &scan_keys {
            let has_global = fields::get_global_ecdh_share(&self.inner, scan_key).is_some();

            if has_global {
                continue;
            }

            // Check all inputs have per-input shares
            let all_inputs_have_shares = self
                .inner
                .inputs
                .iter()
                .all(|input| fields::get_input_ecdh_share(input, scan_key).is_some());

            if !all_inputs_have_shares {
                return Err(PsbtError::IncompleteSigning);
            }
        }

        Ok(SpPsbtExtractor { inner: self.inner })
    }
}

/// Helper: hex-encode a public key for error messages.
fn hex_encode_pubkey(key: &PublicKey) -> String {
    key.serialize().iter().map(|b| format!("{b:02x}")).collect()
}

/// BIP 375 Extractor role: aggregates ECDH shares and computes SP output scripts.
///
/// Created via [`SpPsbtSigner::into_extractor`] after all required ECDH shares
/// are present. The [`extract`](SpPsbtExtractor::extract) method computes
/// the final SP output scripts and returns the updated PSBT.
pub struct SpPsbtExtractor {
    inner: Psbt,
}

impl SpPsbtExtractor {
    /// Compute SP output scripts from aggregated ECDH shares.
    ///
    /// For each unique scan key in SP output info fields:
    /// 1. Aggregates ECDH shares (global share directly, or per-input via `combine_keys`)
    /// 2. The aggregated share IS the BIP 352 shared secret point
    /// 3. Derives output pubkeys via `compute_output_pubkey(spend_key, shared_secret, k)`
    /// 4. Builds P2TR scripts from the output pubkeys
    ///
    /// Returns the updated PSBT (with computed scriptPubKeys) and extracted output metadata.
    pub fn extract(
        mut self,
        _secp: &bitcoin::secp256k1::Secp256k1<impl bitcoin::secp256k1::Verification>,
    ) -> Result<(Psbt, Vec<crate::output::ExtractedSpOutput>), PsbtError> {
        // Collect SP output info: (output_index, scan_key, spend_key)
        let sp_outputs: Vec<(usize, PublicKey, PublicKey)> = self
            .inner
            .outputs
            .iter()
            .enumerate()
            .filter_map(|(idx, output)| {
                fields::get_sp_output_info(output).map(|(scan, spend)| (idx, scan, spend))
            })
            .collect();

        if sp_outputs.is_empty() {
            return Err(PsbtError::NoSpOutputs);
        }

        // Group outputs by scan_key, preserving order for k-counter
        let mut groups: Vec<(PublicKey, Vec<(usize, PublicKey)>)> = Vec::new();
        for (idx, scan_key, spend_key) in &sp_outputs {
            if let Some(group) = groups.iter_mut().find(|(sk, _)| sk == scan_key) {
                group.1.push((*idx, *spend_key));
            } else {
                groups.push((*scan_key, vec![(*idx, *spend_key)]));
            }
        }

        let mut extracted: Vec<crate::output::ExtractedSpOutput> = Vec::new();

        for (scan_key, outputs_for_key) in &groups {
            // Aggregate ECDH share for this scan key
            let shared_secret = self.aggregate_ecdh_share(scan_key)?;

            // Derive output pubkeys for each output with this scan key
            for (k, (output_index, spend_key)) in outputs_for_key.iter().enumerate() {
                let output_pubkey = silent_payments_core::crypto::compute_output_pubkey(
                    &SpendPublicKey::from(*spend_key),
                    &shared_secret,
                    k as u32,
                )?;

                // Build P2TR script from the output pubkey's x-only key
                let (x_only, _parity) = output_pubkey.x_only_public_key();
                let tweaked = bitcoin::key::TweakedPublicKey::dangerous_assume_tweaked(x_only);
                let script = bitcoin::ScriptBuf::new_p2tr_tweaked(tweaked);

                // Update the PSBT output's scriptPubKey
                self.inner.unsigned_tx.output[*output_index].script_pubkey = script.clone();

                extracted.push(crate::output::ExtractedSpOutput::new(
                    script,
                    x_only,
                    *output_index,
                    *scan_key,
                ));
            }
        }

        Ok((self.inner, extracted))
    }

    /// Aggregate ECDH shares for a scan key.
    ///
    /// If a global share exists, returns it directly.
    /// Otherwise, aggregates per-input shares via `PublicKey::combine_keys`.
    fn aggregate_ecdh_share(&self, scan_key: &PublicKey) -> Result<PublicKey, PsbtError> {
        // Check for global share first
        if let Some(global_share) = fields::get_global_ecdh_share(&self.inner, scan_key) {
            return Ok(global_share);
        }

        // Collect per-input shares
        let shares: Vec<PublicKey> = self
            .inner
            .inputs
            .iter()
            .filter_map(|input| fields::get_input_ecdh_share(input, scan_key))
            .collect();

        if shares.is_empty() {
            return Err(PsbtError::MissingEcdhShare {
                scan_key: hex_encode_pubkey(scan_key),
            });
        }

        // Aggregate via combine_keys (sum of points)
        let refs: Vec<&PublicKey> = shares.iter().collect();
        PublicKey::combine_keys(&refs).map_err(|e| {
            PsbtError::Crypto(silent_payments_core::CryptoError::Secp256k1(e.to_string()))
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::absolute::LockTime;
    use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey};
    use bitcoin::transaction::{Transaction, Version};
    use bitcoin::{Amount, ScriptBuf, TxIn, TxOut};

    /// Helper: create a test public key from a scalar byte.
    fn test_pubkey(byte: u8) -> PublicKey {
        let secp = Secp256k1::new();
        let mut key_bytes = [0u8; 32];
        key_bytes[31] = byte;
        let sk = SecretKey::from_slice(&key_bytes).unwrap();
        sk.public_key(&secp)
    }

    /// Helper: create a test secret key from a scalar byte.
    #[cfg(feature = "experimental-dleq")]
    fn test_secret_key(byte: u8) -> SecretKey {
        let mut key_bytes = [0u8; 32];
        key_bytes[31] = byte;
        SecretKey::from_slice(&key_bytes).unwrap()
    }

    /// Helper: create a minimal PSBT with the given number of inputs and outputs.
    fn test_psbt(n_inputs: usize, n_outputs: usize) -> Psbt {
        let tx = Transaction {
            version: Version::TWO,
            lock_time: LockTime::ZERO,
            input: (0..n_inputs).map(|_| TxIn::default()).collect(),
            output: (0..n_outputs)
                .map(|_| TxOut {
                    value: Amount::from_sat(50_000),
                    script_pubkey: ScriptBuf::new(),
                })
                .collect(),
        };
        Psbt::from_unsigned_tx(tx).unwrap()
    }

    /// Helper: create a minimal PSBT with 1 input and the given number of outputs.
    fn test_psbt_with_outputs(n_outputs: usize) -> Psbt {
        test_psbt(1, n_outputs)
    }

    /// Helper: build a signer from a PSBT with SP output info already set.
    fn build_signer(
        n_inputs: usize,
        n_outputs: usize,
        scan_byte: u8,
        spend_byte: u8,
    ) -> SpPsbtSigner {
        let psbt = test_psbt(n_inputs, n_outputs);
        let scan = ScanPublicKey::from(test_pubkey(scan_byte));
        let spend = SpendPublicKey::from(test_pubkey(spend_byte));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();
        let updater = constructor.into_updater().unwrap();
        updater.into_signer()
    }

    // --- Constructor tests ---

    #[test]
    fn test_constructor_add_sp_output() {
        let psbt = test_psbt_with_outputs(1);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();

        // Transition to updater so we can verify via inner()
        let updater = constructor.into_updater().unwrap();
        let (got_scan, got_spend) =
            fields::get_sp_output_info(&updater.inner().outputs[0]).expect("should read back");
        assert_eq!(got_scan, *scan.as_inner());
        assert_eq!(got_spend, *spend.as_inner());
    }

    #[test]
    fn test_constructor_output_index_out_of_bounds() {
        let psbt = test_psbt_with_outputs(1);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let result = constructor.add_sp_output(5, &scan, &spend);

        assert!(
            matches!(
                result,
                Err(PsbtError::OutputIndexOutOfBounds { index: 5, count: 1 })
            ),
            "should reject out-of-bounds index, got: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_constructor_into_updater_requires_sp_output() {
        let psbt = test_psbt_with_outputs(1);
        let constructor = SpPsbtConstructor::new(psbt).unwrap();

        let result = constructor.into_updater();
        assert!(
            matches!(result, Err(PsbtError::NoSpOutputs)),
            "should require at least one SP output, got: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_constructor_into_updater_succeeds() {
        let psbt = test_psbt_with_outputs(1);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();

        let updater = constructor.into_updater();
        assert!(updater.is_ok(), "should transition after adding SP output");
    }

    #[test]
    fn test_constructor_add_sp_label() {
        let psbt = test_psbt_with_outputs(1);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();
        let constructor = constructor.add_sp_label(0, 42).unwrap();

        let updater = constructor.into_updater().unwrap();
        let label = fields::get_sp_output_label(&updater.inner().outputs[0]);
        assert_eq!(label, Some(42));
    }

    // --- Updater tests ---

    #[test]
    fn test_updater_inner_mut_access() {
        let psbt = test_psbt_with_outputs(2);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();
        let mut updater = constructor.into_updater().unwrap();

        // Wallet can modify the PSBT via inner_mut
        let psbt_mut = updater.inner_mut();
        assert_eq!(psbt_mut.outputs.len(), 2, "should have 2 outputs");

        // Modify something to prove write access works
        let original_len = psbt_mut.inputs.len();
        assert!(original_len > 0, "should have at least one input");
    }

    #[test]
    fn test_updater_into_signer() {
        let psbt = test_psbt_with_outputs(1);
        let scan = ScanPublicKey::from(test_pubkey(1));
        let spend = SpendPublicKey::from(test_pubkey(2));

        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan, &spend).unwrap();
        let updater = constructor.into_updater().unwrap();

        // Should transition without error
        let _signer = updater.into_signer();
    }

    // --- Signer tests ---

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_signer_add_ecdh_share_validates_dleq() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));
        // Signer's secret key (the input key)
        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        // Generate a valid DLEQ proof
        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let signer = build_signer(1, 1, 10, 20);
        let signer = signer
            .add_ecdh_share(0, &scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .expect("valid DLEQ proof should be accepted");

        // Verify the share was stored
        let result = signer.into_extractor();
        assert!(result.is_ok(), "should transition with all shares present");
    }

    #[test]
    fn test_signer_rejects_invalid_dleq() {
        let secp = Secp256k1::new();
        let scan_pk = ScanPublicKey::from(test_pubkey(10));
        let signer_pk = test_pubkey(5);
        let ecdh_share = test_pubkey(30); // Not a real ECDH share

        // Garbage DLEQ proof
        let proof = DleqProof::from_bytes([0xAB; 64]);

        let signer = build_signer(1, 1, 10, 20);
        let result = signer.add_ecdh_share(0, &scan_pk, &ecdh_share, &proof, &signer_pk, &secp);

        assert!(
            matches!(result, Err(PsbtError::InvalidProof(_))),
            "should reject invalid DLEQ proof, got: {:?}",
            result.err()
        );
    }

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_signer_rejects_duplicate_global_and_per_input() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        // Generate a second proof for global share
        let aux_rand2 = [0x43u8; 32];
        let (ecdh_share2, proof2) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand2, &secp).unwrap();

        let signer = build_signer(1, 1, 10, 20);

        // Add global share first
        let signer = signer
            .add_global_ecdh_share(&scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        // Try to add per-input share for same scan key -- should fail
        let result = signer.add_ecdh_share(0, &scan_pk, &ecdh_share2, &proof2, &signer_pk, &secp);
        assert!(
            matches!(result, Err(PsbtError::DuplicateShareType { .. })),
            "should reject per-input share when global exists, got: {:?}",
            result.err()
        );
    }

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_signer_rejects_global_when_per_input_exists() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let aux_rand2 = [0x43u8; 32];
        let (ecdh_share2, proof2) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand2, &secp).unwrap();

        let signer = build_signer(1, 1, 10, 20);

        // Add per-input share first
        let signer = signer
            .add_ecdh_share(0, &scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        // Try to add global share for same scan key -- should fail
        let result =
            signer.add_global_ecdh_share(&scan_pk, &ecdh_share2, &proof2, &signer_pk, &secp);
        assert!(
            matches!(result, Err(PsbtError::DuplicateShareType { .. })),
            "should reject global share when per-input exists, got: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_signer_into_extractor_requires_all_shares() {
        // Signer with no shares added at all
        let signer = build_signer(1, 1, 10, 20);

        let result = signer.into_extractor();
        assert!(
            matches!(result, Err(PsbtError::IncompleteSigning)),
            "should require all ECDH shares, got: {:?}",
            result.err()
        );
    }

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_signer_into_extractor_succeeds_with_global_share() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let signer = build_signer(1, 1, 10, 20);
        let signer = signer
            .add_global_ecdh_share(&scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        let result = signer.into_extractor();
        assert!(
            result.is_ok(),
            "should transition with global share, got: {:?}",
            result.err()
        );
    }

    // --- Extractor tests ---

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_extractor_single_signer_flow() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));
        let spend_pk = SpendPublicKey::from(test_pubkey(20));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        // Build the full flow: Constructor -> Updater -> Signer -> Extractor
        let psbt = test_psbt(1, 1);
        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan_pk, &spend_pk).unwrap();
        let updater = constructor.into_updater().unwrap();
        let signer = updater.into_signer();

        // Generate and add global ECDH share
        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let signer = signer
            .add_global_ecdh_share(&scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        let extractor = signer.into_extractor().unwrap();
        let (psbt, outputs) = extractor.extract(&secp).unwrap();

        // Verify outputs
        assert_eq!(outputs.len(), 1, "should produce one extracted output");
        assert_eq!(outputs[0].output_index(), 0);
        assert_eq!(outputs[0].scan_pubkey(), scan_pk.as_inner());

        // Verify the P2TR script is valid (34 bytes: OP_1 + PUSH32 + 32-byte key)
        let script = outputs[0].script_pubkey();
        assert_eq!(script.len(), 34, "P2TR script should be 34 bytes");
        assert_eq!(script.as_bytes()[0], 0x51, "should start with OP_1");
        assert_eq!(script.as_bytes()[1], 0x20, "should have PUSH32");

        // Verify PSBT output scriptPubKey was updated
        assert_eq!(
            psbt.unsigned_tx.output[0].script_pubkey, *script,
            "PSBT output scriptPubKey should be updated"
        );
    }

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_extractor_updates_psbt_outputs() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));
        let spend_pk = SpendPublicKey::from(test_pubkey(20));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        let psbt = test_psbt(1, 2);
        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan_pk, &spend_pk).unwrap();
        let updater = constructor.into_updater().unwrap();
        let signer = updater.into_signer();

        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let signer = signer
            .add_global_ecdh_share(&scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        let extractor = signer.into_extractor().unwrap();
        let (psbt, _outputs) = extractor.extract(&secp).unwrap();

        // Output 0 should have a P2TR script (updated by extractor)
        let script0 = &psbt.unsigned_tx.output[0].script_pubkey;
        assert_eq!(script0.len(), 34, "SP output should have P2TR script");

        // Output 1 should remain unchanged (empty script, no SP info)
        let script1 = &psbt.unsigned_tx.output[1].script_pubkey;
        assert!(script1.is_empty(), "non-SP output should be unchanged");
    }

    #[cfg(feature = "experimental-dleq")]
    #[test]
    fn test_extractor_returns_metadata() {
        let secp = Secp256k1::new();
        let scan_sk = test_secret_key(10);
        let scan_pk = ScanPublicKey::from(scan_sk.public_key(&secp));
        let spend_pk = SpendPublicKey::from(test_pubkey(20));

        let signer_sk = test_secret_key(5);
        let signer_pk = signer_sk.public_key(&secp);

        let psbt = test_psbt(1, 1);
        let constructor = SpPsbtConstructor::new(psbt).unwrap();
        let constructor = constructor.add_sp_output(0, &scan_pk, &spend_pk).unwrap();
        let updater = constructor.into_updater().unwrap();
        let signer = updater.into_signer();

        let aux_rand = [0x42u8; 32];
        let (ecdh_share, proof) =
            DleqProof::generate(&signer_sk, scan_pk.as_inner(), &aux_rand, &secp).unwrap();

        let signer = signer
            .add_global_ecdh_share(&scan_pk, &ecdh_share, &proof, &signer_pk, &secp)
            .unwrap();

        let extractor = signer.into_extractor().unwrap();
        let (_psbt, outputs) = extractor.extract(&secp).unwrap();

        // Verify metadata
        assert_eq!(outputs.len(), 1);
        let output = &outputs[0];
        assert_eq!(output.output_index(), 0, "should reference output index 0");
        assert_eq!(
            output.scan_pubkey(),
            scan_pk.as_inner(),
            "should have correct scan pubkey"
        );

        // x_only_pubkey should be 32 bytes
        assert_eq!(output.x_only_pubkey().serialize().len(), 32);

        // script_pubkey should be a valid P2TR script containing the x_only_pubkey
        let script = output.script_pubkey();
        let xonly_bytes = output.x_only_pubkey().serialize();
        assert_eq!(
            &script.as_bytes()[2..34],
            &xonly_bytes,
            "P2TR script should embed the x-only pubkey"
        );
    }
}