[randstruct] Add randomize structure layout support

The Randstruct feature is a compile-time hardening technique that
randomizes the field layout for designated structures of a code base.
Admittedly, this is mostly useful for closed-source releases of code,
since the randomization seed would need to be available for public and
open source applications.

Why implement it? This patch set enhances Clang’s feature parity with
that of GCC which already has the Randstruct feature. It's used by the
Linux kernel in certain structures to help thwart attacks that depend on
structure layouts in memory.

This patch set is a from-scratch reimplementation of the Randstruct
feature that was originally ported to GCC. The patches for the GCC
implementation can be found here:

  https://www.openwall.com/lists/kernel-hardening/2017/04/06/14

Link: https://lists.llvm.org/pipermail/cfe-dev/2019-March/061607.html
Co-authored-by: Cole Nixon <nixontcole@gmail.com>
Co-authored-by: Connor Kuehl <cipkuehl@gmail.com>
Co-authored-by: James Foster <jafosterja@gmail.com>
Co-authored-by: Jeff Takahashi <jeffrey.takahashi@gmail.com>
Co-authored-by: Jordan Cantrell <jordan.cantrell@mail.com>
Co-authored-by: Nikk Forbus <nicholas.forbus@gmail.com>
Co-authored-by: Tim Pugh <nwtpugh@gmail.com>
Co-authored-by: Bill Wendling <isanbard@gmail.com>
Signed-off-by: Bill Wendling <isanbard@gmail.com>

Reviewed By: aaron.ballman

Differential Revision: https://reviews.llvm.org/D121556

8 files changed, 308 insertions, 0 deletions

diff --git a/clang/lib/AST/CMakeLists.txt b/clang/lib/AST/CMakeLists.txt
index 9365a22..24c8cb8 100644
--- a/clang/lib/AST/CMakeLists.txt
+++ b/clang/lib/AST/CMakeLists.txt
@@ -97,6 +97,7 @@ add_clang_library(clangAST
   ParentMap.cpp
   PrintfFormatString.cpp
   QualTypeNames.cpp
+  Randstruct.cpp
   RawCommentList.cpp
   RecordLayout.cpp
   RecordLayoutBuilder.cpp
diff --git a/clang/lib/AST/Decl.cpp b/clang/lib/AST/Decl.cpp
index 6af7a37..c651fa2 100644
--- a/clang/lib/AST/Decl.cpp
+++ b/clang/lib/AST/Decl.cpp
@@ -30,6 +30,7 @@
 #include "clang/AST/ODRHash.h"
 #include "clang/AST/PrettyDeclStackTrace.h"
 #include "clang/AST/PrettyPrinter.h"
+#include "clang/AST/Randstruct.h"
 #include "clang/AST/Redeclarable.h"
 #include "clang/AST/Stmt.h"
 #include "clang/AST/TemplateBase.h"
@@ -4583,6 +4584,7 @@ RecordDecl::RecordDecl(Kind DK, TagKind TK, const ASTContext &C,
   setHasNonTrivialToPrimitiveCopyCUnion(false);
   setParamDestroyedInCallee(false);
   setArgPassingRestrictions(APK_CanPassInRegs);
+  setIsRandomized(false);
 }
 
 RecordDecl *RecordDecl::Create(const ASTContext &C, TagKind TK, DeclContext *DC,
@@ -4666,6 +4668,12 @@ bool RecordDecl::isMsStruct(const ASTContext &C) const {
   return hasAttr<MSStructAttr>() || C.getLangOpts().MSBitfields == 1;
 }
 
+void RecordDecl::reorderFields(const SmallVectorImpl<Decl *> &Fields) {
+  std::tie(FirstDecl, LastDecl) = DeclContext::BuildDeclChain(Fields, false);
+  LastDecl->NextInContextAndBits.setPointer(nullptr);
+  setIsRandomized(true);
+}
+
 void RecordDecl::LoadFieldsFromExternalStorage() const {
   ExternalASTSource *Source = getASTContext().getExternalSource();
   assert(hasExternalLexicalStorage() && Source && "No external storage?");
diff --git a/clang/lib/AST/Randstruct.cpp b/clang/lib/AST/Randstruct.cpp
new file mode 100644
index 0000000..f6d7ba1
--- /dev/null
+++ b/clang/lib/AST/Randstruct.cpp
@@ -0,0 +1,223 @@
+//===--- 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/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)) {
+      // 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, StringRef Name,
+                              ArrayRef<Decl *> Fields,
+                              SmallVectorImpl<Decl *> &FinalOrdering) {
+  SmallVector<FieldDecl *, 64> RandomizedFields;
+
+  unsigned TotalNumFields = 0;
+  for (Decl *D : Fields) {
+    ++TotalNumFields;
+    if (auto *FD = dyn_cast<FieldDecl>(D))
+      RandomizedFields.push_back(FD);
+    else
+      FinalOrdering.push_back(D);
+  }
+
+  if (RandomizedFields.empty())
+    return false;
+
+  // Struct might end with a variable-length array or an array of size 0 or 1,
+  // in which case we don't want to randomize it.
+  FieldDecl *VLA = nullptr;
+  const auto *CA =
+      dyn_cast<ConstantArrayType>(RandomizedFields.back()->getType());
+  if ((CA && (CA->getSize().sle(2) || CA->isIncompleteArrayType())) ||
+      llvm::any_of(Fields, [](Decl *D) {
+        if (const FieldDecl *FD = dyn_cast<FieldDecl>(D)) {
+          const Type *FDTy = FD->getType().getTypePtr();
+          if (const RecordType *FDTTy = FDTy->getAs<RecordType>())
+            return FDTTy->getDecl()->hasFlexibleArrayMember();
+        }
+        return false;
+      }))
+    VLA = RandomizedFields.pop_back_val();
+
+  std::string Seed = (Context.getLangOpts().RandstructSeed + Name).str();
+  std::seed_seq SeedSeq(Seed.begin(), Seed.end());
+  std::mt19937 RNG(SeedSeq);
+
+  randomizeStructureLayoutImpl(Context, RandomizedFields, RNG);
+  if (VLA)
+    RandomizedFields.push_back(VLA);
+
+  FinalOrdering.insert(FinalOrdering.end(), RandomizedFields.begin(),
+                       RandomizedFields.end());
+
+  assert(TotalNumFields == FinalOrdering.size() &&
+         "Decl count has been altered after Randstruct randomization!");
+  return true;
+}
+
+} // end namespace randstruct
+} // end namespace clang
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index baac938..1601dba 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -5907,6 +5907,14 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
     CmdArgs.push_back(
         Args.MakeArgString("-fmessage-length=" + Twine(MessageLength)));
 
