[-Wunsafe-buffer-usage] Check safe assignment patterns

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

@@ -181,6 +181,11 @@ class UnsafeBufferUsageHandler {
                                           ASTContext &Ctx) {
     handleUnsafeOperation(Assign, IsRelatedToDecl, Ctx);
   }
+
+  virtual void handleUnsafeCountAttributedPointerAssignment(
+      const BinaryOperator *Assign, 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

@@ -14289,6 +14289,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_unsafe_count_attributed_pointer_assignment : Warning<
+  "unsafe assignment to count-attributed pointer">,
+  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

@@ -5947,6 +5947,60 @@ static bool checkMissingAndDuplicatedAssignments(
   return IsGroupSafe;
 }
 
+// Checks if each assignment to count-attributed pointer in the group is safe.
+static bool
+checkAssignmentPatterns(const BoundsAttributedAssignmentGroup &Group,
+                        UnsafeBufferUsageHandler &Handler, ASTContext &Ctx) {
+  // Collect dependent values.
+  DependentValuesTy DependentValues;
+  for (size_t I = 0, N = Group.AssignedObjects.size(); I < N; ++I) {
+    const ValueDecl *VD = Group.AssignedObjects[I].Decl;
+    const auto *Attr = VD->getAttr<DependerDeclsAttr>();
+    if (!Attr)
+      continue;
+
+    const BinaryOperator *Assign = Group.Assignments[I];
+    const Expr *Value = Assign->getRHS();
+
+    [[maybe_unused]] bool Inserted =
+        DependentValues.insert({{VD, Attr->getIsDeref()}, Value}).second;
+    // Previous checks in `checkBoundsAttributedGroup` should have validated
+    // that we have only a single assignment.
+    assert(Inserted);
+  }
+
+  bool IsGroupSafe = true;
+
+  // Check every pointer in the group.
+  for (size_t I = 0, N = Group.AssignedObjects.size(); I < N; ++I) {
+    const ValueDecl *VD = Group.AssignedObjects[I].Decl;
+
+    QualType Ty = VD->getType();
+    const auto *CAT = Ty->getAs<CountAttributedType>();
+    if (!CAT && Ty->isPointerType())
+      CAT = Ty->getPointeeType()->getAs<CountAttributedType>();
+    if (!CAT)
+      continue;
+
+    const BinaryOperator *Assign = Group.Assignments[I];
+
+    // TODO: Move this logic to isCountAttributedPointerArgumentSafeImpl.
+    const Expr *CountArg =
+        DependentValues.size() == 1 ? DependentValues.begin()->second : nullptr;
+
+    bool IsSafe = isCountAttributedPointerArgumentSafeImpl(
+        Ctx, Assign->getRHS(), CountArg, CAT, CAT->getCountExpr(),
+        CAT->isCountInBytes(), CAT->isOrNull(), &DependentValues);
+    if (!IsSafe) {
+      Handler.handleUnsafeCountAttributedPointerAssignment(
+          Assign, /*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,8 +6010,7 @@ checkBoundsAttributedGroup(const BoundsAttributedAssignmentGroup &Group,
     return false;
   if (!checkMissingAndDuplicatedAssignments(Group, Handler, Ctx))
     return false;
-  // TODO: Add more checks.
-  return true;
+  return checkAssignmentPatterns(Group, Handler, Ctx);
 }
 
 static void

clang/lib/Sema/AnalysisBasedWarnings.cpp

@@ -2685,6 +2685,14 @@ class UnsafeBufferUsageReporter : public UnsafeBufferUsageHandler {
     S.Diag(PrevAssign->getOperatorLoc(),
            diag::note_bounds_safety_previous_assignment);
   }
+
+  void
+  handleUnsafeCountAttributedPointerAssignment(const BinaryOperator *Assign,
+                                               bool IsRelatedToDecl,
+                                               ASTContext &Ctx) override {
+    S.Diag(Assign->getOperatorLoc(),
+           diag::warn_unsafe_count_attributed_pointer_assignment);
+  }
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   void handleUnsafeVariableGroup(const VarDecl *Variable,