escape symbol names in global asm

Author: usamoi

compiler/rustc_codegen_llvm/src/asm.rs

@@ -414,7 +414,7 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
                                 llvm::LLVMRustGetMangledName(llval, s);
                             })
                             .expect("symbol is not valid UTF-8");
-                            template_str.push_str(&symbol);
+                            template_str.push_str(&escape_symbol_name(self, symbol));
                         }
                         GlobalAsmOperandRef::SymStatic { def_id } => {
                             let llval = self
@@ -428,7 +428,7 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
                                 llvm::LLVMRustGetMangledName(llval, s);
                             })
                             .expect("symbol is not valid UTF-8");
-                            template_str.push_str(&symbol);
+                            template_str.push_str(&escape_symbol_name(self, symbol));
                         }
                     }
                 }
@@ -1390,3 +1390,36 @@ fn llvm_fixup_output_type<'ll, 'tcx>(
         _ => layout.llvm_type(cx),
     }
 }
+
+fn escape_symbol_name<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, symbol: String) -> String {
+    use rustc_target::spec::{Arch, BinaryFormat};
+    if !symbol.is_empty()
+        && symbol.chars().all(|c| matches!(c, '0'..='9' | 'A'..='Z' | 'a'..='z' | '_' | '$' | '.'))
+    {
+        return symbol;
+    }
+    if cx.tcx.sess.target.binary_format == BinaryFormat::Xcoff {
+        cx.tcx.sess.dcx().fatal(format!(
+            "symbol escaping is not supported for the binary format {}",
+            cx.tcx.sess.target.binary_format
+        ));
+    }
+    if matches!(cx.tcx.sess.target.arch, Arch::Nvptx64) {
+        cx.tcx.sess.dcx().fatal(format!(
+            "symbol escaping is not supported for the architecture {}",
+            cx.tcx.sess.target.arch
+        ));
+    }
+    let mut escaped_symbol = String::new();
+    escaped_symbol.push('\"');
+    for c in symbol.chars() {
+        match c {
+            '\n' => escaped_symbol.push_str("\\\n"),
+            '"' => escaped_symbol.push_str("\\\""),
+            '\\' => escaped_symbol.push_str("\\\\"),
+            c => escaped_symbol.push(c),
+        }
+    }
+    escaped_symbol.push('\"');
+    escaped_symbol
+}

tests/ui/asm/x86_64/global_asm_escape.rs

@@ -0,0 +1,16 @@
+//@ run-pass
+//@ only-x86_64-unknown-linux-gnu
+
+// https://github.com/rust-lang/rust/issues/151950
+
+unsafe extern "C" {
+    #[link_name = "exit@GLIBC_2.2.5"]
+    safe fn exit(status: i32) -> !;
+    safe fn my_exit(status: i32) -> !;
+}
+
+core::arch::global_asm!(".global my_exit", "my_exit:", "jmp {}", sym exit);
+
+fn main() {
+    my_exit(0);
+}

tests/ui/asm/x86_64/naked_asm_escape.rs

@@ -0,0 +1,31 @@
+//@ build-fail
+//@ only-x86_64-unknown-linux-gnu
+//@ dont-check-compiler-stderr
+//@ dont-check-compiler-stdout
+
+// https://github.com/rust-lang/rust/issues/151950
+
+unsafe extern "C" {
+    #[link_name = "memset]; mov eax, 1; #"]
+    unsafe fn inject();
+}
+
+#[unsafe(export_name = "memset]; mov eax, 1; #")]
+extern "C" fn inject_() {}
+
+#[unsafe(naked)]
+extern "C" fn print_0() -> usize {
+    core::arch::naked_asm!("lea rax, [{}]", "ret", sym inject)
+}
+
+#[unsafe(naked)]
+extern "C" fn print_1() -> usize {
+    core::arch::naked_asm!("lea rax, [{}]", "ret", sym inject_)
+}
+
+fn main() {
+    dbg!(print_0());
+    dbg!(print_1());
+}
+
+//~? ERROR linking