From 3943a32212a424c6dc0c0c97e3f15b0f627e937c Mon Sep 17 00:00:00 2001
From: Tom Kirchner <tjk@amazon.com>
Date: Wed, 17 Feb 2021 11:15:27 -0800
Subject: [PATCH] models: fix TOCTOU in build.rs by safely swapping links into
 place

---
 sources/Cargo.lock        |  1 +
 sources/models/Cargo.toml |  1 +
 sources/models/build.rs   | 50 ++++++++++++++++++++++++++-------------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/sources/Cargo.lock b/sources/Cargo.lock
index 90261674de4..589cc34c1dd 100644
--- a/sources/Cargo.lock
+++ b/sources/Cargo.lock
@@ -1909,6 +1909,7 @@ dependencies = [
  "cargo-readme",
  "lazy_static",
  "model-derive",
+ "rand 0.8.1",
  "regex",
  "semver 0.11.0",
  "serde",
diff --git a/sources/models/Cargo.toml b/sources/models/Cargo.toml
index e048ccc963e..28de3ba22ac 100644
--- a/sources/models/Cargo.toml
+++ b/sources/models/Cargo.toml
@@ -24,6 +24,7 @@ url = "2.1"
 
 [build-dependencies]
 cargo-readme = "3.1"
+rand = "0.8"
 
 [lib]
 # We're picking the current *model* with build.rs, so users shouldn't think
diff --git a/sources/models/build.rs b/sources/models/build.rs
index 6e8408f43a6..9c9557ba9e4 100644
--- a/sources/models/build.rs
+++ b/sources/models/build.rs
@@ -4,6 +4,7 @@
 //
 // See README.md to understand the symlink setup.
 
+use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use std::env;
 use std::fs::{self, File};
 use std::io::{self, Write};
@@ -32,21 +33,6 @@ fn main() {
     link_current_variant();
 }
 
-fn symlink_force<P1, P2>(target: P1, link: P2) -> io::Result<()>
-where
-    P1: AsRef<Path>,
-    P2: AsRef<Path>,
-{
-    // Remove link if it already exists
-    if let Err(e) = fs::remove_file(&link) {
-        if e.kind() != io::ErrorKind::NotFound {
-            return Err(e);
-        }
-    }
-    // Link to requested target
-    symlink(&target, &link)
-}
-
 fn link_current_variant() {
     // The VARIANT variable is originally BUILDSYS_VARIANT, set in the top-level Makefile.toml,
     // and is passed through as VARIANT by the top-level Dockerfile.  It represents which OS
@@ -67,7 +53,7 @@ fn link_current_variant() {
     }
 
     // Create the symlink for the following `cargo build` to use for its source code
-    symlink_force(&variant_target, VARIANT_LINK).unwrap_or_else(|e| {
+    symlink_safe(&variant_target, VARIANT_LINK).unwrap_or_else(|e| {
         eprintln!("Failed to create symlink at '{}' pointing to '{}' - we need this to support different API models for different variants.  Error: {}", VARIANT_LINK, variant_target, e);
         process::exit(1);
     });
@@ -75,7 +61,7 @@ fn link_current_variant() {
     // Also create the link for mod.rs so Rust can import source from the "current" link
     // created above.
     let mod_target = "../variant_mod.rs";
-    symlink_force(&mod_target, MOD_LINK).unwrap_or_else(|e| {
+    symlink_safe(&mod_target, MOD_LINK).unwrap_or_else(|e| {
         eprintln!("Failed to create symlink at '{}' pointing to '{}' - we need this to build a Rust module structure through the `current` link.  Error: {}", MOD_LINK, mod_target, e);
         process::exit(1);
     });
@@ -106,3 +92,33 @@ fn generate_readme() {
     let mut readme = File::create("README.md").unwrap();
     readme.write_all(content.as_bytes()).unwrap();
 }
+
+// Creates the requested symlink through an atomic swap, so it doesn't matter if the link path
+// already exists or not; like --force but fewer worries about reentrancy and retries.
+fn symlink_safe<P1, P2>(target: P1, link: P2) -> io::Result<()>
+where
+    P1: AsRef<Path>,
+    P2: AsRef<Path>,
+{
+    // Create the link at a temporary path.
+    let temp_link = link.as_ref().with_file_name(format!(".{}", rando()));
+    symlink(&target, &temp_link)?;
+
+    // Swap the temporary link into the real location
+    if let Err(e) = fs::rename(&temp_link, &link) {
+        // If we couldn't, for whatever reason, clean up the temporary path and return the error.
+        let _ = fs::remove_file(&temp_link);
+        return Err(e);
+    }
+
+    Ok(())
+}
+
+// Generates a random ID, affectionately known as a 'rando'.
+fn rando() -> String {
+    thread_rng()
+        .sample_iter(&Alphanumeric)
+        .take(16)
+        .map(char::from)
+        .collect()
+}