Fix race condition when allocating source files in SourceMap

Author: Zoxc

src/librustc_span/lib.rs

@@ -1075,7 +1075,7 @@ impl SourceFile {
         unmapped_path: FileName,
         mut src: String,
         start_pos: BytePos,
-    ) -> Result<SourceFile, OffsetOverflowError> {
+    ) -> Self {
         let normalized_pos = normalize_src(&mut src, start_pos);
 
         let src_hash = {
@@ -1089,14 +1089,12 @@ impl SourceFile {
             hasher.finish::<u128>()
         };
         let end_pos = start_pos.to_usize() + src.len();
-        if end_pos > u32::max_value() as usize {
-            return Err(OffsetOverflowError);
-        }
+        assert!(end_pos <= u32::max_value() as usize);
 
         let (lines, multibyte_chars, non_narrow_chars) =
             analyze_source_file::analyze_source_file(&src[..], start_pos);
 
-        Ok(SourceFile {
+        SourceFile {
             name,
             name_was_remapped,
             unmapped_path: Some(unmapped_path),
@@ -1111,7 +1109,7 @@ impl SourceFile {
             non_narrow_chars,
             normalized_pos,
             name_hash,
-        })
+        }
     }
 
     /// Returns the `BytePos` of the beginning of the current line.

src/librustc_span/source_map.rs

@@ -12,10 +12,12 @@ pub use crate::*;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::StableHasher;
-use rustc_data_structures::sync::{Lock, LockGuard, Lrc, MappedLockGuard};
+use rustc_data_structures::sync::{AtomicU32, Lock, LockGuard, Lrc, MappedLockGuard};
 use std::cmp;
+use std::convert::TryFrom;
 use std::hash::Hash;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::Ordering;
 
 use log::debug;
 use std::env;
@@ -131,6 +133,9 @@ pub(super) struct SourceMapFiles {
 }
 
 pub struct SourceMap {
+    /// The address space below this value is currently used by the files in the source map.
+    used_address_space: AtomicU32,
+
     files: Lock<SourceMapFiles>,
     file_loader: Box<dyn FileLoader + Sync + Send>,
     // This is used to apply the file path remapping as specified via
@@ -140,14 +145,24 @@ pub struct SourceMap {
 
 impl SourceMap {
     pub fn new(path_mapping: FilePathMapping) -> SourceMap {
-        SourceMap { files: Default::default(), file_loader: Box::new(RealFileLoader), path_mapping }
+        SourceMap {
+            used_address_space: AtomicU32::new(0),
+            files: Default::default(),
+            file_loader: Box::new(RealFileLoader),
+            path_mapping,
+        }
     }
 
     pub fn with_file_loader(
         file_loader: Box<dyn FileLoader + Sync + Send>,
         path_mapping: FilePathMapping,
     ) -> SourceMap {
-        SourceMap { files: Default::default(), file_loader, path_mapping }
+        SourceMap {
+            used_address_space: AtomicU32::new(0),
+            files: Default::default(),
+            file_loader,
+            path_mapping,
+        }
     }
 
     pub fn path_mapping(&self) -> &FilePathMapping {
@@ -194,12 +209,25 @@ impl SourceMap {
         self.files.borrow().stable_id_to_source_file.get(&stable_id).map(|sf| sf.clone())
     }
 
-    fn next_start_pos(&self) -> usize {
-        match self.files.borrow().source_files.last() {
-            None => 0,
-            // Add one so there is some space between files. This lets us distinguish
-            // positions in the `SourceMap`, even in the presence of zero-length files.
-            Some(last) => last.end_pos.to_usize() + 1,
+    fn allocate_address_space(&self, size: usize) -> Result<usize, OffsetOverflowError> {
+        let size = u32::try_from(size).map_err(|_| OffsetOverflowError)?;
+
+        loop {
+            let current = self.used_address_space.load(Ordering::Relaxed);
+            let next = current
+                .checked_add(size)
+                // Add one so there is some space between files. This lets us distinguish
+                // positions in the `SourceMap`, even in the presence of zero-length files.
+                .and_then(|next| next.checked_add(1))
+                .ok_or(OffsetOverflowError)?;
+
+            if self
+                .used_address_space
+                .compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed)
+                .is_ok()
+            {
+                return Ok(usize::try_from(current).unwrap());
+            }
         }
     }
 
@@ -218,8 +246,6 @@ impl SourceMap {
         filename: FileName,
         src: String,
     ) -> Result<Lrc<SourceFile>, OffsetOverflowError> {
-        let start_pos = self.next_start_pos();
-
         // The path is used to determine the directory for loading submodules and
         // include files, so it must be before remapping.
         // Note that filename may not be a valid path, eg it may be `<anon>` etc,
@@ -241,13 +267,15 @@ impl SourceMap {
         let lrc_sf = match self.source_file_by_stable_id(file_id) {
             Some(lrc_sf) => lrc_sf,
             None => {
+                let start_pos = self.allocate_address_space(src.len())?;
+
                 let source_file = Lrc::new(SourceFile::new(
                     filename,
                     was_remapped,
                     unmapped_path,
                     src,
                     Pos::from_usize(start_pos),
-                )?);
+                ));
 
                 let mut files = self.files.borrow_mut();
 
@@ -277,7 +305,9 @@ impl SourceMap {
         mut file_local_non_narrow_chars: Vec<NonNarrowChar>,
         mut file_local_normalized_pos: Vec<NormalizedPos>,
     ) -> Lrc<SourceFile> {
-        let start_pos = self.next_start_pos();
+        let start_pos = self
+            .allocate_address_space(source_len)
+            .expect("not enough address space for imported source file");
 
         let end_pos = Pos::from_usize(start_pos + source_len);
         let start_pos = Pos::from_usize(start_pos);