228 lines · 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file contains the implementation for Clang's structure field layout10// 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 StaticAssertDecl19#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 a37// 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(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 they74 // 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 means83 // that they can't fit in our current bucket, and we need to start a new84 // 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()) {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 a107 // 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 processing136 // 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 llvm::append_range(FinalOrder, RandFields);163 }164 165 FieldsOut = FinalOrder;166}167 168} // anonymous namespace169 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>(D))182 RandomizedFields.push_back(FD);183 else if (isa<StaticAssertDecl>(D) || isa<IndirectFieldDecl>(D))184 PostRandomizedFields.push_back(D);185 else186 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 llvm::append_range(FinalOrdering, RandomizedFields);212 213 // Add fields that belong towards the end of the RecordDecl.214 llvm::append_range(FinalOrdering, PostRandomizedFields);215 216 // Add back the flexible array.217 if (FlexibleArray)218 FinalOrdering.push_back(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 randstruct227} // end namespace clang228