[-Wunsafe-buffer-usage] Check count-attributed objects that are used and assigned in the same group

Check if the bounds-attributed group has assignments to objects that are
also used in the same group. In those cases, the correctness of the
group might depend on the order of assignments. We conservatively
disallow such assignments.

In the example below, the bounds-check in `sp.first()` uses the value of
`b` before the later update, which can lead to OOB if `b` was less than
42.

```
  void foo(int *__counted_by(a + b) p, int a, int b, std::span<int> sp) {
    p = sp.first(b + 42).data();
    b = 42; // b is assigned and used
    a = b;
  }
```

rdar://161608319

Author: patrykstefanski

clang/include/clang/Analysis/Analyses/UnsafeBufferUsage.h

@@ -181,6 +181,12 @@ class UnsafeBufferUsageHandler {
                                           ASTContext &Ctx) {
     handleUnsafeOperation(Assign, IsRelatedToDecl, Ctx);
   }
+
+  virtual void handleAssignedAndUsed(const BinaryOperator *Assign,
+                                     const Expr *Use, const ValueDecl *VD,
+                                     bool IsRelatedToDecl, ASTContext &Ctx) {
+    handleUnsafeOperation(Assign, IsRelatedToDecl, Ctx);
+  }
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   /// Invoked when a fix is suggested against a variable. This function groups

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -14291,6 +14291,9 @@ def warn_missing_assignments_in_bounds_attributed_group : Warning<
 def warn_duplicated_assignment_in_bounds_attributed_group : Warning<
   "duplicated assignment to %select{variable|parameter|member}0 %1 in "
   "bounds-attributed group">, InGroup<UnsafeBufferUsage>, DefaultIgnore;
+def warn_assigned_and_used_in_bounds_attributed_group : Warning<
+  "%select{variable|parameter|member}0 %1 is assigned and used in the same "
+  "bounds-attributed group">, InGroup<UnsafeBufferUsage>, DefaultIgnore;
 #ifndef NDEBUG
 // Not a user-facing diagnostic. Useful for debugging false negatives in
 // -fsafe-buffer-usage-suggestions (i.e. lack of -Wunsafe-buffer-usage fixits).

clang/lib/Analysis/UnsafeBufferUsage.cpp

@@ -5620,6 +5620,16 @@ struct BoundsAttributedObject {
   const ValueDecl *Decl = nullptr;
   const Expr *MemberBase = nullptr;
   int DerefLevel = 0;
+
+  bool operator==(const BoundsAttributedObject &Other) const {
+    if (Other.Decl != Decl || Other.DerefLevel != DerefLevel)
+      return false;
+    if (Other.MemberBase == MemberBase)
+      return true;
+    if (Other.MemberBase == nullptr || MemberBase == nullptr)
+      return false;
+    return isSameMemberBase(Other.MemberBase, MemberBase);
+  }
 };
 
 static std::optional<BoundsAttributedObject>
@@ -5663,7 +5673,7 @@ struct BoundsAttributedAssignmentGroup {
   DependentDeclSetTy DeclClosure;
   llvm::SmallVector<const BinaryOperator *, 4> Assignments;
   llvm::SmallVector<BoundsAttributedObject, 4> AssignedObjects;
-  using DeclUseTy = std::pair<const ValueDecl *, const Expr *>;
+  llvm::SmallDenseMap<const ValueDecl *, const Expr *, 4> Uses;
   const Expr *MemberBase = nullptr;
 
   void init(const BoundsAttributedObject &Object) {
@@ -5675,6 +5685,7 @@ struct BoundsAttributedAssignmentGroup {
     DeclClosure.clear();
     Assignments.clear();
     AssignedObjects.clear();
+    Uses.clear();
     MemberBase = nullptr;
   }
 
@@ -5816,6 +5827,22 @@ struct BoundsAttributedGroupFinder
     if (DA.has_value())
       TooComplexAssignHandler(E, DA->Decl);
   }
+
+  // Collect uses of decls belonging to the current group by visiting all DREs
+  // and MemberExprs.
+
+  void addUseIfPartOfCurGroup(const Expr *E) {
+    const auto DA = getBoundsAttributedObject(E);
+    if (DA.has_value() && CurGroup.isPartOfGroup(*DA))
+      CurGroup.Uses.insert({DA->Decl, E});
+  }
+
+  void VisitDeclRefExpr(const DeclRefExpr *DRE) { addUseIfPartOfCurGroup(DRE); }
+
+  void VisitMemberExpr(const MemberExpr *ME) {
+    addUseIfPartOfCurGroup(ME);
+    Visit(ME->getBase());
+  }
 };
 
 // Checks if the bounds-attributed group does not assign to implicitly
