llvm-project

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/ASTDiagnostic.h"
17
#include "clang/AST/Attr.h"
18
#include "clang/AST/Decl.h"
19
#include "clang/AST/DeclCXX.h" // For StaticAssertDecl
20
#include "clang/Basic/Diagnostic.h"
21
#include "llvm/ADT/SmallVector.h"
22

23
#include <algorithm>
24
#include <random>
25
#include <set>
26
#include <sstream>
27
#include <string>
28

29
using clang::ASTContext;
30
using clang::FieldDecl;
31
using llvm::SmallVector;
32

33
namespace {
34

35
// FIXME: Replace this with some discovery once that mechanism exists.
36
enum { CACHE_LINE = 64 };
37

38
// The Bucket class holds the struct fields we're trying to fill to a
39
// cache-line.
40
class Bucket {
41
  SmallVector<FieldDecl *, 64> Fields;
42
  int Size = 0;
43

44
public:
45
  virtual ~Bucket() = default;
46

47
  SmallVector<FieldDecl *, 64> &fields() { return Fields; }
48
  void addField(FieldDecl *Field, int FieldSize);
49
  virtual bool canFit(int FieldSize) const {
50
    return Size + FieldSize <= CACHE_LINE;
51
  }
52
  virtual bool isBitfieldRun() const { return false; }
53
  bool full() const { return Size >= CACHE_LINE; }
54
};
55

56
void Bucket::addField(FieldDecl *Field, int FieldSize) {
57
  Size += FieldSize;
58
  Fields.push_back(Field);
59
}
60

61
struct BitfieldRunBucket : public Bucket {
62
  bool canFit(int FieldSize) const override { return true; }
63
  bool isBitfieldRun() const override { return true; }
64
};
65

66
void randomizeStructureLayoutImpl(const ASTContext &Context,
67
                                  llvm::SmallVectorImpl<FieldDecl *> &FieldsOut,
68
                                  std::mt19937 &RNG) {
69
  // All of the Buckets produced by best-effort cache-line algorithm.
70
  SmallVector<std::unique_ptr<Bucket>, 16> Buckets;
71

72
  // The current bucket of fields that we are trying to fill to a cache-line.
73
  std::unique_ptr<Bucket> CurrentBucket;
74

75
  // The current bucket containing the run of adjacent bitfields to ensure they
76
  // remain adjacent.
77
  std::unique_ptr<BitfieldRunBucket> CurrentBitfieldRun;
78

79
  // Tracks the number of fields that we failed to fit to the current bucket,
80
  // and thus still need to be added later.
81
  size_t Skipped = 0;
82

83
  while (!FieldsOut.empty()) {
84
    // If we've Skipped more fields than we have remaining to place, that means
85
    // that they can't fit in our current bucket, and we need to start a new
86
    // one.
87
    if (Skipped >= FieldsOut.size()) {
88
      Skipped = 0;
89
      Buckets.push_back(std::move(CurrentBucket));
90
    }
91

92
    // Take the first field that needs to be put in a bucket.
93
    auto FieldIter = FieldsOut.begin();
94
    FieldDecl *FD = *FieldIter;
95

96
    if (FD->isBitField() && !FD->isZeroLengthBitField(Context)) {
97
      // Start a bitfield run if this is the first bitfield we have found.
98
      if (!CurrentBitfieldRun)
99
        CurrentBitfieldRun = std::make_unique<BitfieldRunBucket>();
100

101
      // We've placed the field, and can remove it from the "awaiting Buckets"
102
      // vector called "Fields."
103
      CurrentBitfieldRun->addField(FD, /*FieldSize is irrelevant here*/ 1);
104
      FieldsOut.erase(FieldIter);
105
      continue;
106
    }
107

108
    // Else, current field is not a bitfield. If we were previously in a
109
    // bitfield run, end it.
110
    if (CurrentBitfieldRun)
111
      Buckets.push_back(std::move(CurrentBitfieldRun));
112

113
    // If we don't have a bucket, make one.
114
    if (!CurrentBucket)
115
      CurrentBucket = std::make_unique<Bucket>();
116

117
    uint64_t Width = Context.getTypeInfo(FD->getType()).Width;
118
    if (Width >= CACHE_LINE) {
119
      std::unique_ptr<Bucket> OverSized = std::make_unique<Bucket>();
120
      OverSized->addField(FD, Width);
121
      FieldsOut.erase(FieldIter);
122
      Buckets.push_back(std::move(OverSized));
123
      continue;
124
    }
125

126
    // If it fits, add it.
127
    if (CurrentBucket->canFit(Width)) {
128
      CurrentBucket->addField(FD, Width);
129
      FieldsOut.erase(FieldIter);
130

131
      // If it's now full, tie off the bucket.
132
      if (CurrentBucket->full()) {
133
        Skipped = 0;
134
        Buckets.push_back(std::move(CurrentBucket));
135
      }
136
    } else {
137
      // We can't fit it in our current bucket. Move to the end for processing
138
      // later.
139
      ++Skipped; // Mark it skipped.
140
      FieldsOut.push_back(FD);
141
      FieldsOut.erase(FieldIter);
142
    }
143
  }
144

145
  // Done processing the fields awaiting a bucket.
146

147
  // If we were filling a bucket, tie it off.
148
  if (CurrentBucket)
149
    Buckets.push_back(std::move(CurrentBucket));
150

151
  // If we were processing a bitfield run bucket, tie it off.
152
  if (CurrentBitfieldRun)
153
    Buckets.push_back(std::move(CurrentBitfieldRun));
154

155
  std::shuffle(std::begin(Buckets), std::end(Buckets), RNG);
156

157
  // Produce the new ordering of the elements from the Buckets.
158
  SmallVector<FieldDecl *, 16> FinalOrder;
159
  for (const std::unique_ptr<Bucket> &B : Buckets) {
160
    llvm::SmallVectorImpl<FieldDecl *> &RandFields = B->fields();
161
    if (!B->isBitfieldRun())
162
      std::shuffle(std::begin(RandFields), std::end(RandFields), RNG);
163

164
    FinalOrder.insert(FinalOrder.end(), RandFields.begin(), RandFields.end());
165
  }
166

167
  FieldsOut = FinalOrder;
168
}
169

170
} // anonymous namespace
171

