//===--- 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)6
) {

      // 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)20
)

      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()2
 : 
nullptr17
;

  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