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
27using clang::ASTContext;
28using clang::FieldDecl;
29using llvm::SmallVector;
30
31namespace {
32
33// FIXME: Replace this with some discovery once that mechanism exists.
34enum { CACHE_LINE = 64 };
35
36// The Bucket class holds the struct fields we're trying to fill to a
37// cache-line.
38class Bucket {
39 SmallVector<FieldDecl *, 64> Fields;
40 int Size = 0;
41
42public:
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
54void Bucket::addField(FieldDecl *Field, int FieldSize) {
55 Size += FieldSize;
56 Fields.push_back(Elt: Field);
57}
58
59struct BitfieldRunBucket : public Bucket {
60 bool canFit(int FieldSize) const override { return true; }
61 bool isBitfieldRun() const override { return true; }
62};
63
64void 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(Elt: 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()) {
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(Field: FD, /*FieldSize is irrelevant here*/ FieldSize: 1);
102 FieldsOut.erase(CI: 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(Elt: 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(T: FD->getType()).Width;
116 if (Width >= CACHE_LINE) {
117 std::unique_ptr<Bucket> OverSized = std::make_unique<Bucket>();
118 OverSized->addField(Field: FD, FieldSize: Width);
119 FieldsOut.erase(CI: FieldIter);
120 Buckets.push_back(Elt: std::move(OverSized));
121 continue;
122 }
123
124 // If it fits, add it.
125 if (CurrentBucket->canFit(FieldSize: Width)) {
126 CurrentBucket->addField(Field: FD, FieldSize: Width);
127 FieldsOut.erase(CI: FieldIter);
128
129 // If it's now full, tie off the bucket.
130 if (CurrentBucket->full()) {
131 Skipped = 0;
132 Buckets.push_back(Elt: 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(Elt: FD);
139 FieldsOut.erase(CI: 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(Elt: std::move(CurrentBucket));
148
149 // If we were processing a bitfield run bucket, tie it off.
150 if (CurrentBitfieldRun)
151 Buckets.push_back(Elt: std::move(CurrentBitfieldRun));
152
153 std::shuffle(first: std::begin(cont&: Buckets), last: std::end(cont&: Buckets), g&: 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(first: std::begin(cont&: RandFields), last: std::end(cont&: RandFields), g&: RNG);
161
162 llvm::append_range(C&: FinalOrder, R&: RandFields);
163 }
164
165 FieldsOut = FinalOrder;
166}
167
168} // anonymous namespace
169
170namespace clang {
171namespace randstruct {
172
173bool 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>(Val: D))
182 RandomizedFields.push_back(Elt: FD);
183 else if (isa<StaticAssertDecl>(Val: D) || isa<IndirectFieldDecl>(Val: D))
184 PostRandomizedFields.push_back(Elt: D);
185 else
186 FinalOrdering.push_back(Elt: 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>(Val: RandomizedFields.back()->getType()))
199 if (CA->getSize().sle(RHS: 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, FieldsOut&: RandomizedFields, RNG);
209
210 // Plorp the randomized decls into the final ordering.
211 llvm::append_range(C&: FinalOrdering, R&: RandomizedFields);
212
213 // Add fields that belong towards the end of the RecordDecl.
214 llvm::append_range(C&: FinalOrdering, R&: PostRandomizedFields);
215
216 // Add back the flexible array.
217 if (FlexibleArray)
218 FinalOrdering.push_back(Elt: FlexibleArray);
219
220 assert(TotalNumFields == FinalOrdering.size() &&
221 "Decl count has been altered after Randstruct randomization!");
222 (void)TotalNumFields;
223 return true;
224}
225
226} // end namespace randstruct
227} // end namespace clang
228