vendor/llvm-project/llvmorg-15-init-15358-g53dc0f107877

1 files changed, 231 insertions, 0 deletions

diff --git a/clang/lib/AST/Randstruct.cpp b/clang/lib/AST/Randstruct.cpp
new file mode 100644
index 000000000000..99c665f420e6
--- /dev/null
+++ b/clang/lib/AST/Randstruct.cpp
@@ -0,0 +1,231 @@
+//===--- Randstruct.cpp ---------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the implementation for Clang's structure field layout
+// randomization.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/AST/Randstruct.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/ASTDiagnostic.h"
+#include "clang/AST/Attr.h"
+#include "clang/AST/Decl.h"
+#include "clang/AST/DeclCXX.h" // For StaticAssertDecl
+#include "clang/Basic/Diagnostic.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <algorithm>
+#include <random>
+#include <set>
+#include <sstream>
+#include <string>
+
+using clang::ASTContext;
+using clang::FieldDecl;
+using llvm::SmallVector;
+
+namespace {
+
+// FIXME: Replace this with some discovery once that mechanism exists.
+enum { CACHE_LINE = 64 };
+
+// The Bucket class holds the struct fields we're trying to fill to a
+// cache-line.
+class Bucket {
+  SmallVector<FieldDecl *, 64> Fields;
+  int Size = 0;
+
+public:
+  virtual ~Bucket() = default;
+
+  SmallVector<FieldDecl *, 64> &fields() { return Fields; }
+  void addField(FieldDecl *Field, int FieldSize);
+  virtual bool canFit(int FieldSize) const {
+    return Size + FieldSize <= CACHE_LINE;
+  }
+  virtual bool isBitfieldRun() const { return false; }
+  bool full() const { return Size >= CACHE_LINE; }
+};
+
+void Bucket::addField(FieldDecl *Field, int FieldSize) {
+  Size += FieldSize;
+  Fields.push_back(Field);
+}
+
+struct BitfieldRunBucket : public Bucket {
+  bool canFit(int FieldSize) const override { return true; }
+  bool isBitfieldRun() const override { return true; }
+};
+
+void randomizeStructureLayoutImpl(const ASTContext &Context,
+                                  llvm::SmallVectorImpl<FieldDecl *> &FieldsOut,
+                                  std::mt19937 &RNG) {
+  // All of the Buckets produced by best-effort cache-line algorithm.
+  SmallVector<std::unique_ptr<Bucket>, 16> Buckets;
+
+  // The current bucket of fields that we are trying to fill to a cache-line.
+  std::unique_ptr<Bucket> CurrentBucket;
+
+  // The current bucket containing the run of adjacent bitfields to ensure they
+  // remain adjacent.
+  std::unique_ptr<BitfieldRunBucket> CurrentBitfieldRun;
+
+  // Tracks the number of fields that we failed to fit to the current bucket,
+  // and thus still need to be added later.
+  size_t Skipped = 0;
+
+  while (!FieldsOut.empty()) {
+    // If we've Skipped more fields than we have remaining to place, that means
+    // that they can't fit in our current bucket, and we need to start a new
+    // one.
+    if (Skipped >= FieldsOut.size()) {
+      Skipped = 0;
+      Buckets.push_back(std::move(CurrentBucket));
+    }
+
+    // Take the first field that needs to be put in a bucket.
+    auto FieldIter = FieldsOut.begin();
+    FieldDecl *FD = *FieldIter;
+
+    if (FD->isBitField() && !FD->isZeroLengthBitField(Context)) {
+      // Start a bitfield run if this is the first bitfield we have found.
+      if (!CurrentBitfieldRun)
+        CurrentBitfieldRun = std::make_unique<BitfieldRunBucket>();
+
+      // We've placed the field, and can remove it from the "awaiting Buckets"
+      // vector called "Fields."
+      CurrentBitfieldRun->addField(FD, /*FieldSize is irrelevant here*/ 1);
+      FieldsOut.erase(FieldIter);
+      continue;
+    }
+
+    // Else, current field is not a bitfield. If we were previously in a
+    // bitfield run, end it.
+    if (CurrentBitfieldRun)
+      Buckets.push_back(std::move(CurrentBitfieldRun));
+
+    // If we don't have a bucket, make one.
+    if (!CurrentBucket)
+      CurrentBucket = std::make_unique<Bucket>();
+
+    uint64_t Width = Context.getTypeInfo(FD->getType()).Width;
+    if (Width >= CACHE_LINE) {
+      std::unique_ptr<Bucket> OverSized = std::make_unique<Bucket>();
+      OverSized->addField(FD, Width);
+      FieldsOut.erase(FieldIter);
+      Buckets.push_back(std::move(OverSized));
+      continue;
+    }
+
+    // If it fits, add it.
+    if (CurrentBucket->canFit(Width)) {
+      CurrentBucket->addField(FD, Width);
+      FieldsOut.erase(FieldIter);
+
+      // If it's now full, tie off the bucket.
+      if (CurrentBucket->full()) {
+        Skipped = 0;
+        Buckets.push_back(std::move(CurrentBucket));
+      }
+    } else {
+      // We can't fit it in our current bucket. Move to the end for processing
+      // later.
+      ++Skipped; // Mark it skipped.
+      FieldsOut.push_back(FD);
+      FieldsOut.erase(FieldIter);
+    }
+  }
+
+  // Done processing the fields awaiting a bucket.
+
+  // If we were filling a bucket, tie it off.
+  if (CurrentBucket)
+    Buckets.push_back(std::move(CurrentBucket));
+
+  // If we were processing a bitfield run bucket, tie it off.
+  if (CurrentBitfieldRun)
+    Buckets.push_back(std::move(CurrentBitfieldRun));
+
+  std::shuffle(std::begin(Buckets), std::end(Buckets), RNG);
+
+  // Produce the new ordering of the elements from the Buckets.
+  SmallVector<FieldDecl *, 16> FinalOrder;
+  for (const std::unique_ptr<Bucket> &B : Buckets) {
+    llvm::SmallVectorImpl<FieldDecl *> &RandFields = B->fields();
+    if (!B->isBitfieldRun())
+      std::shuffle(std::begin(RandFields), std::end(RandFields), RNG);
+
+    FinalOrder.insert(FinalOrder.end(), RandFields.begin(), RandFields.end());
+  }
+
+  FieldsOut = FinalOrder;
+}
+
+} // anonymous namespace
+
+namespace clang {
+namespace randstruct {
+
+bool randomizeStructureLayout(const ASTContext &Context, RecordDecl *RD,
+                              SmallVectorImpl<Decl *> &FinalOrdering) {
+  SmallVector<FieldDecl *, 64> RandomizedFields;
+  SmallVector<Decl *, 8> PostRandomizedFields;
+
+  unsigned TotalNumFields = 0;
+  for (Decl *D : RD->decls()) {
+    ++TotalNumFields;
+    if (auto *FD = dyn_cast<FieldDecl>(D))
+      RandomizedFields.push_back(FD);
+    else if (isa<StaticAssertDecl>(D) || isa<IndirectFieldDecl>(D))
+      PostRandomizedFields.push_back(D);
+    else
+      FinalOrdering.push_back(D);
+  }
+
+  if (RandomizedFields.empty())
+    return false;
+
+  // Struct might end with a flexible array or an array of size 0 or 1,
+  // in which case we don't want to randomize it.
+  FieldDecl *FlexibleArray =
+      RD->hasFlexibleArrayMember() ? RandomizedFields.pop_back_val() : nullptr;
+  if (!FlexibleArray) {
+    if (const auto *CA =
+            dyn_cast<ConstantArrayType>(RandomizedFields.back()->getType()))
+      if (CA->getSize().sle(2))
+        FlexibleArray = RandomizedFields.pop_back_val();
+  }
+
+  std::string Seed =
+      Context.getLangOpts().RandstructSeed + RD->getNameAsString();
+  std::seed_seq SeedSeq(Seed.begin(), Seed.end());
+  std::mt19937 RNG(SeedSeq);
+
+  randomizeStructureLayoutImpl(Context, RandomizedFields, RNG);
+
+  // Plorp the randomized decls into the final ordering.
+  FinalOrdering.insert(FinalOrdering.end(), RandomizedFields.begin(),
+                       RandomizedFields.end());
+
+  // Add fields that belong towards the end of the RecordDecl.
+  FinalOrdering.insert(FinalOrdering.end(), PostRandomizedFields.begin(),
+                       PostRandomizedFields.end());
+
+  // Add back the flexible array.
+  if (FlexibleArray)
+    FinalOrdering.push_back(FlexibleArray);
+
+  assert(TotalNumFields == FinalOrdering.size() &&
+         "Decl count has been altered after Randstruct randomization!");
+  (void)TotalNumFields;
+  return true;
+}
+
+} // end namespace randstruct
+} // end namespace clang