+  if (Arg *A = Args.getLastArg(options::OPT_frandomize_layout_seed_EQ))
+    CmdArgs.push_back(
+        Args.MakeArgString("-frandomize-layout-seed=" + Twine(A->getValue(0))));
+
+  if (Arg *A = Args.getLastArg(options::OPT_frandomize_layout_seed_file_EQ))
+    CmdArgs.push_back(Args.MakeArgString("-frandomize-layout-seed-file=" +
+                                         Twine(A->getValue(0))));
+
   // -fvisibility= and -fvisibility-ms-compat are of a piece.
   if (const Arg *A = Args.getLastArg(options::OPT_fvisibility_EQ,
                                      options::OPT_fvisibility_ms_compat)) {
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 83de27b..5181999 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -94,6 +94,7 @@
 #include <cassert>
 #include <cstddef>
 #include <cstring>
+#include <fstream>
 #include <memory>
 #include <string>
 #include <tuple>
@@ -3685,6 +3686,9 @@ void CompilerInvocation::GenerateLangArgs(const LangOptions &Opts,
 
   for (const auto &MP : Opts.MacroPrefixMap)
     GenerateArg(Args, OPT_fmacro_prefix_map_EQ, MP.first + "=" + MP.second, SA);
+
+  if (!Opts.RandstructSeed.empty())
+    GenerateArg(Args, OPT_frandomize_layout_seed_EQ, Opts.RandstructSeed, SA);
 }
 
 bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
@@ -4237,6 +4241,19 @@ bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
       Diags.Report(diag::err_cc1_unbounded_vscale_min);
   }
 
+  if (const Arg *A = Args.getLastArg(OPT_frandomize_layout_seed_file_EQ)) {
+    std::ifstream SeedFile(A->getValue(0));
+
+    if (!SeedFile.is_open())
+      Diags.Report(diag::err_drv_cannot_open_randomize_layout_seed_file)
+          << A->getValue(0);
+
+    std::getline(SeedFile, Opts.RandstructSeed);
+  }
+
+  if (const Arg *A = Args.getLastArg(OPT_frandomize_layout_seed_EQ))
+    Opts.RandstructSeed = A->getValue(0);
+
   return Diags.getNumErrors() == NumErrorsBefore;
 }
 