@@ -5947,6 +5974,52 @@ static bool checkMissingAndDuplicatedAssignments(
   return IsGroupSafe;
 }
 
+// Checks if the bounds-attributed group has assignments to objects that are
+// also used in the same group. In those cases, the correctness of the group
+// might depend on the order of assignments. We conservatively disallow such
+// assignments.
+//
+// In the example below, the bounds-check in `sp.first()` uses the value of `b`
+// before the later update, which can lead to OOB if `b` was less than 42.
+//   void foo(int *__counted_by(a + b) p, int a, int b, std::span<int> sp) {
+//     p = sp.first(b + 42).data();
+//     b = 42; // b is assigned and used
+//     a = b;
+//   }
+static bool checkAssignedAndUsed(const BoundsAttributedAssignmentGroup &Group,
+                                 UnsafeBufferUsageHandler &Handler,
+                                 ASTContext &Ctx) {
+  const auto &Uses = Group.Uses;
+  if (Uses.empty())
+    return true;
+
+  bool IsGroupSafe = true;
+
+  for (size_t I = 0, N = Group.AssignedObjects.size(); I < N; ++I) {
+    const BoundsAttributedObject &LHSObj = Group.AssignedObjects[I];
+    const BinaryOperator *Assign = Group.Assignments[I];
+
+    // Ignore self assignments, because they don't matter, since the value stays
+    // the same.
+    const auto RHSObj = getBoundsAttributedObject(Assign->getRHS());
+    bool IsSelfAssign = RHSObj.has_value() && *RHSObj == LHSObj;
+    if (IsSelfAssign)
+      continue;
+
+    const ValueDecl *VD = LHSObj.Decl;
+    const auto It = Uses.find(VD);
+    if (It == Uses.end())
+      continue;
+
+    const Expr *Use = It->second;
+    Handler.handleAssignedAndUsed(Assign, Use, VD,
+                                  /*IsRelatedToDecl=*/false, Ctx);
+    IsGroupSafe = false;
+  }
+
+  return IsGroupSafe;
+}
+
 // Checks if the bounds-attributed group is safe. This function returns false
 // iff the assignment group is unsafe and diagnostics have been emitted.
 static bool
@@ -5956,6 +6029,8 @@ checkBoundsAttributedGroup(const BoundsAttributedAssignmentGroup &Group,
     return false;
   if (!checkMissingAndDuplicatedAssignments(Group, Handler, Ctx))
     return false;
+  if (!checkAssignedAndUsed(Group, Handler, Ctx))
+    return false;
   // TODO: Add more checks.
   return true;
 }

clang/lib/Sema/AnalysisBasedWarnings.cpp

@@ -2685,6 +2685,15 @@ class UnsafeBufferUsageReporter : public UnsafeBufferUsageHandler {
     S.Diag(PrevAssign->getOperatorLoc(),
            diag::note_bounds_safety_previous_assignment);
   }
+
+  void handleAssignedAndUsed(const BinaryOperator *Assign, const Expr *Use,
+                             const ValueDecl *VD, bool IsRelatedToDecl,
+                             ASTContext &Ctx) override {
+    S.Diag(Assign->getOperatorLoc(),
+           diag::warn_assigned_and_used_in_bounds_attributed_group)
+        << getBoundsAttributedObjectKind(VD) << VD;
+    S.Diag(Use->getBeginLoc(), diag::note_used_here);
+  }
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   void handleUnsafeVariableGroup(const VarDecl *Variable,

clang/test/SemaCXX/warn-unsafe-buffer-usage-count-attributed-pointer-assignment.cpp

@@ -339,6 +339,24 @@ void duplicated_complex(int *__counted_by(a + b) p,
   q = nullptr; // expected-warning{{duplicated assignment to parameter 'q' in bounds-attributed group}}
 }
 
+// Assigned and used
+
+void good_assigned_and_used(int *__counted_by(count) p, int count, std::span<int> sp) {
+  p = sp.first(count).data();
+  count = count;
+}
+
+void bad_assigned_and_used(int *__counted_by(count) p, int count, std::span<int> sp, int new_count) {
+  p = sp.first(count).data(); // expected-note{{used here}}
+  count = new_count;          // expected-warning{{parameter 'count' is assigned and used in the same bounds-attributed group}}
+}
+
+void bad_assigned_and_used2(int *__counted_by(a + b) p, int a, int b, std::span<int> sp) {
+  p = sp.first(b + 42).data(); // expected-note{{used here}}
+  b = 42;                      // expected-warning{{parameter 'b' is assigned and used in the same bounds-attributed group}}
+  a = b;
+}
+
 // Assigns to bounds-attributed that we consider too complex to analyze.
 
 void too_complex_assign_to_ptr(int *__counted_by(count) p, int count, int *q) {