172
namespace clang {
173
namespace randstruct {
174

175
bool randomizeStructureLayout(const ASTContext &Context, RecordDecl *RD,
176
                              SmallVectorImpl<Decl *> &FinalOrdering) {
177
  SmallVector<FieldDecl *, 64> RandomizedFields;
178
  SmallVector<Decl *, 8> PostRandomizedFields;
179

180
  unsigned TotalNumFields = 0;
181
  for (Decl *D : RD->decls()) {
182
    ++TotalNumFields;
183
    if (auto *FD = dyn_cast<FieldDecl>(D))
184
      RandomizedFields.push_back(FD);
185
    else if (isa<StaticAssertDecl>(D) || isa<IndirectFieldDecl>(D))
186
      PostRandomizedFields.push_back(D);
187
    else
188
      FinalOrdering.push_back(D);
189
  }
190

191
  if (RandomizedFields.empty())
192
    return false;
193

194
  // Struct might end with a flexible array or an array of size 0 or 1,
195
  // in which case we don't want to randomize it.
196
  FieldDecl *FlexibleArray =
197
      RD->hasFlexibleArrayMember() ? RandomizedFields.pop_back_val() : nullptr;
198
  if (!FlexibleArray) {
199
    if (const auto *CA =
200
            dyn_cast<ConstantArrayType>(RandomizedFields.back()->getType()))
201
      if (CA->getSize().sle(2))
202
        FlexibleArray = RandomizedFields.pop_back_val();
203
  }
204

205
  std::string Seed =
206
      Context.getLangOpts().RandstructSeed + RD->getNameAsString();
207
  std::seed_seq SeedSeq(Seed.begin(), Seed.end());
208
  std::mt19937 RNG(SeedSeq);
209

210
  randomizeStructureLayoutImpl(Context, RandomizedFields, RNG);
211

212
  // Plorp the randomized decls into the final ordering.
213
  FinalOrdering.insert(FinalOrdering.end(), RandomizedFields.begin(),
214
                       RandomizedFields.end());
215

216
  // Add fields that belong towards the end of the RecordDecl.
217
  FinalOrdering.insert(FinalOrdering.end(), PostRandomizedFields.begin(),
218
                       PostRandomizedFields.end());
219

220
  // Add back the flexible array.
221
  if (FlexibleArray)
222
    FinalOrdering.push_back(FlexibleArray);
223

224
  assert(TotalNumFields == FinalOrdering.size() &&
225
         "Decl count has been altered after Randstruct randomization!");
226
  (void)TotalNumFields;
227
  return true;
228
}
229

230
} // end namespace randstruct
231
} // end namespace clang
232