diff --git a/clang/lib/Sema/SemaCast.cpp b/clang/lib/Sema/SemaCast.cpp
index 7d580dc..7b5bc7c 100644
--- a/clang/lib/Sema/SemaCast.cpp
+++ b/clang/lib/Sema/SemaCast.cpp
@@ -3129,6 +3129,23 @@ void CastOperation::CheckCStyleCast() {
     Self.Diag(OpRange.getBegin(), diag::warn_cast_function_type)
         << SrcType << DestType << OpRange;
 
+  if (isa<PointerType>(SrcType) && isa<PointerType>(DestType)) {
+    QualType SrcTy = cast<PointerType>(SrcType)->getPointeeType();
+    QualType DestTy = cast<PointerType>(DestType)->getPointeeType();
+
+    const RecordDecl *SrcRD = SrcTy->getAsRecordDecl();
+    const RecordDecl *DestRD = DestTy->getAsRecordDecl();
+
+    if (SrcRD && DestRD && SrcRD->hasAttr<RandomizeLayoutAttr>() &&
+        SrcRD != DestRD) {
+      // The struct we are casting the pointer from was randomized.
+      Self.Diag(OpRange.getBegin(), diag::err_cast_from_randomized_struct)
+          << SrcType << DestType;
+      SrcExpr = ExprError();
+      return;
+    }
+  }
+
   DiagnoseCastOfObjCSEL(Self, SrcExpr, DestType);
   DiagnoseCallingConvCast(Self, SrcExpr, DestType, OpRange);
   DiagnoseBadFunctionCast(Self, SrcExpr, DestType);
diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index 427b257..ce1eb2e 100644
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -24,6 +24,7 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/NonTrivialTypeVisitor.h"
+#include "clang/AST/Randstruct.h"
 #include "clang/AST/StmtCXX.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/PartialDiagnostic.h"
@@ -17844,6 +17845,18 @@ void Sema::ActOnFields(Scope *S, SourceLocation RecLoc, Decl *EnclosingDecl,
     // Handle attributes before checking the layout.
     ProcessDeclAttributeList(S, Record, Attrs);
 
+    // Maybe randomize the field order.
+    if (!getLangOpts().CPlusPlus && Record->hasAttr<RandomizeLayoutAttr>() &&
+        !Record->isUnion() && !getLangOpts().RandstructSeed.empty() &&
+        !Record->isRandomized()) {
+      SmallVector<Decl *, 32> OrigFieldOrdering(Record->fields());
+      SmallVector<Decl *, 32> NewFieldOrdering;
+      if (randstruct::randomizeStructureLayout(
+              Context, Record->getNameAsString(), OrigFieldOrdering,
+              NewFieldOrdering))
+        Record->reorderFields(NewFieldOrdering);
+    }
+
     // We may have deferred checking for a deleted destructor. Check now.
     if (CXXRecord) {
       auto *Dtor = CXXRecord->getDestructor();
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 4b5201d..843ddbd 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -5121,6 +5121,21 @@ static void handleLifetimeCategoryAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   }
 }
 
+static void handleRandomizeLayoutAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
+  if (checkAttrMutualExclusion<NoRandomizeLayoutAttr>(S, D, AL))
+    return;
+  if (!D->hasAttr<RandomizeLayoutAttr>())
+    D->addAttr(::new (S.Context) RandomizeLayoutAttr(S.Context, AL));
+}
+
+static void handleNoRandomizeLayoutAttr(Sema &S, Decl *D,
+                                        const ParsedAttr &AL) {
+  if (checkAttrMutualExclusion<RandomizeLayoutAttr>(S, D, AL))
+    return;
+  if (!D->hasAttr<NoRandomizeLayoutAttr>())
+    D->addAttr(::new (S.Context) NoRandomizeLayoutAttr(S.Context, AL));
+}
+
 bool Sema::CheckCallingConvAttr(const ParsedAttr &Attrs, CallingConv &CC,
                                 const FunctionDecl *FD) {
   if (Attrs.isInvalid())
@@ -8635,6 +8650,12 @@ static void ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D,
   case ParsedAttr::AT_Section:
     handleSectionAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_RandomizeLayout:
+    handleRandomizeLayoutAttr(S, D, AL);
+    break;
+  case ParsedAttr::AT_NoRandomizeLayout:
+    handleNoRandomizeLayoutAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_CodeSeg:
     handleCodeSegAttr(S, D, AL);
     break;