Skip to main content

forge_lint/sol/med/
weak_prng.rs

1use super::WeakPrng;
2use crate::{
3    linter::{EarlyLintPass, LintContext},
4    sol::{Severity, SolLint},
5};
6use solar::{
7    ast::{BinOp, BinOpKind, Expr, ExprKind, IndexKind, LitKind, SourceUnit, visit::Visit},
8    interface::SpannedOption,
9};
10use std::ops::ControlFlow;
11
12declare_forge_lint!(
13    WEAK_PRNG,
14    Severity::Med,
15    "weak-prng",
16    "weak randomness derived from a predictable on-chain value"
17);
18
19impl<'ast> EarlyLintPass<'ast> for WeakPrng {
20    fn check_full_source_unit(&mut self, ctx: &LintContext<'ast, '_>, ast: &'ast SourceUnit<'ast>) {
21        if ctx.is_lint_enabled(WEAK_PRNG.id) {
22            let mut checker = WeakPrngChecker { ctx };
23            let _ = checker.visit_source_unit(ast);
24        }
25    }
26}
27
28struct WeakPrngChecker<'a, 's> {
29    ctx: &'a LintContext<'s, 'a>,
30}
31
32impl<'ast> Visit<'ast> for WeakPrngChecker<'_, '_> {
33    type BreakValue = ();
34
35    fn visit_expr(&mut self, expr: &'ast Expr<'ast>) -> ControlFlow<Self::BreakValue> {
36        if is_randomness_expr(expr) {
37            self.ctx.emit(&WEAK_PRNG, expr.span);
38            ControlFlow::Continue(())
39        } else {
40            self.walk_expr(expr)
41        }
42    }
43}
44
45fn is_randomness_expr(expr: &Expr<'_>) -> bool {
46    match &expr.peel_parens().kind {
47        ExprKind::Binary(lhs, BinOp { kind: BinOpKind::Rem, .. }, rhs) => {
48            is_randomness_modulo(lhs, rhs)
49        }
50        ExprKind::Call(callee, args) => {
51            let callee = callee.peel_parens();
52            is_keccak256(callee) && args.exprs().any(contains_predictable_source)
53        }
54        _ => false,
55    }
56}
57
58fn is_randomness_modulo(lhs: &Expr<'_>, rhs: &Expr<'_>) -> bool {
59    if is_timestamp_time_bucket(lhs, rhs) {
60        return false;
61    }
62    contains_predictable_source(lhs) || contains_predictable_source(rhs)
63}
64
65fn is_timestamp_time_bucket(lhs: &Expr<'_>, rhs: &Expr<'_>) -> bool {
66    is_block_timestamp(lhs.peel_parens()) && is_time_bucket_modulus(rhs)
67}
68
69fn is_timestamp_time_bucket_expr(expr: &Expr<'_>) -> bool {
70    matches!(
71        &expr.peel_parens().kind,
72        ExprKind::Binary(lhs, BinOp { kind: BinOpKind::Rem, .. }, rhs)
73            if is_timestamp_time_bucket(lhs, rhs)
74    )
75}
76
77fn is_time_bucket_modulus(expr: &Expr<'_>) -> bool {
78    const SECONDS_PER_DAY: u64 = 24 * 60 * 60;
79
80    const_eval_u64(expr)
81        .is_some_and(|value| value >= SECONDS_PER_DAY && value % SECONDS_PER_DAY == 0)
82}
83
84fn const_eval_u64(expr: &Expr<'_>) -> Option<u64> {
85    match &expr.peel_parens().kind {
86        ExprKind::Lit(lit, sub) => {
87            if let LitKind::Number(value) = &lit.kind {
88                let base = u64::try_from(value).ok()?;
89                let multiplier = sub.map(|s| s.value()).unwrap_or(1);
90                base.checked_mul(multiplier)
91            } else {
92                None
93            }
94        }
95        ExprKind::Binary(lhs, BinOp { kind, .. }, rhs) => {
96            let lhs = const_eval_u64(lhs)?;
97            let rhs = const_eval_u64(rhs)?;
98            match kind {
99                BinOpKind::Add => lhs.checked_add(rhs),
100                BinOpKind::Sub => lhs.checked_sub(rhs),
101                BinOpKind::Mul => lhs.checked_mul(rhs),
102                BinOpKind::Div => lhs.checked_div(rhs),
103                _ => None,
104            }
105        }
106        _ => None,
107    }
108}
109
110fn contains_predictable_source(expr: &Expr<'_>) -> bool {
111    let expr = expr.peel_parens();
112    if is_timestamp_time_bucket_expr(expr) {
113        return false;
114    }
115    if is_predictable_source(expr) {
116        return true;
117    }
118
119    match &expr.kind {
120        ExprKind::Array(elems) => elems.iter().any(|elem| contains_predictable_source(elem)),
121        ExprKind::Assign(lhs, _, rhs) | ExprKind::Binary(lhs, _, rhs) => {
122            contains_predictable_source(lhs) || contains_predictable_source(rhs)
123        }
124        ExprKind::Call(callee, args) if is_abi_encode(callee.peel_parens()) => {
125            args.exprs().any(contains_predictable_source)
126        }
127        ExprKind::Call(callee, args) => {
128            contains_predictable_source(callee) || args.exprs().any(contains_predictable_source)
129        }
130        ExprKind::CallOptions(callee, options) => {
131            contains_predictable_source(callee)
132                || options.iter().any(|option| contains_predictable_source(option.value))
133        }
134        ExprKind::Member(inner, _) | ExprKind::Unary(_, inner) => {
135            contains_predictable_source(inner)
136        }
137        ExprKind::Index(base, index) => {
138            contains_predictable_source(base)
139                || match index {
140                    IndexKind::Index(Some(index)) => contains_predictable_source(index),
141                    IndexKind::Range(start, end) => {
142                        start.as_ref().is_some_and(|start| contains_predictable_source(start))
143                            || end.as_ref().is_some_and(|end| contains_predictable_source(end))
144                    }
145                    _ => false,
146                }
147        }
148        ExprKind::Payable(args) => args.exprs().any(contains_predictable_source),
149        ExprKind::Ternary(cond, then_expr, else_expr) => {
150            contains_predictable_source(cond)
151                || contains_predictable_source(then_expr)
152                || contains_predictable_source(else_expr)
153        }
154        ExprKind::Tuple(elems) => elems.iter().any(|elem| {
155            if let SpannedOption::Some(inner) = elem.as_ref() {
156                contains_predictable_source(inner)
157            } else {
158                false
159            }
160        }),
161        _ => false,
162    }
163}
164
165fn is_predictable_source(expr: &Expr<'_>) -> bool {
166    is_block_member(expr) || is_blockhash_call(expr)
167}
168
169fn is_block_member(expr: &Expr<'_>) -> bool {
170    is_block_member_with(expr, |member| {
171        matches!(member, "timestamp" | "number" | "coinbase" | "prevrandao" | "difficulty")
172    })
173}
174
175fn is_block_timestamp(expr: &Expr<'_>) -> bool {
176    is_block_member_with(expr, |member| member == "timestamp")
177}
178
179fn is_block_member_with(expr: &Expr<'_>, predicate: impl FnOnce(&str) -> bool) -> bool {
180    matches!(
181        &expr.kind,
182        ExprKind::Member(base, member)
183            if predicate(member.as_str())
184            && is_ident(base.peel_parens(), "block")
185    )
186}
187
188fn is_blockhash_call(expr: &Expr<'_>) -> bool {
189    matches!(
190        &expr.kind,
191        ExprKind::Call(callee, _) if is_ident(callee.peel_parens(), "blockhash")
192    )
193}
194
195fn is_keccak256(expr: &Expr<'_>) -> bool {
196    is_ident(expr, "keccak256")
197}
198
199fn is_abi_encode(expr: &Expr<'_>) -> bool {
200    matches!(
201        &expr.kind,
202        ExprKind::Member(base, member)
203            if matches!(
204                member.as_str(),
205                "encode"
206                    | "encodePacked"
207                    | "encodeWithSelector"
208                    | "encodeWithSignature"
209                    | "encodeCall"
210            ) && is_ident(base.peel_parens(), "abi")
211    )
212}
213
214fn is_ident(expr: &Expr<'_>, name: &str) -> bool {
215    matches!(&expr.kind, ExprKind::Ident(ident) if ident.as_str() == name)
216}