cairo_vm/vm/runners/builtin_runner/
signature.rs

1use crate::air_private_input::{PrivateInput, PrivateInputSignature, SignatureInput};
2use crate::math_utils::div_mod;
3use crate::stdlib::{cell::RefCell, collections::HashMap, prelude::*, rc::Rc};
4
5use crate::types::builtin_name::BuiltinName;
6use crate::types::instance_definitions::ecdsa_instance_def::CELLS_PER_SIGNATURE;
7use crate::vm::errors::runner_errors::RunnerError;
8use crate::vm::runners::cairo_pie::BuiltinAdditionalData;
9use crate::Felt252;
10use crate::{
11    types::relocatable::{MaybeRelocatable, Relocatable},
12    vm::{
13        errors::memory_errors::MemoryError,
14        vm_memory::{
15            memory::{Memory, ValidationRule},
16            memory_segments::MemorySegmentManager,
17        },
18    },
19};
20use lazy_static::lazy_static;
21use num_bigint::{BigInt, Sign};
22use num_integer::div_ceil;
23use num_traits::{Num, One};
24use starknet_crypto::{verify, Signature};
25
26lazy_static! {
27    static ref EC_ORDER: BigInt = BigInt::from_str_radix(
28        "3618502788666131213697322783095070105526743751716087489154079457884512865583",
29        10
30    )
31    .unwrap();
32}
33
34#[derive(Debug, Clone)]
35pub struct SignatureBuiltinRunner {
36    pub(crate) included: bool,
37    ratio: Option<u32>,
38    base: usize,
39    pub(crate) stop_ptr: Option<usize>,
40    signatures: Rc<RefCell<HashMap<Relocatable, Signature>>>,
41}
42
43impl SignatureBuiltinRunner {
44    pub(crate) fn new(ratio: Option<u32>, included: bool) -> Self {
45        SignatureBuiltinRunner {
46            base: 0,
47            included,
48            ratio,
49            stop_ptr: None,
50            signatures: Rc::new(RefCell::new(HashMap::new())),
51        }
52    }
53
54    pub fn add_signature(
55        &mut self,
56        relocatable: Relocatable,
57        (r, s): &(Felt252, Felt252),
58    ) -> Result<(), MemoryError> {
59        let r_be_bytes = r.to_bytes_be();
60        let s_be_bytes = s.to_bytes_be();
61        let (r_felt, s_felt) = (
62            Felt252::from_bytes_be(&r_be_bytes),
63            Felt252::from_bytes_be(&s_be_bytes),
64        );
65
66        let signature = Signature {
67            r: r_felt,
68            s: s_felt,
69        };
70
71        self.signatures
72            .borrow_mut()
73            .entry(relocatable)
74            .or_insert(signature);
75
76        Ok(())
77    }
78}
79
80impl SignatureBuiltinRunner {
81    pub fn initialize_segments(&mut self, segments: &mut MemorySegmentManager) {
82        self.base = segments.add().segment_index as usize // segments.add() always returns a positive index
83    }
84
85    pub fn initial_stack(&self) -> Vec<MaybeRelocatable> {
86        if self.included {
87            vec![MaybeRelocatable::from((self.base as isize, 0))]
88        } else {
89            vec![]
90        }
91    }
92
93    pub fn base(&self) -> usize {
94        self.base
95    }
96    pub fn add_validation_rule(&self, memory: &mut Memory) {
97        let cells_per_instance = CELLS_PER_SIGNATURE;
98        let signatures = Rc::clone(&self.signatures);
99        let rule: ValidationRule = ValidationRule(Box::new(
100            move |memory: &Memory, addr: Relocatable| -> Result<Vec<Relocatable>, MemoryError> {
101                let cell_index = addr.offset % cells_per_instance as usize;
102
103                let (pubkey_addr, message_addr) = match cell_index {
104                    0 => (addr, (addr + 1)?),
105                    1 => match addr - 1 {
106                        Ok(prev_addr) => (prev_addr, addr),
107                        Err(_) => return Ok(vec![]),
108                    },
109                    _ => return Ok(vec![]),
110                };
111
112                let pubkey = match memory.get_integer(pubkey_addr) {
113                    Ok(num) => num,
114                    Err(_) if cell_index == 1 => return Ok(vec![]),
115                    _ => return Err(MemoryError::PubKeyNonInt(Box::new(pubkey_addr))),
116                };
117
118                let msg = match memory.get_integer(message_addr) {
119                    Ok(num) => num,
120                    Err(_) if cell_index == 0 => return Ok(vec![]),
121                    _ => return Err(MemoryError::MsgNonInt(Box::new(message_addr))),
122                };
123
124                let signatures_map = signatures.borrow();
125                let signature = signatures_map
126                    .get(&pubkey_addr)
127                    .ok_or_else(|| MemoryError::SignatureNotFound(Box::new(pubkey_addr)))?;
128
129                let public_key = Felt252::from_bytes_be(&pubkey.to_bytes_be());
130                let (r, s) = (signature.r, signature.s);
131                let message = Felt252::from_bytes_be(&msg.to_bytes_be());
132                match verify(&public_key, &message, &r, &s) {
133                    Ok(true) => Ok(vec![]),
134                    _ => Err(MemoryError::InvalidSignature(Box::new((
135                        format!("({}, {})", signature.r, signature.s),
136                        pubkey.into_owned(),
137                        msg.into_owned(),
138                    )))),
139                }
140            },
141        ));
142        memory.add_validation_rule(self.base, rule);
143    }
144
145    pub fn ratio(&self) -> Option<u32> {
146        self.ratio
147    }
148
149    pub fn get_used_cells(&self, segments: &MemorySegmentManager) -> Result<usize, MemoryError> {
150        segments
151            .get_segment_used_size(self.base)
152            .ok_or(MemoryError::MissingSegmentUsedSizes)
153    }
154
155    pub fn get_used_instances(
156        &self,
157        segments: &MemorySegmentManager,
158    ) -> Result<usize, MemoryError> {
159        let used_cells = self.get_used_cells(segments)?;
160        Ok(div_ceil(used_cells, CELLS_PER_SIGNATURE as usize))
161    }
162
163    pub fn get_additional_data(&self) -> BuiltinAdditionalData {
164        // Convert signatures to Felt tuple
165        let signatures: HashMap<Relocatable, (Felt252, Felt252)> = self
166            .signatures
167            .borrow()
168            .iter()
169            .map(|(k, v)| {
170                (
171                    *k,
172                    (
173                        Felt252::from_bytes_be(&v.r.to_bytes_be()),
174                        Felt252::from_bytes_be(&v.s.to_bytes_be()),
175                    ),
176                )
177            })
178            .collect();
179        BuiltinAdditionalData::Signature(signatures)
180    }
181
182    pub fn extend_additional_data(
183        &mut self,
184        additional_data: &BuiltinAdditionalData,
185    ) -> Result<(), RunnerError> {
186        let additional_data = match additional_data {
187            BuiltinAdditionalData::Signature(d) => d,
188            BuiltinAdditionalData::Empty(_) => return Ok(()),
189            _ => return Err(RunnerError::InvalidAdditionalData(BuiltinName::ecdsa)),
190        };
191        for (addr, (r, s)) in additional_data {
192            if addr.segment_index != self.base as isize {
193                return Err(RunnerError::InvalidAdditionalData(BuiltinName::ecdsa));
194            }
195            self.signatures.borrow_mut().insert(
196                *addr,
197                Signature {
198                    r: Felt252::from_bytes_be(&r.to_bytes_be()),
199                    s: Felt252::from_bytes_be(&s.to_bytes_be()),
200                },
201            );
202        }
203        Ok(())
204    }
205
206    pub fn air_private_input(&self, memory: &Memory) -> Vec<PrivateInput> {
207        let mut private_inputs = vec![];
208
209        // Collect and sort the signatures by their index before the loop
210        let binding = self.signatures.borrow();
211        let mut sorted_signatures: Vec<_> = binding.iter().collect();
212        sorted_signatures.sort_by_key(|(addr, _)| {
213            addr.offset
214                .checked_div(CELLS_PER_SIGNATURE as usize)
215                .unwrap_or_default()
216        });
217
218        for (addr, signature) in sorted_signatures {
219            if let (Ok(pubkey), Some(msg)) = (
220                memory.get_integer(*addr),
221                (*addr + 1_usize)
222                    .ok()
223                    .and_then(|addr| memory.get_integer(addr).ok()),
224            ) {
225                private_inputs.push(PrivateInput::Signature(PrivateInputSignature {
226                    index: addr
227                        .offset
228                        .checked_div(CELLS_PER_SIGNATURE as usize)
229                        .unwrap_or_default(),
230                    pubkey: *pubkey,
231                    msg: *msg,
232                    signature_input: SignatureInput {
233                        r: Felt252::from_bytes_be(&signature.r.to_bytes_be()),
234                        w: Felt252::from(
235                            &div_mod(
236                                &BigInt::one(),
237                                &BigInt::from_bytes_be(Sign::Plus, &signature.s.to_bytes_be()),
238                                &EC_ORDER,
239                            )
240                            .unwrap_or_default(),
241                        ),
242                    },
243                }))
244            }
245        }
246        private_inputs
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::{
254        relocatable,
255        types::builtin_name::BuiltinName,
256        utils::test_utils::*,
257        vm::{
258            errors::{
259                memory_errors::{InsufficientAllocatedCellsError, MemoryError},
260                runner_errors::RunnerError,
261            },
262            runners::builtin_runner::BuiltinRunner,
263            vm_memory::{memory::Memory, memory_segments::MemorySegmentManager},
264        },
265    };
266
267    use crate::felt_str;
268    #[cfg(target_arch = "wasm32")]
269    use wasm_bindgen_test::*;
270
271    #[test]
272    fn get_used_cells_and_allocated_size_valid() {
273        let builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(10), true).into();
274        let mut vm = vm!();
275        vm.current_step = 110;
276        vm.segments.segment_used_sizes = Some(vec![1]);
277        assert_eq!(builtin.get_used_cells_and_allocated_size(&vm), Ok((1, 22)));
278    }
279
280    #[test]
281    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
282    fn initialize_segments_for_ecdsa() {
283        let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
284        let mut segments = MemorySegmentManager::new();
285        builtin.initialize_segments(&mut segments);
286        assert_eq!(builtin.base, 0);
287    }
288
289    #[test]
290    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
291    fn get_used_instances() {
292        let builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
293
294        let mut vm = vm!();
295        vm.segments.segment_used_sizes = Some(vec![1]);
296
297        assert_eq!(builtin.get_used_instances(&vm.segments), Ok(1));
298    }
299
300    #[test]
301    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
302    fn final_stack() {
303        let mut builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
304
305        let mut vm = vm!();
306
307        vm.segments = segments![
308            ((0, 0), (0, 0)),
309            ((0, 1), (0, 1)),
310            ((2, 0), (0, 0)),
311            ((2, 1), (0, 0))
312        ];
313
314        vm.segments.segment_used_sizes = Some(vec![0]);
315
316        let pointer = Relocatable::from((2, 2));
317
318        assert_eq!(
319            builtin.final_stack(&vm.segments, pointer).unwrap(),
320            Relocatable::from((2, 1))
321        );
322    }
323
324    #[test]
325    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
326    fn final_stack_error_stop_pointer() {
327        let mut builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
328
329        let mut vm = vm!();
330
331        vm.segments = segments![
332            ((0, 0), (0, 0)),
333            ((0, 1), (0, 1)),
334            ((2, 0), (0, 0)),
335            ((2, 1), (0, 0))
336        ];
337
338        vm.segments.segment_used_sizes = Some(vec![998]);
339
340        let pointer = Relocatable::from((2, 2));
341
342        assert_eq!(
343            builtin.final_stack(&vm.segments, pointer),
344            Err(RunnerError::InvalidStopPointer(Box::new((
345                BuiltinName::ecdsa,
346                relocatable!(0, 998),
347                relocatable!(0, 0)
348            ))))
349        );
350    }
351
352    #[test]
353    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
354    fn final_stack_error_non_relocatable() {
355        let mut builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
356
357        let mut vm = vm!();
358
359        vm.segments = segments![
360            ((0, 0), (0, 0)),
361            ((0, 1), (0, 1)),
362            ((2, 0), (0, 0)),
363            ((2, 1), 2)
364        ];
365
366        vm.segments.segment_used_sizes = Some(vec![0]);
367
368        let pointer = Relocatable::from((2, 2));
369
370        assert_eq!(
371            builtin.final_stack(&vm.segments, pointer),
372            Err(RunnerError::NoStopPointer(Box::new(BuiltinName::ecdsa)))
373        );
374    }
375
376    #[test]
377    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
378    fn get_used_cells_missing_segment_used_sizes() {
379        let builtin = BuiltinRunner::Signature(SignatureBuiltinRunner::new(Some(512), true));
380        let vm = vm!();
381
382        assert_eq!(
383            builtin.get_used_cells(&vm.segments),
384            Err(MemoryError::MissingSegmentUsedSizes)
385        );
386    }
387
388    #[test]
389    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
390    fn get_used_cells_empty() {
391        let builtin = BuiltinRunner::Signature(SignatureBuiltinRunner::new(Some(512), true));
392        let mut vm = vm!();
393
394        vm.segments.segment_used_sizes = Some(vec![0]);
395        assert_eq!(builtin.get_used_cells(&vm.segments), Ok(0));
396    }
397
398    #[test]
399    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
400    fn get_used_cells() {
401        let builtin = BuiltinRunner::Signature(SignatureBuiltinRunner::new(Some(512), true));
402        let mut vm = vm!();
403
404        vm.segments.segment_used_sizes = Some(vec![4]);
405        assert_eq!(builtin.get_used_cells(&vm.segments), Ok(4));
406    }
407
408    #[test]
409    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
410    fn get_initial_stack_for_range_check_with_base() {
411        let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
412        builtin.base = 1;
413        let initial_stack = builtin.initial_stack();
414        assert_eq!(
415            initial_stack[0].clone(),
416            MaybeRelocatable::RelocatableValue((builtin.base() as isize, 0).into())
417        );
418        assert_eq!(initial_stack.len(), 1);
419    }
420
421    #[test]
422    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
423    fn initial_stack_not_included_test() {
424        let ecdsa_builtin = SignatureBuiltinRunner::new(Some(512), false);
425        assert_eq!(ecdsa_builtin.initial_stack(), Vec::new())
426    }
427
428    #[test]
429    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
430    fn deduce_memory_cell_test() {
431        let memory = Memory::new();
432        let builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
433        let result = builtin.deduce_memory_cell(Relocatable::from((0, 5)), &memory);
434        assert_eq!(result, Ok(None));
435    }
436
437    #[test]
438    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
439    fn test_ratio() {
440        let builtin = SignatureBuiltinRunner::new(Some(512), true);
441        assert_eq!(builtin.ratio(), Some(512));
442    }
443
444    #[test]
445    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
446    fn test_base() {
447        let builtin = SignatureBuiltinRunner::new(Some(512), true);
448        assert_eq!(builtin.base(), 0);
449    }
450
451    #[test]
452    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
453    fn get_allocated_memory_min_step_not_reached() {
454        let builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
455        let mut vm = vm!();
456        vm.current_step = 500;
457        assert_eq!(
458            builtin.get_allocated_memory_units(&vm),
459            Err(MemoryError::InsufficientAllocatedCells(
460                InsufficientAllocatedCellsError::MinStepNotReached(Box::new((
461                    512,
462                    BuiltinName::ecdsa
463                )))
464            ))
465        )
466    }
467
468    #[test]
469    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
470    fn get_used_cells_and_allocated_size_insufficient_allocated() {
471        let builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
472        let mut vm = vm!();
473        vm.segments.segment_used_sizes = Some(vec![50]);
474        vm.current_step = 512;
475        assert_eq!(
476            builtin.get_used_cells_and_allocated_size(&vm),
477            Err(MemoryError::InsufficientAllocatedCells(
478                InsufficientAllocatedCellsError::BuiltinCells(Box::new((
479                    BuiltinName::ecdsa,
480                    50,
481                    2
482                )))
483            ))
484        )
485    }
486
487    #[test]
488    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
489    fn final_stack_invalid_stop_pointer() {
490        let mut builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
491        let mut vm = vm!();
492        vm.segments = segments![((0, 0), (1, 0))];
493        assert_eq!(
494            builtin.final_stack(&vm.segments, (0, 1).into()),
495            Err(RunnerError::InvalidStopPointerIndex(Box::new((
496                BuiltinName::ecdsa,
497                relocatable!(1, 0),
498                0
499            ))))
500        )
501    }
502
503    #[test]
504    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
505    fn final_stack_no_used_instances() {
506        let mut builtin: BuiltinRunner = SignatureBuiltinRunner::new(Some(512), true).into();
507        let mut vm = vm!();
508        vm.segments = segments![((0, 0), (0, 0))];
509        assert_eq!(
510            builtin.final_stack(&vm.segments, (0, 1).into()),
511            Err(RunnerError::Memory(MemoryError::MissingSegmentUsedSizes))
512        )
513    }
514
515    #[test]
516    fn get_additional_data() {
517        let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
518        let signatures = HashMap::from([(
519            Relocatable::from((4, 0)),
520            Signature {
521                r: Felt252::from_dec_str("45678").unwrap(),
522                s: Felt252::from_dec_str("1239").unwrap(),
523            },
524        )]);
525        builtin.signatures = Rc::new(RefCell::new(signatures));
526        let signatures = HashMap::from([(
527            Relocatable::from((4, 0)),
528            (felt_str!("45678"), felt_str!("1239")),
529        )]);
530        assert_eq!(
531            builtin.get_additional_data(),
532            BuiltinAdditionalData::Signature(signatures)
533        )
534    }
535
536    #[test]
537    fn get_and_extend_additional_data() {
538        let mut builtin_a = SignatureBuiltinRunner::new(Some(512), true);
539        let signatures = HashMap::from([(
540            Relocatable::from((0, 0)),
541            Signature {
542                r: Felt252::from_dec_str("45678").unwrap(),
543                s: Felt252::from_dec_str("1239").unwrap(),
544            },
545        )]);
546        builtin_a.signatures = Rc::new(RefCell::new(signatures));
547        let additional_data = builtin_a.get_additional_data();
548        let mut builtin_b = SignatureBuiltinRunner::new(Some(512), true);
549        builtin_b.extend_additional_data(&additional_data).unwrap();
550        // Signature doesn't implement PartialEq so we can't comapre the list of signatures directly
551        let signatures_a = builtin_a.signatures.borrow();
552        let signatures_b = builtin_b.signatures.borrow();
553        assert_eq!(signatures_a.len(), signatures_b.len());
554        for ((addr_a, signature_a), (addr_b, signature_b)) in
555            signatures_a.iter().zip(signatures_b.iter())
556        {
557            assert_eq!(addr_a, addr_b);
558            assert_eq!(signature_a.r, signature_b.r);
559            assert_eq!(signature_a.s, signature_b.s);
560        }
561    }
562    #[test]
563    fn get_air_private_input() {
564        let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
565
566        builtin.base = 0;
567
568        let signature1_r = Felt252::from(1234);
569        let signature1_s = Felt252::from(5678);
570        let signature2_r = Felt252::from(8765);
571        let signature2_s = Felt252::from(4321);
572
573        let sig1_addr = Relocatable::from((builtin.base as isize, 0));
574        let sig2_addr = Relocatable::from((builtin.base as isize, CELLS_PER_SIGNATURE as usize));
575
576        builtin
577            .add_signature(sig1_addr, &(signature1_r, signature1_s))
578            .unwrap();
579        builtin
580            .add_signature(sig2_addr, &(signature2_r, signature2_s))
581            .unwrap();
582
583        let pubkey1 = Felt252::from(1111);
584        let msg1 = Felt252::from(2222);
585        let pubkey2 = Felt252::from(3333);
586        let msg2 = Felt252::from(4444);
587
588        let segments = segments![
589            ((0, 0), 1111),
590            ((0, 1), 2222),
591            ((0, 2), 3333),
592            ((0, 3), 4444)
593        ];
594        let w1 =
595            Felt252::from(&div_mod(&BigInt::one(), &signature1_s.to_bigint(), &EC_ORDER).unwrap());
596
597        let w2 =
598            Felt252::from(&div_mod(&BigInt::one(), &signature2_s.to_bigint(), &EC_ORDER).unwrap());
599
600        let expected_private_inputs = vec![
601            PrivateInput::Signature(PrivateInputSignature {
602                index: 0,
603                pubkey: pubkey1,
604                msg: msg1,
605                signature_input: SignatureInput {
606                    r: signature1_r,
607                    w: w1,
608                },
609            }),
610            PrivateInput::Signature(PrivateInputSignature {
611                index: 1,
612                pubkey: pubkey2,
613                msg: msg2,
614                signature_input: SignatureInput {
615                    r: signature2_r,
616                    w: w2,
617                },
618            }),
619        ];
620
621        let private_inputs = builtin.air_private_input(&segments.memory);
622
623        assert_eq!(private_inputs, expected_private_inputs);
624    }
625}
cairo_vm/vm/runners/builtin_runner/signature.rs

cairo_vm/vm/runners/builtin_runner/
signature.rs
