Skip to main content

forge_lint/sol/med/
weak_prng.rs

1use super::WeakPrng;
2use crate::{
3    linter::{LateLintPass, LintContext},
4    sol::{Severity, SolLint},
5};
6use alloy_primitives::{U256, uint};
7use solar::{
8    ast::{BinOp, BinOpKind},
9    sema::{
10        Gcx,
11        builtins::Builtin,
12        hir::{Expr, ExprKind, Hir, SourceId, Visit},
13    },
14};
15use std::ops::ControlFlow;
16
17declare_forge_lint!(
18    WEAK_PRNG,
19    Severity::Med,
20    "weak-prng",
21    "weak randomness derived from a predictable on-chain value"
22);
23
24impl<'gcx> LateLintPass<'gcx> for WeakPrng {
25    fn check_nested_source(&mut self, ctx: &LintContext, gcx: Gcx<'gcx>, id: SourceId) {
26        if ctx.is_lint_enabled(WEAK_PRNG.id) {
27            let _ = WeakPrngChecker { ctx, gcx }.visit_nested_source(id);
28        }
29    }
30}
31
32struct WeakPrngChecker<'a, 's, 'gcx> {
33    ctx: &'a LintContext<'s, 'a>,
34    gcx: Gcx<'gcx>,
35}
36
37impl<'gcx> Visit<'gcx> for WeakPrngChecker<'_, '_, 'gcx> {
38    type BreakValue = ();
39
40    fn hir(&self) -> &'gcx Hir<'gcx> {
41        &self.gcx.hir
42    }
43
44    /// Emits once per outermost `<..> % <..>` or `keccak256(..)` fed by a predictable source.
45    fn visit_expr(&mut self, expr: &'gcx Expr<'gcx>) -> ControlFlow<()> {
46        let is_randomness = match &expr.peel_parens().kind {
47            ExprKind::Binary(lhs, BinOp { kind: BinOpKind::Rem, .. }, rhs) => {
48                !is_timestamp_time_bucket(self.gcx, lhs, rhs)
49                    && (contains_predictable_source(self.gcx, lhs)
50                        || contains_predictable_source(self.gcx, rhs))
51            }
52            ExprKind::Call(callee, args, _) => {
53                self.gcx.resolved_builtin(callee) == Some(Builtin::Keccak256)
54                    && args.exprs().any(|arg| contains_predictable_source(self.gcx, arg))
55            }
56            _ => false,
57        };
58        if is_randomness {
59            self.ctx.emit(&WEAK_PRNG, expr.span);
60            return ControlFlow::Continue(());
61        }
62        self.walk_expr(expr)
63    }
64}
65
66fn contains_predictable_source<'gcx>(gcx: Gcx<'gcx>, expr: &'gcx Expr<'gcx>) -> bool {
67    PredictableSourceFinder { gcx }.visit_expr(expr).is_break()
68}
69
70struct PredictableSourceFinder<'gcx> {
71    gcx: Gcx<'gcx>,
72}
73
74impl<'gcx> Visit<'gcx> for PredictableSourceFinder<'gcx> {
75    type BreakValue = ();
76
77    fn hir(&self) -> &'gcx Hir<'gcx> {
78        &self.gcx.hir
79    }
80
81    fn visit_expr(&mut self, expr: &'gcx Expr<'gcx>) -> ControlFlow<()> {
82        match &expr.peel_parens().kind {
83            // `block.timestamp % 1 days` is a time bucket, not a random draw.
84            ExprKind::Binary(lhs, BinOp { kind: BinOpKind::Rem, .. }, rhs)
85                if is_timestamp_time_bucket(self.gcx, lhs, rhs) =>
86            {
87                ControlFlow::Continue(())
88            }
89            _ if matches!(
90                self.gcx.resolved_builtin(expr),
91                Some(
92                    Builtin::BlockTimestamp
93                        | Builtin::BlockNumber
94                        | Builtin::BlockCoinbase
95                        | Builtin::BlockPrevrandao
96                        | Builtin::BlockDifficulty
97                )
98            ) =>
99            {
100                ControlFlow::Break(())
101            }
102            ExprKind::Call(callee, ..)
103                if self.gcx.resolved_builtin(callee) == Some(Builtin::Blockhash) =>
104            {
105                ControlFlow::Break(())
106            }
107            _ => self.walk_expr(expr),
108        }
109    }
110}
111
112/// `block.timestamp % <multiple of one day>`.
113fn is_timestamp_time_bucket(gcx: Gcx<'_>, lhs: &Expr<'_>, rhs: &Expr<'_>) -> bool {
114    const SECONDS_PER_DAY: U256 = uint!(86400_U256);
115    gcx.resolved_builtin(lhs) == Some(Builtin::BlockTimestamp)
116        && gcx
117            .try_eval_const(rhs)
118            .ok()
119            .and_then(|v| v.as_u256())
120            .is_some_and(|v| v >= SECONDS_PER_DAY && v % SECONDS_PER_DAY == U256::ZERO)
121}