feat(lint): handle comparison of references in functions

linting/src/strict_balance_equality.rs

@@ -13,8 +13,8 @@
 // limitations under the License.
 use clippy_utils::{
     diagnostics::span_lint_hir_and_then,
+    match_any_def_paths,
     match_def_path,
-    source::snippet_opt,
 };
 use if_chain::if_chain;
 use rustc_errors::Applicability;
@@ -65,10 +65,7 @@ use rustc_session::{
     declare_lint,
     declare_lint_pass,
 };
-use rustc_span::{
-    source_map::BytePos,
-    Span,
-};
+use rustc_span::Span;
 use std::collections::{
     HashMap,
     HashSet,
@@ -405,10 +402,36 @@ impl<'tcx> TransferFunction<'_, 'tcx> {
             acc
         });
 
+        let fn_def_id =
+            if let mir_ty::TyKind::FnDef(fn_def_id, _) = func.literal.ty().kind() {
+                fn_def_id
+            } else {
+                return
+            };
+
+        // Handle `PartialEq` functions that implement comparsion for non-primitive types,
+        // including references like `&i32`.
+        if_chain! {
+            if init_taints.len() == 2;
+            if init_taints.iter().any(|&tainted| tainted);
+            if match_any_def_paths(
+                self.cx,
+                *fn_def_id,
+                &[
+                    &["core", "cmp", "PartialEq", "ne"],
+                    &["core", "cmp", "PartialEq", "eq"],
+                ],
+            )
+            .is_some();
+            then {
+                self.state.insert(destination.local);
+                return
+            }
+        }
+
         let fn_mir = if_chain! {
-            if let mir_ty::TyKind::FnDef(id, _) = func.literal.ty().kind();
-            if self.fn_is_defined_in_user_code(id);
-            then { self.cx.tcx.optimized_mir(id) } else { return }
+            if self.fn_is_defined_in_user_code(fn_def_id);
+            then { self.cx.tcx.optimized_mir(fn_def_id) } else { return }
         };
 
         // Run the dataflow analysis if the function hasn't been analyzed yet
@@ -618,17 +641,15 @@ impl<'tcx> StrictBalanceEquality {
                     .as_ref()
                     .assert_crate_local()
                     .lint_root;
-                if let Some(snip) = snippet_opt(cx, span);
-                if let Some(op) = snip.rfind("==").or(snip.rfind("!="));
                 then {
-                    let op_pos = span.lo() + BytePos(op as u32);
                     let sugg_span = Span::new(
-                        op_pos,
-                        op_pos + BytePos("==".len() as u32),
-                        // We have to use a span different from `span`, since it is resulted after
-                        // macro expansion and therefore cannot be used to emit diagnostics.
-                        fn_span.ctxt(),
-                        fn_span.parent()
+                            span.lo(),
+                            span.hi(),
+                            // We have to use a span different from `span`, since it is resulted
+                            // after macro expansion and therefore cannot be used to emit
+                            // diagnostics.
+                            fn_span.ctxt(),
+                            fn_span.parent()
                     );
                     span_lint_hir_and_then(
                         cx,
@@ -645,7 +666,6 @@ impl<'tcx> StrictBalanceEquality {
                             );
                         },
                     )
-
                 }
             }
         }

linting/ui/fail/strict_balance_equality.rs

@@ -9,6 +9,7 @@ pub mod strict_balance_equality {
             Self {}
         }
 
+        // Return value tainted with balance
         fn get_balance_1(&self) -> Balance {
             self.env().balance()
         }
@@ -27,13 +28,19 @@ pub mod strict_balance_equality {
                 self.env().balance()
             }
         }
+
+        // Return the result of comparison with balance
         fn cmp_balance_1(&self, value: &Balance) -> bool {
             *value == self.env().balance()
         }
         fn cmp_balance_2(&self, value: &Balance, threshold: &Balance) -> bool {
-            value != threshold
+             value != threshold
+        }
+        fn cmp_balance_3(&self, value: Balance, threshold: Balance) -> bool {
+             value != threshold
         }
 
