Skip to main content

keymaster_multisig/
multisig.rs

1use crate::error::{MultisigError, Result};
2use crate::types::{PrivateKey, PublicKey, Transaction};
3use k256::{
4    ecdsa::{
5        signature::{hazmat::PrehashSigner, SignatureEncoding},
6        Signature as EcdsaSignature, SigningKey,
7    },
8    SecretKey,
9};
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12
13const OP_0: u8 = 0x00;
14const OP_CHECKMULTISIG: u8 = 0xae;
15const SIGHASH_ALL_FORKID: u8 = 0x41;
16
17/// Variable length integer encoding (Bitcoin style)
18#[derive(Debug, Clone)]
19struct VarInt(pub u64);
20
21impl VarInt {
22    pub fn serialize(&self) -> Vec<u8> {
23        match self.0 {
24            0x00..=0xFC => vec![self.0 as u8],
25            0xFD..=0xFFFF => {
26                let mut v = vec![0xFD];
27                v.extend_from_slice(&(self.0 as u16).to_le_bytes());
28                v
29            }
30            0x10000..=0xFFFFFFFF => {
31                let mut v = vec![0xFE];
32                v.extend_from_slice(&(self.0 as u32).to_le_bytes());
33                v
34            }
35            _ => {
36                let mut v = vec![0xFF];
37                v.extend_from_slice(&self.0.to_le_bytes());
38                v
39            }
40        }
41    }
42}
43
44#[derive(Serialize, Deserialize, Debug)]
45pub struct Multisig {
46    private_keys: Option<Vec<PrivateKey>>,
47    public_keys: Vec<PublicKey>,
48    m: usize,
49    n: usize,
50    sig_hash_type: u8,
51}
52
53impl Multisig {
54    pub fn new(
55        private_keys: Option<Vec<PrivateKey>>,
56        public_keys: Vec<PublicKey>,
57        m: usize,
58    ) -> Result<Self> {
59        if public_keys.is_empty() || public_keys.len() > 20 {
60            return Err(MultisigError::InvalidPublicKeys);
61        }
62
63        if m == 0 || m > public_keys.len() {
64            return Err(MultisigError::InvalidM(format!(
65                "m={} must be between 1 and n={}",
66                m,
67                public_keys.len()
68            )));
69        }
70
71        if let Some(ref keys) = private_keys {
72            if keys.len() < m {
73                return Err(MultisigError::NoPrivateKeys);
74            }
75        }
76
77        let n = public_keys.len();
78        Ok(Multisig {
79            private_keys,
80            public_keys,
81            m,
82            n,
83            sig_hash_type: SIGHASH_ALL_FORKID,
84        })
85    }
86
87    pub fn lock(&self) -> Result<Vec<u8>> {
88        if self.m == 0 || self.m > self.n {
89            return Err(MultisigError::InvalidM(format!(
90                "m={} must be between 1 and n={}",
91                self.m, self.n
92            )));
93        }
94        if self.n == 0 || self.n > 20 {
95            return Err(MultisigError::InvalidPublicKeys);
96        }
97
98        let mut script = Vec::new();
99
100        script.push(0x01 + (self.m as u8) - 1);
101
102        for pub_key in &self.public_keys {
103            script.push(pub_key.key.len() as u8);
104            script.extend(&pub_key.key);
105        }
106
107        script.push(0x01 + (self.n as u8) - 1);
108        script.push(OP_CHECKMULTISIG);
109
110        Ok(script)
111    }
112
113    pub fn sign(&self, tx: &Transaction, input_index: usize) -> Result<Vec<Vec<u8>>> {
114        if let Some(ref priv_keys) = self.private_keys {
115            if priv_keys.len() < self.m {
116                return Err(MultisigError::NoPrivateKeys);
117            }
118
119            let mut signatures = Vec::new();
120
121            for private_key in priv_keys.iter().take(self.m) {
122                let sig = self.sign_one(tx, input_index, private_key)?;
123                signatures.push(sig);
124            }
125
126            Ok(signatures)
127        } else {
128            Err(MultisigError::NoPrivateKeys)
129        }
130    }
131
132    pub fn sign_one(
133        &self,
134        tx: &Transaction,
135        input_index: usize,
136        private_key: &PrivateKey,
137    ) -> Result<Vec<u8>> {
138        if input_index >= tx.inputs.len() {
139            return Err(MultisigError::TransactionError(
140                "Input index out of bounds".to_string(),
141            ));
142        }
143
144        let sighash = self.calculate_signature_hash(tx, input_index)?;
145
146        let signature = self.generate_signature(&sighash, private_key)?;
147
148        Ok(signature)
149    }
150
151    fn calculate_signature_hash(&self, tx: &Transaction, input_index: usize) -> Result<Vec<u8>> {
152        // Simplified signature hash calculation for Bitcoin SV
153        let mut hash_input = Vec::new();
154        hash_input.extend_from_slice(&tx.version.to_le_bytes());
155
156        // Serialize inputs
157        let inputs_count = VarInt(tx.inputs.len() as u64);
158        hash_input.extend(inputs_count.serialize());
159
160        for (i, input) in tx.inputs.iter().enumerate() {
161            hash_input.extend_from_slice(
162                &hex::decode(&input.source_txid).map_err(|_| {
163                    MultisigError::TransactionError("Invalid source txid".to_string())
164                })?,
165            );
166            hash_input.extend_from_slice(&input.source_output_index.to_le_bytes());
167
168            if i == input_index {
169                // For the input being signed, use empty unlocking script for SIGHASH calculation
170                hash_input.extend(VarInt(0).serialize());
171            } else {
172                // For other inputs, use placeholder script
173                hash_input.extend(VarInt(0).serialize());
174            }
175
176            hash_input.extend_from_slice(&input.sequence.to_le_bytes());
177        }
178
179        // Serialize outputs
180        let outputs_count = VarInt(tx.outputs.len() as u64);
181        hash_input.extend(outputs_count.serialize());
182
183        for output in &tx.outputs {
184            hash_input.extend_from_slice(&output.satoshis.to_le_bytes());
185            let script_len = VarInt(output.locking_script.len() as u64);
186            hash_input.extend(script_len.serialize());
187            hash_input.extend(&output.locking_script);
188        }
189
190        hash_input.extend_from_slice(&tx.lock_time.to_le_bytes());
191        hash_input.push(self.sig_hash_type);
192
193        // Double SHA256 for Bitcoin
194        let hash1 = Sha256::digest(&hash_input);
195        let hash2 = Sha256::digest(hash1);
196        Ok(hash2.to_vec())
197    }
198
199    fn generate_signature(&self, sighash: &[u8], private_key: &PrivateKey) -> Result<Vec<u8>> {
200        // Convert private key bytes to SecretKey
201        let secret_key = SecretKey::from_slice(&private_key.key)
202            .map_err(|_| MultisigError::InvalidPrivateKey)?;
203
204        let signing_key = SigningKey::from(secret_key);
205        let signature: EcdsaSignature = signing_key
206            .sign_prehash(sighash)
207            .map_err(|_| MultisigError::SignatureError("Failed to create signature".to_string()))?;
208
209        // Convert to DER format and add SIGHASH type
210        let der_sig = signature.to_der();
211        let mut sig_with_hash = der_sig.to_vec();
212        sig_with_hash.push(self.sig_hash_type);
213
214        Ok(sig_with_hash)
215    }
216
217    pub fn estimate_length(&self) -> usize {
218        1 + self.m * (71 + 1)
219    }
220
221    pub fn create_fake_sign(&self) -> Result<Vec<u8>> {
222        let mut script = vec![OP_0];
223
224        for _ in 0..self.m {
225            script.extend(vec![0u8; 72]);
226            script.push(self.sig_hash_type);
227        }
228
229        Ok(script)
230    }
231
232    pub fn build_sign_script(&self, signatures: &[Vec<u8>]) -> Result<Vec<u8>> {
233        let mut script = vec![OP_0];
234
235        for sig in signatures {
236            script.push(sig.len() as u8);
237            script.extend(sig);
238        }
239
240        Ok(script)
241    }
242
243    pub fn get_m(&self) -> usize {
244        self.m
245    }
246
247    pub fn get_n(&self) -> usize {
248        self.n
249    }
250
251    pub fn get_sig_hash_type(&self) -> u8 {
252        self.sig_hash_type
253    }
254
255    pub fn get_public_keys(&self) -> &[PublicKey] {
256        &self.public_keys
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::Multisig;
263    use crate::types::{PrivateKey, PublicKey, Transaction, TransactionInput, TransactionOutput};
264
265    #[test]
266    fn supports_all_two_of_three_signature_pairs() {
267        let public_keys = vec![
268            PublicKey::new(vec![0x02; 33]),
269            PublicKey::new(vec![0x03; 33]),
270            PublicKey::new(vec![0x04; 33]),
271        ];
272        let transaction = Transaction::new(
273            1,
274            vec![TransactionInput::new("aa".repeat(32), 0, 1)],
275            vec![TransactionOutput::new(1000, vec![0x51])],
276            0,
277        );
278        let signer = Multisig::new(None, public_keys, 2).unwrap();
279        let buyer = signer
280            .sign_one(&transaction, 0, &PrivateKey::new(vec![1; 32]))
281            .unwrap();
282        let seller = signer
283            .sign_one(&transaction, 0, &PrivateKey::new(vec![2; 32]))
284            .unwrap();
285        let arbiter = signer
286            .sign_one(&transaction, 0, &PrivateKey::new(vec![3; 32]))
287            .unwrap();
288
289        let buyer_seller = signer
290            .build_sign_script(&[buyer.clone(), seller.clone()])
291            .unwrap();
292        let buyer_arbiter = signer
293            .build_sign_script(&[buyer.clone(), arbiter.clone()])
294            .unwrap();
295        let seller_arbiter = signer.build_sign_script(&[seller, arbiter]).unwrap();
296
297        assert_eq!(buyer_seller[0], 0);
298        assert_eq!(buyer_arbiter[0], 0);
299        assert_eq!(seller_arbiter[0], 0);
300        assert_ne!(buyer_seller, buyer_arbiter);
301        assert_ne!(buyer_arbiter, seller_arbiter);
302    }
303}