implement lint6v2 to better handle more than 2 identical accounts

Author: oslfmt

lints/dup_mutable_accounts_2/src/lib.rs

@@ -14,7 +14,7 @@ use if_chain::if_chain;
 use rustc_hir::{
     def_id::DefId,
     intravisit::{walk_expr, FnKind, Visitor},
-    Body, Expr, ExprKind, FnDecl, HirId, Mutability,
+    BinOpKind, Body, Expr, ExprKind, FnDecl, HirId, Mutability,
 };
 use rustc_lint::{LateContext, LateLintPass};
 use rustc_middle::ty::TyKind;
@@ -57,71 +57,80 @@ impl<'tcx> LateLintPass<'tcx> for DupMutableAccounts2 {
         _: HirId,
     ) {
         if !span.from_expansion() {
-            let accounts = get_referenced_accounts(cx, body);
-
-            accounts.values().for_each(|exprs| {
-                // TODO: figure out handling of >2 accounts
-                match exprs.len() {
-                    2 => {
-                        let first = exprs[0];
-                        let second = exprs[1];
-                        if !contains_key_call(cx, body, first) {
-                            span_lint_and_help(
-                                cx,
-                                DUP_MUTABLE_ACCOUNTS_2,
-                                first.span,
-                                "this expression does not have a key check but has the same account type as another expression",
-                                Some(second.span),
-                                "add a key check to make sure the accounts have different keys, e.g., x.key() != y.key()",
-                            );
-                        }
-                        if !contains_key_call(cx, body, second) {
-                            span_lint_and_help(
-                                cx,
-                                DUP_MUTABLE_ACCOUNTS_2,
-                                second.span,
-                                "this expression does not have a key check but has the same account type as another expression",
-                                Some(first.span),
-                                "add a key check to make sure the accounts have different keys, e.g., x.key() != y.key()",
-                            );
+            // get all mutable references to Accounts and if_statements in body
+            let mut values = Values::new(cx);
+            values.get_referenced_accounts_and_if_statements(cx, body);
+            // println!("{:#?}", values.if_statements);
+
+            values.accounts.values().for_each(|exprs| {
+                if exprs.len() > 1 {
+                    for current in 0..exprs.len() - 1 {
+                        for next in current + 1..exprs.len() {
+                            if !values.check_key_constraint(exprs[current], exprs[next]) {
+                                span_lint_and_help(
+                                    cx,
+                                    DUP_MUTABLE_ACCOUNTS_2,
+                                    exprs[current].span,
+                                    "the following expressions have equivalent Account types, yet do not contain a proper key check.",
+                                    Some(exprs[next].span),
+                                    "add a key check to make sure the accounts have different keys, e.g., x.key() != y.key()",
+                                );
+                            }
                         }
-                    },
-                    n if n > 2 => {
-                        span_lint_and_note(
-                            cx,
-                            DUP_MUTABLE_ACCOUNTS_2,
-                            exprs[0].span,
-                            &format!("the following expression has the same account type as {} other accounts", exprs.len()),
-                            None,
-                            "might not check that each account has a unique key"
-                        )
-                    },
-                    _ => {}
+                    }
                 }
             });
         }
     }
 }
 
-struct AccountUses<'cx, 'tcx> {
+struct Values<'cx, 'tcx> {
     cx: &'cx LateContext<'tcx>,
-    uses: HashMap<DefId, Vec<&'tcx Expr<'tcx>>>,
+    accounts: HashMap<DefId, Vec<&'tcx Expr<'tcx>>>,
+    if_statements: Vec<(&'tcx Expr<'tcx>, &'tcx Expr<'tcx>)>,
 }
 
-fn get_referenced_accounts<'tcx>(
-    cx: &LateContext<'tcx>,
-    body: &'tcx Body<'tcx>,
-) -> HashMap<DefId, Vec<&'tcx Expr<'tcx>>> {
-    let mut accounts = AccountUses {
-        cx,
-        uses: HashMap::new(),
-    };
-
-    accounts.visit_expr(&body.value);
-    accounts.uses
+impl<'cx, 'tcx> Values<'cx, 'tcx> {
+    fn new(cx: &'cx LateContext<'tcx>) -> Self {
+        Values {
+            cx,
+            accounts: HashMap::new(),
+            if_statements: Vec::new(),
+        }
+    }
+
+    fn get_referenced_accounts_and_if_statements(
+        &mut self,
+        cx: &'cx LateContext<'tcx>,
+        body: &'tcx Body<'tcx>,
+    ) -> &Self {
+        self.visit_expr(&body.value);
+        self
+    }
+
+    /// Checks if there is a valid key constraint for `first_account` and `second_account`.
+    /// NOTE: currently only considers `first.key() == second.key()` or the symmetric relation as valid constraints.
+    /// TODO: if == relation used, should return some error in the THEN block
+    fn check_key_constraint(&self, first_account: &Expr<'_>, second_account: &Expr<'_>) -> bool {
+        for (left, right) in &self.if_statements {
+            if_chain! {
+                if let ExprKind::MethodCall(path_seg_left, exprs_left, _span) = left.kind;
+                if let ExprKind::MethodCall(path_seg_right, exprs_right, _span) = right.kind;
+                if path_seg_left.ident.name.as_str() == "key" && path_seg_right.ident.name.as_str() == "key";
+                if !exprs_left.is_empty() && !exprs_right.is_empty();
+                let mut spanless_eq = SpanlessEq::new(self.cx);
+                if (spanless_eq.eq_expr(&exprs_left[0], first_account) && spanless_eq.eq_expr(&exprs_right[0], second_account)) 
+                || (spanless_eq.eq_expr(&exprs_left[0], second_account) && spanless_eq.eq_expr(&exprs_right[0], first_account));
+                then {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
 }
 
-impl<'cx, 'tcx> Visitor<'tcx> for AccountUses<'cx, 'tcx> {
+impl<'cx, 'tcx> Visitor<'tcx> for Values<'cx, 'tcx> {
     fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
         if_chain! {
             // get mutable reference expressions
@@ -137,44 +146,58 @@ impl<'cx, 'tcx> Visitor<'tcx> for AccountUses<'cx, 'tcx> {
             if let Some(adt_def) = account_type.ty_adt_def();
             then {
                 let def_id = adt_def.did();
-                if let Some(exprs) = self.uses.get_mut(&def_id) {
+                if let Some(exprs) = self.accounts.get_mut(&def_id) {
                     let mut spanless_eq = SpanlessEq::new(self.cx);
                     // check that expr is not a duplicate within its particular key-pair
                     if exprs.iter().all(|e| !spanless_eq.eq_expr(e, mut_expr)) {
                         exprs.push(mut_expr);
                     }
                 } else {
-                    self.uses.insert(def_id, vec![mut_expr]);
+                    self.accounts.insert(def_id, vec![mut_expr]);
                 }
             }
         }
-        walk_expr(self, expr);
-    }
-}
 
-/// Performs a walk on `body`, checking whether there exists an expression that contains
-/// a `key()` method call on `account_expr`.
-fn contains_key_call<'tcx>(
-    cx: &LateContext<'tcx>,
-    body: &'tcx Body<'tcx>,
-    account_expr: &Expr<'tcx>,
-) -> bool {
-    visit_expr_no_bodies(&body.value, |expr| {
+        // get if statements
         if_chain! {
-            if let ExprKind::MethodCall(path_seg, exprs, _span) = expr.kind;
-            if path_seg.ident.name.as_str() == "key";
-            if !exprs.is_empty();
-            let mut spanless_eq = SpanlessEq::new(cx);
-            if spanless_eq.eq_expr(&exprs[0], account_expr);
+            if let ExprKind::If(wrapped_if_expr, then, _else_opt) = expr.kind;
+            if let ExprKind::DropTemps(if_expr) = wrapped_if_expr.kind;
+            if let ExprKind::Binary(op, left, right) = if_expr.kind;
+            // TODO: leaves out || or &&. Could implement something that pulls apart
+            // an if expr that is of this form into individual == or != comparisons
+            if let BinOpKind::Ne | BinOpKind::Eq = op.node;
             then {
-                true
-            } else {
-                false
+                // println!("{:#?}, {:#?}", expr, then);
+                self.if_statements.push((left, right));
             }
         }
-    })
+        walk_expr(self, expr);
+    }
 }
 
+// /// Performs a walk on `body`, checking whether there exists an expression that contains
+// /// a `key()` method call on `account_expr`.
+// fn contains_key_call<'tcx>(
+//     cx: &LateContext<'tcx>,
+//     body: &'tcx Body<'tcx>,
+//     account_expr: &Expr<'tcx>,
+// ) -> bool {
+//     visit_expr_no_bodies(&body.value, |expr| {
+//         if_chain! {
+//             if let ExprKind::MethodCall(path_seg, exprs, _span) = expr.kind;
+//             if path_seg.ident.name.as_str() == "key";
+//             if !exprs.is_empty();
+//             let mut spanless_eq = SpanlessEq::new(cx);
+//             if spanless_eq.eq_expr(&exprs[0], account_expr);
+//             then {
+//                 true
+//             } else {
+//                 false
+//             }
+//         }
+//     })
+// }
+
 #[test]
 fn insecure() {
     dylint_testing::ui_test_example(env!("CARGO_PKG_NAME"), "insecure");

lints/dup_mutable_accounts_2/ui/secure/src/lib.rs

@@ -14,6 +14,11 @@ pub mod duplicate_mutable_accounts_secure {
         if ctx.accounts.user_a.key() == ctx.accounts.user_b.key() {
             return Err(ProgramError::InvalidArgument);
         }
+
+        // if ctx.accounts.user_a.key() != ctx.accounts.user_b.key() {
+        //     // do program stuff
+        // }
+
         let user_a = &mut ctx.accounts.user_a;
         let user_b = &mut ctx.accounts.user_b;