summaryrefslogtreecommitdiffstats
path: root/third_party/rust/embed-manifest/src/embed/mod.rs
blob: b054878ac6756b8a166c02f7a338053f659ffd6e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
use std::env;
use std::fs::{self, File};
use std::io::{self, stdout, BufWriter, Cursor, Write};
use std::path::{Path, PathBuf};

use crate::manifest::ManifestBuilder;

use self::coff::{resource_data_entry, resource_directory_id_entry, resource_directory_table, CoffWriter, MachineType};
use self::error::Error;

mod coff;
pub mod error;

#[cfg(test)]
mod test;

/// Embeds the manifest described by `manifest` by converting it to XML,
/// then saving it to a file and passing the correct options to the linker
/// on MSVC targets, or by building a static library and instructing Cargo
/// to link the executable against it on GNU targets.
pub fn embed_manifest(manifest: ManifestBuilder) -> Result<(), Error> {
    let out_dir = get_out_dir()?;
    let target = get_target()?;
    if matches!(target.os, TargetOs::WindowsMsvc) {
        let manifest_file = out_dir.join("manifest.xml");
        write!(BufWriter::new(File::create(&manifest_file)?), "{}", manifest)?;
        link_manifest_msvc(&manifest_file, &mut stdout().lock())
    } else {
        let manifest_data = manifest.to_string();
        link_manifest_gnu(manifest_data.as_bytes(), &out_dir, target.arch, &mut stdout().lock())
    }
}

/// Directly embeds the manifest in the provided `file` by passing the correct
/// options to the linker on MSVC targets, or by building a static library
/// and instructing Cargo to link the executable against it on GNU targets.
pub fn embed_manifest_file<P: AsRef<Path>>(file: P) -> Result<(), io::Error> {
    let out_dir = get_out_dir()?;
    let target = get_target()?;
    if matches!(target.os, TargetOs::WindowsMsvc) {
        Ok(link_manifest_msvc(file.as_ref(), &mut stdout().lock())?)
    } else {
        let manifest = fs::read(file.as_ref())?;
        Ok(link_manifest_gnu(&manifest, &out_dir, target.arch, &mut stdout().lock())?)
    }
}

fn get_out_dir() -> Result<PathBuf, io::Error> {
    match env::var_os("OUT_DIR") {
        Some(dir) => Ok(PathBuf::from(dir)),
        None => env::current_dir(),
    }
}

enum TargetOs {
    WindowsGnu,
    WindowsMsvc,
}

struct Target {
    arch: MachineType,
    os: TargetOs,
}

fn get_target() -> Result<Target, Error> {
    match env::var("TARGET") {
        Ok(target) => parse_target(&target),
        _ => Err(Error::unknown_target()),
    }
}

fn parse_target(target: &str) -> Result<Target, Error> {
    let mut iter = target.splitn(3, '-');
    let arch = match iter.next() {
        Some("i686") => MachineType::I386,
        Some("aarch64") => MachineType::Aarch64,
        Some("x86_64") => MachineType::X86_64,
        _ => return Err(Error::unknown_target()),
    };
    if iter.next() != Some("pc") {
        return Err(Error::unknown_target());
    }
    let os = match iter.next() {
        Some("windows-gnu") => TargetOs::WindowsGnu,
        Some("windows-gnullvm") => TargetOs::WindowsGnu,
        Some("windows-msvc") => TargetOs::WindowsMsvc,
        _ => return Err(Error::unknown_target()),
    };
    Ok(Target { arch, os })
}

fn link_manifest_msvc<W: Write>(manifest_path: &Path, out: &mut W) -> Result<(), Error> {
    writeln!(out, "cargo:rustc-link-arg-bins=/MANIFEST:EMBED")?;
    writeln!(
        out,
        "cargo:rustc-link-arg-bins=/MANIFESTINPUT:{}",
        manifest_path.canonicalize()?.display()
    )?;
    writeln!(out, "cargo:rustc-link-arg-bins=/MANIFESTUAC:NO")?;
    Ok(())
}

fn link_manifest_gnu<W: Write>(manifest: &[u8], out_dir: &Path, arch: MachineType, out: &mut W) -> Result<(), Error> {
    // Generate a COFF object file containing the manifest in a .rsrc section.
    let object_data = create_object_file(manifest, arch)?;
    let path = out_dir.join("embed-manifest.o");
    fs::create_dir_all(out_dir)?;
    fs::write(&path, object_data)?;

    // Link the manifest with the executable.
    writeln!(out, "cargo:rustc-link-arg-bins={}", path.display())?;
    Ok(())
}

fn create_object_file(manifest: &[u8], arch: MachineType) -> Result<Vec<u8>, io::Error> {
    // Start object file with .rsrc section.
    let mut obj = CoffWriter::new(Cursor::new(Vec::with_capacity(4096)), arch)?;

    // Create resource directories for type ID 24, name ID 1, language ID 1033.
    obj.add_data(&resource_directory_table(1))?;
    obj.add_data(&resource_directory_id_entry(24, 24, true))?;
    obj.add_data(&resource_directory_table(1))?;
    obj.add_data(&resource_directory_id_entry(1, 48, true))?;
    obj.add_data(&resource_directory_table(1))?;
    obj.add_data(&resource_directory_id_entry(1033, 72, false))?;

    // Add resource data entry with relocated address.
    let address = obj.add_data(&resource_data_entry(88, manifest.len() as u32))?;

    // Add the manifest data.
    obj.add_data(manifest)?;
    obj.align_to(8)?;

    // Write manifest data relocation at the end of the section.
    obj.add_relocation(address)?;

    // Finish writing file and return the populated object data.
    Ok(obj.finish()?.into_inner())
}