clang/lib/AST/Randstruct.cpp

   1 //===--- Randstruct.cpp ---------------------------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file contains the implementation for Clang's structure field layout
  10 // randomization.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "clang/AST/Randstruct.h"
  15 #include "clang/AST/ASTContext.h"
  16 #include "clang/AST/Attr.h"
  17 #include "clang/AST/Decl.h"
  18 #include "clang/AST/DeclCXX.h" // For StaticAssertDecl
  19 #include "clang/Basic/Diagnostic.h"
  20 #include "llvm/ADT/SmallVector.h"
  21
  22 #include <algorithm>
  23 #include <random>
  24 #include <set>
  25 #include <string>
  26
  27 using clang::ASTContext;
  28 using clang::FieldDecl;
  29 using llvm::SmallVector;
  30
  31 namespace {
  32
  33 // FIXME: Replace this with some discovery once that mechanism exists.
  34 enum { CACHE_LINE = 64 };
  35
  36 // The Bucket class holds the struct fields we're trying to fill to a
  37 // cache-line.
  38 class Bucket {
  39   SmallVector<FieldDecl *, 64> Fields;
  40   int Size = 0;
  41
  42 public:
  43   virtual ~Bucket() = default;
  44
  45   SmallVector<FieldDecl *, 64> &fields() { return Fields; }
  46   void addField(FieldDecl *Field, int FieldSize);
  47   virtual bool canFit(int FieldSize) const {
  48     return Size + FieldSize <= CACHE_LINE;
  49   }
  50   virtual bool isBitfieldRun() const { return false; }
  51   bool full() const { return Size >= CACHE_LINE; }
  52 };
  53
  54 void Bucket::addField(FieldDecl *Field, int FieldSize) {
  55   Size += FieldSize;
  56   Fields.push_back(Field);
  57 }
  58
  59 struct BitfieldRunBucket : public Bucket {
  60   bool canFit(int FieldSize) const override { return true; }
  61   bool isBitfieldRun() const override { return true; }
  62 };
  63
  64 void randomizeStructureLayoutImpl(const ASTContext &Context,
  65                                   llvm::SmallVectorImpl<FieldDecl *> &FieldsOut,
  66                                   std::mt19937 &RNG) {
  67   // All of the Buckets produced by best-effort cache-line algorithm.
  68   SmallVector<std::unique_ptr<Bucket>, 16> Buckets;
  69
  70   // The current bucket of fields that we are trying to fill to a cache-line.
  71   std::unique_ptr<Bucket> CurrentBucket;
  72
  73   // The current bucket containing the run of adjacent bitfields to ensure they
  74   // remain adjacent.
  75   std::unique_ptr<BitfieldRunBucket> CurrentBitfieldRun;
  76
  77   // Tracks the number of fields that we failed to fit to the current bucket,
  78   // and thus still need to be added later.
  79   size_t Skipped = 0;
  80
  81   while (!FieldsOut.empty()) {
  82     // If we've Skipped more fields than we have remaining to place, that means
  83     // that they can't fit in our current bucket, and we need to start a new
  84     // one.
  85     if (Skipped >= FieldsOut.size()) {
  86       Skipped = 0;
  87       Buckets.push_back(std::move(CurrentBucket));
  88     }
  89
  90     // Take the first field that needs to be put in a bucket.
  91     auto FieldIter = FieldsOut.begin();
  92     FieldDecl *FD = *FieldIter;
  93
  94     if (FD->isBitField() && !FD->isZeroLengthBitField(Context)) {
  95       // Start a bitfield run if this is the first bitfield we have found.
  96       if (!CurrentBitfieldRun)
  97         CurrentBitfieldRun = std::make_unique<BitfieldRunBucket>();
  98
  99       // We've placed the field, and can remove it from the "awaiting Buckets"
 100       // vector called "Fields."
 101       CurrentBitfieldRun->addField(FD, /*FieldSize is irrelevant here*/ 1);
 102       FieldsOut.erase(FieldIter);
 103       continue;
 104     }
 105
 106     // Else, current field is not a bitfield. If we were previously in a
 107     // bitfield run, end it.
 108     if (CurrentBitfieldRun)
 109       Buckets.push_back(std::move(CurrentBitfieldRun));
 110
 111     // If we don't have a bucket, make one.
 112     if (!CurrentBucket)
 113       CurrentBucket = std::make_unique<Bucket>();
 114
 115     uint64_t Width = Context.getTypeInfo(FD->getType()).Width;
 116     if (Width >= CACHE_LINE) {
 117       std::unique_ptr<Bucket> OverSized = std::make_unique<Bucket>();
 118       OverSized->addField(FD, Width);
 119       FieldsOut.erase(FieldIter);
 120       Buckets.push_back(std::move(OverSized));
 121       continue;
 122     }
 123
 124     // If it fits, add it.
 125     if (CurrentBucket->canFit(Width)) {
 126       CurrentBucket->addField(FD, Width);
 127       FieldsOut.erase(FieldIter);
 128
 129       // If it's now full, tie off the bucket.
 130       if (CurrentBucket->full()) {
 131         Skipped = 0;
 132         Buckets.push_back(std::move(CurrentBucket));
 133       }
 134     } else {
 135       // We can't fit it in our current bucket. Move to the end for processing
 136       // later.
 137       ++Skipped; // Mark it skipped.
 138       FieldsOut.push_back(FD);
 139       FieldsOut.erase(FieldIter);
 140     }
 141   }
 142
 143   // Done processing the fields awaiting a bucket.
 144
 145   // If we were filling a bucket, tie it off.
 146   if (CurrentBucket)
 147     Buckets.push_back(std::move(CurrentBucket));
 148
 149   // If we were processing a bitfield run bucket, tie it off.
 150   if (CurrentBitfieldRun)
 151     Buckets.push_back(std::move(CurrentBitfieldRun));
 152
 153   std::shuffle(std::begin(Buckets), std::end(Buckets), RNG);
 154
 155   // Produce the new ordering of the elements from the Buckets.
 156   SmallVector<FieldDecl *, 16> FinalOrder;
 157   for (const std::unique_ptr<Bucket> &B : Buckets) {
 158     llvm::SmallVectorImpl<FieldDecl *> &RandFields = B->fields();
 159     if (!B->isBitfieldRun())
 160       std::shuffle(std::begin(RandFields), std::end(RandFields), RNG);
 161
 162     FinalOrder.insert(FinalOrder.end(), RandFields.begin(), RandFields.end());
 163   }
 164
 165   FieldsOut = FinalOrder;
 166 }
 167
 168 } // anonymous namespace
 169
 170 namespace clang {
 171 namespace randstruct {
 172
 173 bool randomizeStructureLayout(const ASTContext &Context, RecordDecl *RD,
 174                               SmallVectorImpl<Decl *> &FinalOrdering) {
 175   SmallVector<FieldDecl *, 64> RandomizedFields;
 176   SmallVector<Decl *, 8> PostRandomizedFields;
 177
 178   unsigned TotalNumFields = 0;
 179   for (Decl *D : RD->decls()) {
 180     ++TotalNumFields;
 181     if (auto *FD = dyn_cast<FieldDecl>(D))
 182       RandomizedFields.push_back(FD);
 183     else if (isa<StaticAssertDecl>(D) || isa<IndirectFieldDecl>(D))
 184       PostRandomizedFields.push_back(D);
 185     else
 186       FinalOrdering.push_back(D);
 187   }
 188
 189   if (RandomizedFields.empty())
 190     return false;
 191
 192   // Struct might end with a flexible array or an array of size 0 or 1,
 193   // in which case we don't want to randomize it.
 194   FieldDecl *FlexibleArray =
 195       RD->hasFlexibleArrayMember() ? RandomizedFields.pop_back_val() : nullptr;
 196   if (!FlexibleArray) {
 197     if (const auto *CA =
 198             dyn_cast<ConstantArrayType>(RandomizedFields.back()->getType()))
 199       if (CA->getSize().sle(2))
 200         FlexibleArray = RandomizedFields.pop_back_val();
 201   }
 202
 203   std::string Seed =
 204       Context.getLangOpts().RandstructSeed + RD->getNameAsString();
 205   std::seed_seq SeedSeq(Seed.begin(), Seed.end());
 206   std::mt19937 RNG(SeedSeq);
 207
 208   randomizeStructureLayoutImpl(Context, RandomizedFields, RNG);
 209
 210   // Plorp the randomized decls into the final ordering.
 211   FinalOrdering.insert(FinalOrdering.end(), RandomizedFields.begin(),
 212                        RandomizedFields.end());
 213
 214   // Add fields that belong towards the end of the RecordDecl.
 215   FinalOrdering.insert(FinalOrdering.end(), PostRandomizedFields.begin(),
 216                        PostRandomizedFields.end());
 217
 218   // Add back the flexible array.
 219   if (FlexibleArray)
 220     FinalOrdering.push_back(FlexibleArray);
 221
 222   assert(TotalNumFields == FinalOrdering.size() &&
 223          "Decl count has been altered after Randstruct randomization!");
 224   (void)TotalNumFields;
 225   return true;
 226 }
 227
 228 } // end namespace randstruct
 229 } // end namespace clang