+        // Tainted `&mut` input argument
         fn get_balance_arg_1(&self, value: &mut Balance) {
             *value = self.env().balance();
         }
@@ -56,10 +63,14 @@ pub mod strict_balance_equality {
             if self.get_balance_2() == 10 { /* ... */ }
             if self.get_balance_3() == 10 { /* ... */ }
             if self.get_balance_recursive(&10) == 10 { /* ... */ }
-            // if self.cmp_balance_1(&10) { /* ... */ } // TODO: false negative
-            // if self.cmp_balance_2(&self.env().balance(), &threshold) { /* ... */ } // TODO: false negative
 
-            // Bad: Strict equality in function: tainted arguments
+            // Bad: Strict equality in function call: return value contains the result of
+            // comparison
+            if self.cmp_balance_1(&10) { /* ... */ }
+            if self.cmp_balance_2(&self.env().balance(), &threshold) { /* ... */ }
+            if self.cmp_balance_3(self.env().balance(), threshold) { /* ... */ }
+
+            // // Bad: Strict equality in function: tainted arguments
             let mut res_1 = 0_u128;
             self.get_balance_arg_1(&mut res_1);
             if res_1 == 10 { /* ... */ }

linting/ui/fail/strict_balance_equality.stderr

@@ -1,58 +1,76 @@
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:50:37
+  --> $DIR/strict_balance_equality.rs:57:16
    |
 LL |             if self.env().balance() == 10 { /* ... */ }
-   |                                     ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
    |
    = note: `-D strict-balance-equality` implied by `-D warnings`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:51:22
+  --> $DIR/strict_balance_equality.rs:58:16
    |
 LL |             if value == 11 { /* ... */ }
-   |                      ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:52:37
+  --> $DIR/strict_balance_equality.rs:59:16
    |
 LL |             if self.env().balance() == threshold { /* ... */ }
-   |                                     ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:55:37
+  --> $DIR/strict_balance_equality.rs:62:16
    |
 LL |             if self.get_balance_1() == 10 { /* ... */ }
-   |                                     ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:56:37
+  --> $DIR/strict_balance_equality.rs:63:16
    |
 LL |             if self.get_balance_2() == 10 { /* ... */ }
-   |                                     ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:57:37
+  --> $DIR/strict_balance_equality.rs:64:16
    |
 LL |             if self.get_balance_3() == 10 { /* ... */ }
-   |                                     ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:58:48
+  --> $DIR/strict_balance_equality.rs:65:16
    |
 LL |             if self.get_balance_recursive(&10) == 10 { /* ... */ }
-   |                                                ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:65:22
+  --> $DIR/strict_balance_equality.rs:69:16
+   |
+LL |             if self.cmp_balance_1(&10) { /* ... */ }
+   |                ^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
+
+error: dangerous strict balance equality
+  --> $DIR/strict_balance_equality.rs:70:16
+   |
+LL |             if self.cmp_balance_2(&self.env().balance(), &threshold) { /* ... */ }
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
+
+error: dangerous strict balance equality
+  --> $DIR/strict_balance_equality.rs:71:16
+   |
+LL |             if self.cmp_balance_3(self.env().balance(), threshold) { /* ... */ }
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
+
+error: dangerous strict balance equality
+  --> $DIR/strict_balance_equality.rs:76:16
    |
 LL |             if res_1 == 10 { /* ... */ }
-   |                      ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
 error: dangerous strict balance equality
-  --> $DIR/strict_balance_equality.rs:68:22
+  --> $DIR/strict_balance_equality.rs:79:16
    |
 LL |             if res_2 == 10 { /* ... */ }
-   |                      ^^ help: consider using non-strict equality operators instead: `<`, `>`
+   |                ^^^^^^^^^^^ help: consider using non-strict equality operators instead: `<`, `>`
 
-error: aborting due to 9 previous errors
+error: aborting due to 12 previous errors