1use std::collections::{HashMap, HashSet};
4
5use sdkt_wasm::ContractSpec;
6use syn::visit::{self, Visit};
7use syn::{Expr, ExprCall, ExprMethodCall, FnArg, Item, Local, Pat};
8
9use crate::error::AuditError;
10use crate::types::{AuditReport, Severity};
11
12pub struct AuditContext<'a> {
15 pub spec: Option<&'a ContractSpec>,
16}
17
18pub trait AuditRule {
23 fn id(&self) -> &'static str;
25 fn severity(&self) -> Severity;
27 fn description(&self) -> &'static str;
29 fn check(&self, scans: &[FnScan], ctx: &AuditContext, report: &mut AuditReport);
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41enum OrderEvent {
42 ExternalCall,
44 StateWrite,
46}
47
48#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
50pub struct FnScan {
51 pub fn_name: String,
52 pub require_auth: usize,
53 pub invoke_contract: usize,
54 pub bound: HashSet<String>,
56 pub usage: HashMap<String, usize>,
58 #[serde(default)]
60 pub division_before_multiplication: bool,
61 #[serde(default)]
67 pub external_call_before_state_write: bool,
68}
69
70impl FnScan {
71 fn new(fn_name: String) -> Self {
72 Self {
73 fn_name,
74 require_auth: 0,
75 invoke_contract: 0,
76 bound: HashSet::new(),
77 usage: HashMap::new(),
78 division_before_multiplication: false,
79 external_call_before_state_write: false,
80 }
81 }
82}
83
84fn is_auth_fn(name: &str) -> bool {
86 matches!(
87 name,
88 "require_auth" | "require_auth_for_args" | "require_auth_for_caller"
89 )
90}
91
92pub(crate) fn is_privileged(name: &str) -> bool {
94 let n = unqualified(name).to_lowercase();
95 n == "initialize"
96 || n.starts_with("initialize_")
97 || n == "init"
98 || n.starts_with("init_")
99 || [
100 "admin",
101 "mint",
102 "burn",
103 "pause",
104 "unpause",
105 "upgrade",
106 "withdraw",
107 "configure",
108 "freeze",
109 "ban",
110 "set_admin",
111 "owner_set",
112 "transfer_ownership",
113 "change_owner",
114 "set_auth",
115 ]
116 .iter()
117 .any(|p| n.contains(p))
118}
119
120pub(crate) fn is_initialize(name: &str) -> bool {
122 let n = unqualified(name).to_lowercase();
123 n == "initialize" || n.starts_with("initialize_") || n == "init" || n.starts_with("init_")
124}
125
126pub(crate) fn unqualified(name: &str) -> &str {
128 name.rsplit("::").next().unwrap_or(name)
129}
130
131struct FnVisitor<'a> {
134 scan: &'a mut FnScan,
135 divided_bindings: HashSet<String>,
136 order_events: &'a mut Vec<OrderEvent>,
138}
139
140impl<'ast, 'a> Visit<'ast> for FnVisitor<'a> {
141 fn visit_local(&mut self, node: &'ast Local) {
142 if let Pat::Ident(pat) = &node.pat {
143 let name = pat.ident.to_string();
144 self.scan.bound.insert(name.clone());
145 if let Some(init) = &node.init {
146 if expr_contains_division(&init.expr)
147 || expr_uses_any(&init.expr, &self.divided_bindings)
148 {
149 self.divided_bindings.insert(name);
150 } else {
151 self.divided_bindings.remove(&name);
152 }
153 }
154 }
155 visit::visit_local(self, node);
156 }
157
158 fn visit_expr_binary(&mut self, node: &'ast syn::ExprBinary) {
159 if matches!(node.op, syn::BinOp::Mul(_))
160 && (expr_contains_division(&node.left)
161 || expr_contains_division(&node.right)
162 || expr_uses_any(&node.left, &self.divided_bindings)
163 || expr_uses_any(&node.right, &self.divided_bindings))
164 {
165 self.scan.division_before_multiplication = true;
166 }
167 visit::visit_expr_binary(self, node);
168 }
169
170 fn visit_expr_call(&mut self, node: &'ast ExprCall) {
171 if let Expr::Path(p) = &*node.func {
172 if let Some(seg) = p.path.segments.last() {
173 let name = seg.ident.to_string();
174 if is_auth_fn(&name) {
175 self.scan.require_auth += 1;
176 }
177 }
178 }
179 for arg in &node.args {
181 self.count_ident(arg);
182 visit::visit_expr(self, arg);
183 }
184 if let Expr::Path(p) = &*node.func {
186 if let Some(seg) = p.path.segments.last() {
187 let name = seg.ident.to_string();
188 if name == "invoke_contract" || name == "invoke_contract_light" {
189 self.scan.invoke_contract += 1;
190 self.order_events.push(OrderEvent::ExternalCall);
191 }
192 }
193 }
194 }
196
197 fn visit_expr_method_call(&mut self, node: &'ast ExprMethodCall) {
198 let method = node.method.to_string();
199 if is_auth_fn(&method) {
200 self.scan.require_auth += 1;
201 }
202
203 visit::visit_expr(self, &node.receiver);
205 for arg in &node.args {
207 self.count_ident(arg);
208 visit::visit_expr(self, arg);
209 }
210
211 if method == "invoke_contract" || method == "invoke_contract_light" {
213 self.scan.invoke_contract += 1;
214 self.order_events.push(OrderEvent::ExternalCall);
215 }
216 self.detect_state_write(&method);
218 }
220}
221
222fn expr_contains_division(expr: &Expr) -> bool {
223 struct Finder(bool);
224 impl<'ast> Visit<'ast> for Finder {
225 fn visit_expr_binary(&mut self, node: &'ast syn::ExprBinary) {
226 if matches!(node.op, syn::BinOp::Div(_))
227 && !(is_constant_expr(&node.left) && is_constant_expr(&node.right))
228 {
229 self.0 = true;
230 }
231 visit::visit_expr_binary(self, node);
232 }
233 }
234 let mut finder = Finder(false);
235 finder.visit_expr(expr);
236 finder.0
237}
238
239fn is_constant_expr(expr: &Expr) -> bool {
240 match expr {
241 Expr::Lit(literal) => matches!(literal.lit, syn::Lit::Int(_)),
242 Expr::Paren(paren) => is_constant_expr(&paren.expr),
243 Expr::Group(group) => is_constant_expr(&group.expr),
244 Expr::Unary(unary) => is_constant_expr(&unary.expr),
245 Expr::Binary(binary) => is_constant_expr(&binary.left) && is_constant_expr(&binary.right),
246 _ => false,
247 }
248}
249
250fn expr_uses_any(expr: &Expr, names: &HashSet<String>) -> bool {
251 struct Finder<'a> {
252 names: &'a HashSet<String>,
253 found: bool,
254 }
255 impl<'ast> Visit<'ast> for Finder<'_> {
256 fn visit_expr_path(&mut self, node: &'ast syn::ExprPath) {
257 if let Some(name) = node.path.get_ident() {
258 self.found |= self.names.contains(&name.to_string());
259 }
260 visit::visit_expr_path(self, node);
261 }
262 }
263 let mut finder = Finder {
264 names,
265 found: false,
266 };
267 finder.visit_expr(expr);
268 finder.found
269}
270
271impl<'a> FnVisitor<'a> {
272 fn count_ident(&mut self, expr: &Expr) {
273 if let Expr::Path(p) = expr {
274 if let Some(ident) = p.path.get_ident() {
275 let s = ident.to_string();
276 if self.scan.bound.contains(&s) {
277 *self.scan.usage.entry(s).or_insert(0) += 1;
278 }
279 }
280 }
281 }
282
283 fn detect_state_write(&mut self, method_name: &str) {
296 let patterns = ["set", "write", "put", "remove", "delete", "extend"];
301 if patterns.iter().any(|p| method_name.eq_ignore_ascii_case(p)) {
302 self.order_events.push(OrderEvent::StateWrite);
303 }
304 }
305}
306
307pub fn scan_all_functions(ast: &syn::File) -> Vec<FnScan> {
312 let mut out = Vec::new();
313 for item in &ast.items {
314 match item {
315 Item::Fn(f) => {
316 let scan = fn_scan(&f.sig, &f.block);
317 out.push(scan);
318 }
319 Item::Impl(imp) => {
320 let type_name = match &*imp.self_ty {
322 syn::Type::Path(p) => p
323 .path
324 .segments
325 .last()
326 .map(|s| s.ident.to_string())
327 .unwrap_or_default(),
328 _ => String::new(),
329 };
330 for inner in &imp.items {
331 if let syn::ImplItem::Fn(m) = inner {
332 let qual = if type_name.is_empty() {
333 m.sig.ident.to_string()
334 } else {
335 format!("{}::{}", type_name, m.sig.ident)
336 };
337 let mut scan = fn_scan(&m.sig, &m.block);
338 scan.fn_name = qual;
339 out.push(scan);
340 }
341 }
342 }
343 _ => {}
344 }
345 }
346 out
347}
348
349fn fn_scan(sig: &syn::Signature, block: &syn::Block) -> FnScan {
351 let mut scan = FnScan::new(sig.ident.to_string());
352 for input in &sig.inputs {
353 if let FnArg::Typed(pt) = input {
354 if let Pat::Ident(pat) = &*pt.pat {
355 scan.bound.insert(pat.ident.to_string());
356 }
357 }
358 }
359 let mut order_events = Vec::new();
360 let mut visitor = FnVisitor {
361 scan: &mut scan,
362 divided_bindings: HashSet::new(),
363 order_events: &mut order_events,
364 };
365 visitor.visit_block(block);
366
367 let mut found_external_call = false;
370 for event in order_events {
371 match event {
372 OrderEvent::ExternalCall => {
373 found_external_call = true;
374 }
375 OrderEvent::StateWrite if found_external_call => {
376 scan.external_call_before_state_write = true;
378 break;
379 }
380 _ => {}
381 }
382 }
383
384 scan
385}
386
387pub fn scan_all_functions_str(src: &str) -> Option<Vec<FnScan>> {
391 let ast = syn::parse_file(src).ok()?;
392 Some(scan_all_functions(&ast))
393}
394
395pub fn all_rules() -> Vec<Box<dyn AuditRule>> {
397 vec![
398 Box::new(crate::rules::Auth001),
399 Box::new(crate::rules::Auth002),
400 Box::new(crate::rules::Auth003),
401 Box::new(crate::rules::Auth004),
402 Box::new(crate::rules::Move001),
403 Box::new(crate::rules::Math001),
404 ]
405}
406
407fn run_rules(
408 ast: &syn::File,
409 ctx: &AuditContext,
410 disabled: &[&str],
411) -> Result<AuditReport, AuditError> {
412 let scans = scan_all_functions(ast);
413 let mut report = AuditReport::default();
414 crate::registry::run_registered(&scans, ctx, disabled, &mut report);
417 Ok(report)
418}
419
420pub fn audit_source(src: &str) -> Result<AuditReport, AuditError> {
422 audit_source_with(src, &[])
423}
424
425pub fn audit_source_with(src: &str, disabled: &[&str]) -> Result<AuditReport, AuditError> {
427 let ast = syn::parse_file(src).map_err(AuditError::Parse)?;
428 let ctx = AuditContext { spec: None };
429 run_rules(&ast, &ctx, disabled)
430}
431
432pub fn audit_source_with_registry(
433 src: &str,
434 reg: &crate::registry::RuleRegistry,
435 disabled: &[&str],
436) -> Result<AuditReport, AuditError> {
437 let ast = syn::parse_file(src).map_err(AuditError::Parse)?;
438 let ctx = AuditContext { spec: None };
439 let scans = scan_all_functions(&ast);
440 let mut report = AuditReport::default();
441 reg.run_all(&scans, &ctx, disabled, &mut report);
442 Ok(report)
443}
444
445pub fn audit_source_with_spec(
447 src: &str,
448 spec: &ContractSpec,
449 disabled: &[&str],
450) -> Result<AuditReport, AuditError> {
451 let ast = syn::parse_file(src).map_err(AuditError::Parse)?;
452 let ctx = AuditContext { spec: Some(spec) };
453 run_rules(&ast, &ctx, disabled)
454}
455
456#[cfg(test)]
457mod tests {
458 use super::*;
459 use crate::rules::{Auth001, Auth002, Auth003, Move001};
460 use sdkt_wasm::spec::ContractFunction;
461
462 fn report_for(src: &str) -> AuditReport {
463 audit_source(src).unwrap()
464 }
465
466 fn has(rep: &AuditReport, id: &str) -> bool {
467 rep.findings.iter().any(|f| f.rule_id == id)
468 }
469
470 #[test]
471 fn audit_source_with_registry_uses_supplied_registry() {
472 struct Stub {
473 id: &'static str,
474 }
475 impl AuditRule for Stub {
476 fn id(&self) -> &'static str {
477 self.id
478 }
479 fn severity(&self) -> Severity {
480 Severity::Info
481 }
482 fn description(&self) -> &'static str {
483 "stub"
484 }
485 fn check(&self, _scans: &[FnScan], _ctx: &AuditContext, report: &mut AuditReport) {
486 report.add(crate::types::Finding {
487 file: None,
488 rule_id: self.id.to_string(),
489 severity: Severity::Info,
490 message: "stub fired".into(),
491 location: None,
492 });
493 }
494 }
495
496 let mut reg = crate::registry::RuleRegistry::new();
497 reg.register_builtin_rules();
498 reg.register_rule(Box::new(Stub { id: "REG-TEST" }) as crate::registry::BoxedRule);
499 let rep = audit_source_with_registry("pub fn mint() { }", ®, &["AUTH-001"]).unwrap();
500 assert!(rep.findings.iter().any(|f| f.rule_id == "REG-TEST"));
501 }
502
503 #[test]
504 fn auth001_fires_on_privileged_without_auth() {
505 let src = "pub fn mint_token(to: Address) { /* no auth */ }";
506 let rep = report_for(src);
507 assert!(has(&rep, "AUTH-001"));
508 }
509
510 #[test]
511 fn auth001_silent_when_auth_present() {
512 let src = "pub fn mint_token(to: Address) { require_auth(); }";
513 let rep = report_for(src);
514 assert!(!has(&rep, "AUTH-001"));
515 }
516
517 #[test]
518 fn auth001_disabled_is_silent() {
519 let src = "pub fn mint_token(to: Address) { }";
520 let rep = audit_source_with(src, &["AUTH-001"]).unwrap();
521 assert!(!has(&rep, "AUTH-001"));
522 }
523
524 #[test]
525 fn auth001_negative_on_normal_fn() {
526 let src = "pub fn balance_of(who: Address) -> u32 { 0 }";
527 let rep = report_for(src);
528 assert!(!has(&rep, "AUTH-001"));
529 }
530
531 #[test]
532 fn auth002_fires_on_invoke_without_auth() {
533 let src = "pub fn relay() { env.invoke_contract(&addr, sym, vec![]); }";
534 let rep = report_for(src);
535 assert!(has(&rep, "AUTH-002"));
536 }
537
538 #[test]
539 fn auth002_silent_when_auth_present() {
540 let src = "pub fn relay() { require_auth(); env.invoke_contract(&addr, sym, vec![]); }";
541 let rep = report_for(src);
542 assert!(!has(&rep, "AUTH-002"));
543 }
544
545 #[test]
546 fn auth003_fires_on_unguarded_initialize() {
547 let src = "pub fn initialize(admin: Address) { /* no auth */ }";
548 let rep = report_for(src);
549 assert!(has(&rep, "AUTH-003"));
550 }
551
552 #[test]
553 fn auth003_silent_when_initialize_has_auth() {
554 let src = "pub fn initialize(admin: Address) { require_auth(); }";
555 let rep = report_for(src);
556 assert!(!has(&rep, "AUTH-003"));
557 }
558
559 #[test]
560 fn auth003_matches_qualified_impl_method_against_unqualified_spec() {
561 let src = "impl Token { pub fn initialize(admin: Address) { } }";
562 let spec = ContractSpec {
563 env_meta: None,
564 functions: vec![ContractFunction {
565 name: "initialize".into(),
566 doc: String::new(),
567 parameters: vec![],
568 outputs: vec![],
569 }],
570 custom_types: vec![],
571 events: vec![],
572 };
573 let rep = audit_source_with_spec(src, &spec, &[]).unwrap();
574 assert!(has(&rep, "AUTH-003"));
575 }
576
577 #[test]
578 fn move001_fires_on_double_arg_use() {
579 let src = "pub fn foo(a: Address) { bar(a); baz(a); }";
580 let rep = report_for(src);
581 assert!(has(&rep, "MOVE-001"));
582 }
583
584 #[test]
585 fn move001_silent_on_single_use() {
586 let src = "pub fn foo(a: Address) { bar(a); }";
587 let rep = report_for(src);
588 assert!(!has(&rep, "MOVE-001"));
589 }
590
591 #[test]
592 fn math001_detects_direct_and_intermediate_division_before_multiplication() {
593 for src in [
594 "pub fn quote(amount: i128, bps: i128) -> i128 { amount / 10_000 * bps }",
595 "pub fn quote(amount: i128, bps: i128) -> i128 { (amount / 10_000) * (bps + 1) }",
596 "pub fn quote(amount: i128, bps: i128) -> i128 { let divided = amount / 10_000; divided * bps }",
597 ] {
598 let report = report_for(src);
599 let finding = report
600 .findings
601 .iter()
602 .find(|finding| finding.rule_id == "MATH-001")
603 .expect("division before multiplication should be reported");
604 assert_eq!(finding.severity, Severity::Warning);
605 assert_eq!(finding.location.as_deref(), Some("quote"));
606 assert!(finding.message.contains("divides before multiplying"));
607 }
608 }
609
610 #[test]
611 fn math001_ignores_unrelated_or_non_multiplicative_arithmetic() {
612 for src in [
613 "pub fn quote(amount: i128, bps: i128) -> i128 { amount * bps }",
614 "pub fn quote(bps: i128) -> i128 { 5 / 2 * bps }",
615 "pub fn quote(amount: i128) -> i128 { amount / 10_000 }",
616 "pub fn quote(amount: i128, bps: i128) -> i128 { amount / 10_000 + bps }",
617 "pub fn quote(amount: i128, other: i128, bps: i128) -> i128 { amount / 10_000 + other * bps }",
618 "pub fn quote(amount: i128, other: i128, bps: i128) -> i128 { let divided = amount / 10_000; other * bps + divided }",
619 ] {
620 assert!(
621 !has(&report_for(src), "MATH-001"),
622 "unexpected math finding for {src}"
623 );
624 }
625 }
626
627 #[test]
628 fn math001_can_be_disabled() {
629 let src = "pub fn quote(amount: i128, bps: i128) -> i128 { amount / 10_000 * bps }";
630 assert!(!has(
631 &audit_source_with(src, &["MATH-001"]).unwrap(),
632 "MATH-001"
633 ));
634 }
635
636 #[test]
637 fn cei001_fires_on_external_call_before_state_write() {
638 let src = r#"
639 pub fn swap(env: Env, to: Address) {
640 env.invoke_contract(&pool, sym!("withdraw"), args);
641 env.storage().set(key, value);
642 }
643 "#;
644 let rep = report_for(src);
645 assert!(has(&rep, "CEI-001"));
646 }
647
648 #[test]
649 fn cei001_detects_nested_invoke_in_set_argument() {
650 let src = r#"
653 pub fn swap(env: Env) {
654 env.storage().set(key, env.invoke_contract(&pool, sym, args));
655 }
656 "#;
657 let rep = report_for(src);
658 assert!(
660 has(&rep, "CEI-001"),
661 "invoke in set args happens before set's effect"
662 );
663 }
664
665 #[test]
666 fn cei001_detects_invoke_wrapping_set_argument() {
667 let src = r#"
670 pub fn swap(env: Env) {
671 env.invoke_contract(&pool, sym, env.storage().set(key, value));
672 }
673 "#;
674 let rep = report_for(src);
675 assert!(
678 !has(&rep, "CEI-001"),
679 "set in invoke args happens before invoke's effect"
680 );
681 }
682
683 #[test]
684 fn cei001_detects_write_call_write_pattern() {
685 let src = r#"
687 pub fn swap(env: Env) {
688 env.storage().set(key1, value1);
689 env.invoke_contract(&pool, sym!("withdraw"), args);
690 env.storage().set(key2, value2);
691 }
692 "#;
693 let rep = report_for(src);
694 assert!(
698 has(&rep, "CEI-001"),
699 "write → call → write should flag on second write"
700 );
701 }
702
703 #[test]
704 fn cei001_still_emits_one_finding_per_function() {
705 let src = r#"
707 pub fn swap(env: Env) {
708 env.storage().set(key1, value1);
709 env.invoke_contract(&pool, sym!("withdraw"), args);
710 env.storage().set(key2, value2);
711 env.invoke_contract(&pool2, sym!("transfer"), args);
712 env.storage().set(key3, value3);
713 }
714 "#;
715 let rep = report_for(src);
716 let cei_findings: Vec<_> = rep
717 .findings
718 .iter()
719 .filter(|f| f.rule_id == "CEI-001")
720 .collect();
721 assert_eq!(
722 cei_findings.len(),
723 1,
724 "should emit only one CEI-001 finding per function"
725 );
726 }
727
728 #[test]
729 fn cei001_silent_on_state_write_before_external_call() {
730 let src = r#"
731 pub fn swap_safe(env: Env, to: Address) {
732 env.storage().set(key, value);
733 env.invoke_contract(&pool, sym!("withdraw"), args);
734 }
735 "#;
736 let rep = report_for(src);
737 assert!(!has(&rep, "CEI-001"));
738 }
739
740 #[test]
741 fn cei001_silent_on_only_external_call() {
742 let src = r#"
743 pub fn relay(env: Env) {
744 env.invoke_contract(&addr, sym, args);
745 }
746 "#;
747 let rep = report_for(src);
748 assert!(!has(&rep, "CEI-001"));
749 }
750
751 #[test]
752 fn cei001_silent_on_only_state_write() {
753 let src = r#"
754 pub fn set_state(env: Env) {
755 env.storage().set(key, value);
756 }
757 "#;
758 let rep = report_for(src);
759 assert!(!has(&rep, "CEI-001"));
760 }
761
762 #[test]
763 fn cei001_can_be_disabled() {
764 let src = r#"
765 pub fn swap(env: Env) {
766 env.invoke_contract(&addr, sym, args);
767 env.storage().set(key, value);
768 }
769 "#;
770 assert!(
771 !has(&audit_source_with(src, &["CEI-001"]).unwrap(), "CEI-001"),
772 "CEI-001 should be suppressible"
773 );
774 }
775
776 #[test]
777 fn clean_source_reports_no_findings() {
778 let src = "pub fn hello(name: String) -> String { require_auth(); hello_helper(name) }";
779 let rep = report_for(src);
780 assert!(rep.is_clean());
781 assert_eq!(rep.summary.total, 0);
782 }
783
784 #[test]
785 fn parse_error_is_reported() {
786 let res = audit_source("this is not valid rust @@@");
787 assert!(matches!(res, Err(AuditError::Parse(_))));
788 }
789
790 #[test]
792 fn rules_implement_trait() {
793 fn assert_rule(_r: &dyn AuditRule) {}
794 assert_rule(&Auth001);
795 assert_rule(&Auth002);
796 assert_rule(&Auth003);
797 assert_rule(&Move001);
798 }
799}