From 09f1542418fb6312e2252b4dead8a16ad0817e30 Mon Sep 17 00:00:00 2001
From: Chase Wilson <me@chasewilson.dev>
Date: Sun, 29 Aug 2021 08:55:29 -0700
Subject: [PATCH] Implemented -Z randomize-layout

---
 Cargo.lock                             | 21 ++++++++---
 compiler/rustc_middle/Cargo.toml       |  2 ++
 compiler/rustc_middle/src/ty/layout.rs | 48 +++++++++++++++++++-------
 compiler/rustc_middle/src/ty/mod.rs    | 38 +++++++++++++++++++-
 compiler/rustc_session/src/options.rs  |  2 ++
 src/tools/tidy/src/deps.rs             |  1 +
 6 files changed, 93 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6b3a714920b..9eff4dd8f6b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1660,7 +1660,7 @@ checksum = "3ca8957e71f04a205cb162508f9326aea04676c8dfd0711220190d6b83664f3f"
 dependencies = [
  "bitmaps",
  "rand_core 0.5.1",
- "rand_xoshiro",
+ "rand_xoshiro 0.4.0",
  "sized-chunks",
  "typenum",
  "version_check",
@@ -2256,7 +2256,7 @@ dependencies = [
  "libc",
  "log",
  "measureme",
- "rand 0.8.3",
+ "rand 0.8.4",
  "rustc-workspace-hack",
  "rustc_version",
  "shell-escape",
@@ -2852,9 +2852,9 @@ dependencies = [
 
 [[package]]
 name = "rand"
-version = "0.8.3"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0ef9e7e66b4468674bfcb0c81af8b7fa0bb154fa9f28eb840da5c447baeb8d7e"
+checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8"
 dependencies = [
  "libc",
  "rand_chacha 0.3.0",
@@ -2945,6 +2945,15 @@ dependencies = [
  "rand_core 0.5.1",
 ]
 
+[[package]]
+name = "rand_xoshiro"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
+dependencies = [
+ "rand_core 0.6.2",
+]
+
 [[package]]
 name = "rayon"
 version = "1.3.1"
@@ -4086,6 +4095,8 @@ dependencies = [
  "either",
  "gsgdt",
  "polonius-engine",
+ "rand 0.8.4",
+ "rand_xoshiro 0.6.0",
  "rustc-rayon-core",
  "rustc_apfloat",
  "rustc_arena",
@@ -5096,7 +5107,7 @@ checksum = "dac1c663cfc93810f88aed9b8941d48cabf856a1b111c29a40439018d870eb22"
 dependencies = [
  "cfg-if 1.0.0",
  "libc",
- "rand 0.8.3",
+ "rand 0.8.4",
  "redox_syscall",
  "remove_dir_all",
  "winapi",
diff --git a/compiler/rustc_middle/Cargo.toml b/compiler/rustc_middle/Cargo.toml
index b1fcc34bee1..d06c593d394 100644
--- a/compiler/rustc_middle/Cargo.toml
+++ b/compiler/rustc_middle/Cargo.toml
@@ -32,3 +32,5 @@ chalk-ir = "0.55.0"
 smallvec = { version = "1.6.1", features = ["union", "may_dangle"] }
 rustc_session = { path = "../rustc_session" }
 rustc_type_ir = { path = "../rustc_type_ir" }
+rand = "0.8.4"
+rand_xoshiro = "0.6.0"
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index cfbbec374a1..b6aeb9122c3 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -24,6 +24,9 @@ use std::iter;
 use std::num::NonZeroUsize;
 use std::ops::Bound;
 
+use rand::{seq::SliceRandom, SeedableRng};
+use rand_xoshiro::Xoshiro128StarStar;
+
 pub fn provide(providers: &mut ty::query::Providers) {
     *providers =
         ty::query::Providers { layout_of, fn_abi_of_fn_ptr, fn_abi_of_instance, ..*providers };
@@ -324,6 +327,10 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
 
         let mut inverse_memory_index: Vec<u32> = (0..fields.len() as u32).collect();
 
+        // `ReprOptions.layout_seed` is a deterministic seed that we can use to
+        // randomize field ordering with
+        let mut rng = Xoshiro128StarStar::seed_from_u64(repr.field_shuffle_seed);
+
         let optimize = !repr.inhibit_struct_field_reordering_opt();
         if optimize {
             let end =
@@ -332,20 +339,35 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
             let field_align = |f: &TyAndLayout<'_>| {
                 if let Some(pack) = pack { f.align.abi.min(pack) } else { f.align.abi }
             };
-            match kind {
-                StructKind::AlwaysSized | StructKind::MaybeUnsized => {
-                    optimizing.sort_by_key(|&x| {
-                        // Place ZSTs first to avoid "interesting offsets",
-                        // especially with only one or two non-ZST fields.
-                        let f = &fields[x as usize];
-                        (!f.is_zst(), cmp::Reverse(field_align(f)))
-                    });
-                }
-                StructKind::Prefixed(..) => {
-                    // Sort in ascending alignment so that the layout stay optimal
-                    // regardless of the prefix
-                    optimizing.sort_by_key(|&x| field_align(&fields[x as usize]));
+
+            // If `-Z randomize-layout` was enabled for the type definition we can shuffle
+            // the field ordering to try and catch some code making assumptions about layouts
+            // we don't guarantee
+            if repr.can_randomize_type_layout() {
+                // Shuffle the ordering of the fields
+                optimizing.shuffle(&mut rng);
+
+            // Otherwise we just leave things alone and actually optimize the type's fields
+            } else {
+                match kind {
+                    StructKind::AlwaysSized | StructKind::MaybeUnsized => {
+                        optimizing.sort_by_key(|&x| {
+                            // Place ZSTs first to avoid "interesting offsets",
+                            // especially with only one or two non-ZST fields.
+                            let f = &fields[x as usize];
+                            (!f.is_zst(), cmp::Reverse(field_align(f)))
+                        });
+                    }
+
+                    StructKind::Prefixed(..) => {
+                        // Sort in ascending alignment so that the layout stays optimal
+                        // regardless of the prefix
+                        optimizing.sort_by_key(|&x| field_align(&fields[x as usize]));
+                    }
                 }
+
+                // FIXME(Kixiron): We can always shuffle fields within a given alignment class
+                //                 regardless of the status of `-Z randomize-layout`
             }
         }
 
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index b3ae76d9871..d04d1565fea 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -1491,6 +1491,9 @@ bitflags! {
         const IS_LINEAR          = 1 << 3;
         // If true, don't expose any niche to type's context.
         const HIDE_NICHE         = 1 << 4;
+        // If true, the type's layout can be randomized using
+        // the seed stored in `ReprOptions.layout_seed`
+        const RANDOMIZE_LAYOUT   = 1 << 5;
         // Any of these flags being set prevent field reordering optimisation.
         const IS_UNOPTIMISABLE   = ReprFlags::IS_C.bits |
                                    ReprFlags::IS_SIMD.bits |
@@ -1505,6 +1508,14 @@ pub struct ReprOptions {
     pub align: Option<Align>,
     pub pack: Option<Align>,
     pub flags: ReprFlags,
+    /// The seed to be used for randomizing a type's layout
+    ///
+    /// Note: This could technically be a `[u8; 16]` (a `u128`) which would
+    /// be the "most accurate" hash as it'd encompass the item and crate
+    /// hash without loss, but it does pay the price of being larger.
+    /// Everything's a tradeoff, a `u64` seed should be sufficient for our
+    /// purposes (primarily `-Z randomize-layout`)
+    pub field_shuffle_seed: u64,
 }
 
 impl ReprOptions {
@@ -1513,6 +1524,11 @@ impl ReprOptions {
         let mut size = None;
         let mut max_align: Option<Align> = None;
         let mut min_pack: Option<Align> = None;
+
+        // Generate a deterministically-derived seed from the item's path hash
+        // to allow for cross-crate compilation to actually work
+        let field_shuffle_seed = tcx.def_path_hash(did).0.to_smaller_hash();
+
         for attr in tcx.get_attrs(did).iter() {
             for r in attr::find_repr_attrs(&tcx.sess, attr) {
                 flags.insert(match r {
@@ -1541,33 +1557,45 @@ impl ReprOptions {
             }
         }
 
+        // If `-Z randomize-layout` was enabled for the type definition then we can
+        // consider performing layout randomization
+        if tcx.sess.opts.debugging_opts.randomize_layout {
+            flags.insert(ReprFlags::RANDOMIZE_LAYOUT);
+        }
+
         // This is here instead of layout because the choice must make it into metadata.
         if !tcx.consider_optimizing(|| format!("Reorder fields of {:?}", tcx.def_path_str(did))) {
             flags.insert(ReprFlags::IS_LINEAR);
         }
-        ReprOptions { int: size, align: max_align, pack: min_pack, flags }
+
+        Self { int: size, align: max_align, pack: min_pack, flags, field_shuffle_seed }
     }
 
     #[inline]
     pub fn simd(&self) -> bool {
         self.flags.contains(ReprFlags::IS_SIMD)
     }
+
     #[inline]
     pub fn c(&self) -> bool {
         self.flags.contains(ReprFlags::IS_C)
     }
+
     #[inline]
     pub fn packed(&self) -> bool {
         self.pack.is_some()
     }
+
     #[inline]
     pub fn transparent(&self) -> bool {
         self.flags.contains(ReprFlags::IS_TRANSPARENT)
     }
+
     #[inline]
     pub fn linear(&self) -> bool {
         self.flags.contains(ReprFlags::IS_LINEAR)
     }
+
     #[inline]
     pub fn hide_niche(&self) -> bool {
         self.flags.contains(ReprFlags::HIDE_NICHE)
@@ -1594,9 +1622,17 @@ impl ReprOptions {
                 return true;
             }
         }
+
         self.flags.intersects(ReprFlags::IS_UNOPTIMISABLE) || self.int.is_some()
     }
 
+    /// Returns `true` if this type is valid for reordering and `-Z randomize-layout`
+    /// was enabled for its declaration crate
+    pub fn can_randomize_type_layout(&self) -> bool {
+        !self.inhibit_struct_field_reordering_opt()
+            && self.flags.contains(ReprFlags::RANDOMIZE_LAYOUT)
+    }
+
     /// Returns `true` if this `#[repr()]` should inhibit union ABI optimisations.
     pub fn inhibit_union_abi_opt(&self) -> bool {
         self.c()
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 8110afe75fa..8ecb7a031ad 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -1246,6 +1246,8 @@ options! {
         "enable queries of the dependency graph for regression testing (default: no)"),
     query_stats: bool = (false, parse_bool, [UNTRACKED],
         "print some statistics about the query system (default: no)"),
+    randomize_layout: bool = (false, parse_bool, [TRACKED],
+        "randomize the layout of types (default: no)"),
     relax_elf_relocations: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "whether ELF relocations can be relaxed"),
     relro_level: Option<RelroLevel> = (None, parse_relro_level, [TRACKED],
diff --git a/src/tools/tidy/src/deps.rs b/src/tools/tidy/src/deps.rs
index 60703384e9e..c1719a9ffe8 100644
--- a/src/tools/tidy/src/deps.rs
+++ b/src/tools/tidy/src/deps.rs
@@ -167,6 +167,7 @@ const PERMITTED_DEPENDENCIES: &[&str] = &[
     "rand_hc",
     "rand_pcg",
     "rand_xorshift",
+    "rand_xoshiro",
     "redox_syscall",
     "regex",
     "regex-automata",