xref: /aosp_15_r20/external/bazelbuild-rules_rust/tools/runfiles/runfiles.rs (revision d4726bddaa87cc4778e7472feed243fa4b6c267f)
1*d4726bddSHONG Yifan //! Runfiles lookup library for Bazel-built Rust binaries and tests.
2*d4726bddSHONG Yifan //!
3*d4726bddSHONG Yifan //! USAGE:
4*d4726bddSHONG Yifan //!
5*d4726bddSHONG Yifan //! 1.  Depend on this runfiles library from your build rule:
6*d4726bddSHONG Yifan //!     ```python
7*d4726bddSHONG Yifan //!       rust_binary(
8*d4726bddSHONG Yifan //!           name = "my_binary",
9*d4726bddSHONG Yifan //!           ...
10*d4726bddSHONG Yifan //!           data = ["//path/to/my/data.txt"],
11*d4726bddSHONG Yifan //!           deps = ["@rules_rust//tools/runfiles"],
12*d4726bddSHONG Yifan //!       )
13*d4726bddSHONG Yifan //!     ```
14*d4726bddSHONG Yifan //!
15*d4726bddSHONG Yifan //! 2.  Import the runfiles library.
16*d4726bddSHONG Yifan //!     ```ignore
17*d4726bddSHONG Yifan //!     extern crate runfiles;
18*d4726bddSHONG Yifan //!
19*d4726bddSHONG Yifan //!     use runfiles::Runfiles;
20*d4726bddSHONG Yifan //!     ```
21*d4726bddSHONG Yifan //!
22*d4726bddSHONG Yifan //! 3.  Create a Runfiles object and use `rlocation!`` to look up runfile paths:
23*d4726bddSHONG Yifan //!     ```ignore -- This doesn't work under rust_doc_test because argv[0] is not what we expect.
24*d4726bddSHONG Yifan //!
25*d4726bddSHONG Yifan //!     use runfiles::{Runfiles, rlocation};
26*d4726bddSHONG Yifan //!
27*d4726bddSHONG Yifan //!     let r = Runfiles::create().unwrap();
28*d4726bddSHONG Yifan //!     let path = rlocation!(r, "my_workspace/path/to/my/data.txt");
29*d4726bddSHONG Yifan //!
30*d4726bddSHONG Yifan //!     let f = File::open(path).unwrap();
31*d4726bddSHONG Yifan //!     // ...
32*d4726bddSHONG Yifan //!     ```
33*d4726bddSHONG Yifan 
34*d4726bddSHONG Yifan use std::collections::HashMap;
35*d4726bddSHONG Yifan use std::env;
36*d4726bddSHONG Yifan use std::fs;
37*d4726bddSHONG Yifan use std::io;
38*d4726bddSHONG Yifan use std::path::Path;
39*d4726bddSHONG Yifan use std::path::PathBuf;
40*d4726bddSHONG Yifan 
41*d4726bddSHONG Yifan const RUNFILES_DIR_ENV_VAR: &str = "RUNFILES_DIR";
42*d4726bddSHONG Yifan const MANIFEST_FILE_ENV_VAR: &str = "RUNFILES_MANIFEST_FILE";
43*d4726bddSHONG Yifan const TEST_SRCDIR_ENV_VAR: &str = "TEST_SRCDIR";
44*d4726bddSHONG Yifan 
45*d4726bddSHONG Yifan #[macro_export]
46*d4726bddSHONG Yifan macro_rules! rlocation {
47*d4726bddSHONG Yifan     ($r:expr, $path:expr) => {
48*d4726bddSHONG Yifan         $r.rlocation_from($path, env!("REPOSITORY_NAME"))
49*d4726bddSHONG Yifan     };
50*d4726bddSHONG Yifan }
51*d4726bddSHONG Yifan 
52*d4726bddSHONG Yifan #[derive(Debug)]
53*d4726bddSHONG Yifan enum Mode {
54*d4726bddSHONG Yifan     DirectoryBased(PathBuf),
55*d4726bddSHONG Yifan     ManifestBased(HashMap<PathBuf, PathBuf>),
56*d4726bddSHONG Yifan }
57*d4726bddSHONG Yifan 
58*d4726bddSHONG Yifan type RepoMappingKey = (String, String);
59*d4726bddSHONG Yifan type RepoMapping = HashMap<RepoMappingKey, String>;
60*d4726bddSHONG Yifan 
61*d4726bddSHONG Yifan #[derive(Debug)]
62*d4726bddSHONG Yifan pub struct Runfiles {
63*d4726bddSHONG Yifan     mode: Mode,
64*d4726bddSHONG Yifan     repo_mapping: RepoMapping,
65*d4726bddSHONG Yifan }
66*d4726bddSHONG Yifan 
67*d4726bddSHONG Yifan impl Runfiles {
68*d4726bddSHONG Yifan     /// Creates a manifest based Runfiles object when
69*d4726bddSHONG Yifan     /// RUNFILES_MANIFEST_FILE environment variable is present,
70*d4726bddSHONG Yifan     /// or a directory based Runfiles object otherwise.
create() -> io::Result<Self>71*d4726bddSHONG Yifan     pub fn create() -> io::Result<Self> {
72*d4726bddSHONG Yifan         let mode = if let Some(manifest_file) = std::env::var_os(MANIFEST_FILE_ENV_VAR) {
73*d4726bddSHONG Yifan             Self::create_manifest_based(Path::new(&manifest_file))?
74*d4726bddSHONG Yifan         } else {
75*d4726bddSHONG Yifan             Mode::DirectoryBased(find_runfiles_dir()?)
76*d4726bddSHONG Yifan         };
77*d4726bddSHONG Yifan 
78*d4726bddSHONG Yifan         let repo_mapping = parse_repo_mapping(raw_rlocation(&mode, "_repo_mapping"))
79*d4726bddSHONG Yifan             .unwrap_or_else(|_| {
80*d4726bddSHONG Yifan                 println!("No repo mapping found!");
81*d4726bddSHONG Yifan                 RepoMapping::new()
82*d4726bddSHONG Yifan             });
83*d4726bddSHONG Yifan 
84*d4726bddSHONG Yifan         Ok(Runfiles { mode, repo_mapping })
85*d4726bddSHONG Yifan     }
86*d4726bddSHONG Yifan 
create_manifest_based(manifest_path: &Path) -> io::Result<Mode>87*d4726bddSHONG Yifan     fn create_manifest_based(manifest_path: &Path) -> io::Result<Mode> {
88*d4726bddSHONG Yifan         let manifest_content = std::fs::read_to_string(manifest_path)?;
89*d4726bddSHONG Yifan         let path_mapping = manifest_content
90*d4726bddSHONG Yifan             .lines()
91*d4726bddSHONG Yifan             .map(|line| {
92*d4726bddSHONG Yifan                 let pair = line
93*d4726bddSHONG Yifan                     .split_once(' ')
94*d4726bddSHONG Yifan                     .expect("manifest file contained unexpected content");
95*d4726bddSHONG Yifan                 (pair.0.into(), pair.1.into())
96*d4726bddSHONG Yifan             })
97*d4726bddSHONG Yifan             .collect::<HashMap<_, _>>();
98*d4726bddSHONG Yifan         Ok(Mode::ManifestBased(path_mapping))
99*d4726bddSHONG Yifan     }
100*d4726bddSHONG Yifan 
101*d4726bddSHONG Yifan     /// Returns the runtime path of a runfile.
102*d4726bddSHONG Yifan     ///
103*d4726bddSHONG Yifan     /// Runfiles are data-dependencies of Bazel-built binaries and tests.
104*d4726bddSHONG Yifan     /// The returned path may not be valid. The caller should check the path's
105*d4726bddSHONG Yifan     /// validity and that the path exists.
106*d4726bddSHONG Yifan     /// @deprecated - this is not bzlmod-aware. Prefer the `rlocation!` macro or `rlocation_from`
rlocation(&self, path: impl AsRef<Path>) -> PathBuf107*d4726bddSHONG Yifan     pub fn rlocation(&self, path: impl AsRef<Path>) -> PathBuf {
108*d4726bddSHONG Yifan         let path = path.as_ref();
109*d4726bddSHONG Yifan         if path.is_absolute() {
110*d4726bddSHONG Yifan             return path.to_path_buf();
111*d4726bddSHONG Yifan         }
112*d4726bddSHONG Yifan         raw_rlocation(&self.mode, path)
113*d4726bddSHONG Yifan     }
114*d4726bddSHONG Yifan 
115*d4726bddSHONG Yifan     /// Returns the runtime path of a runfile.
116*d4726bddSHONG Yifan     ///
117*d4726bddSHONG Yifan     /// Runfiles are data-dependencies of Bazel-built binaries and tests.
118*d4726bddSHONG Yifan     /// The returned path may not be valid. The caller should check the path's
119*d4726bddSHONG Yifan     /// validity and that the path exists.
120*d4726bddSHONG Yifan     ///
121*d4726bddSHONG Yifan     /// Typically this should be used via the `rlocation!` macro to properly set source_repo.
rlocation_from(&self, path: impl AsRef<Path>, source_repo: &str) -> PathBuf122*d4726bddSHONG Yifan     pub fn rlocation_from(&self, path: impl AsRef<Path>, source_repo: &str) -> PathBuf {
123*d4726bddSHONG Yifan         let path = path.as_ref();
124*d4726bddSHONG Yifan         if path.is_absolute() {
125*d4726bddSHONG Yifan             return path.to_path_buf();
126*d4726bddSHONG Yifan         }
127*d4726bddSHONG Yifan 
128*d4726bddSHONG Yifan         let path_str = path.to_str().expect("Should be valid UTF8");
129*d4726bddSHONG Yifan         let (repo_alias, repo_path): (&str, Option<&str>) = match path_str.split_once('/') {
130*d4726bddSHONG Yifan             Some((name, alias)) => (name, Some(alias)),
131*d4726bddSHONG Yifan             None => (path_str, None),
132*d4726bddSHONG Yifan         };
133*d4726bddSHONG Yifan         let key: (String, String) = (source_repo.into(), repo_alias.into());
134*d4726bddSHONG Yifan         if let Some(target_repo_directory) = self.repo_mapping.get(&key) {
135*d4726bddSHONG Yifan             match repo_path {
136*d4726bddSHONG Yifan                 Some(repo_path) => {
137*d4726bddSHONG Yifan                     raw_rlocation(&self.mode, format!("{target_repo_directory}/{repo_path}"))
138*d4726bddSHONG Yifan                 }
139*d4726bddSHONG Yifan                 None => raw_rlocation(&self.mode, target_repo_directory),
140*d4726bddSHONG Yifan             }
141*d4726bddSHONG Yifan         } else {
142*d4726bddSHONG Yifan             raw_rlocation(&self.mode, path)
143*d4726bddSHONG Yifan         }
144*d4726bddSHONG Yifan     }
145*d4726bddSHONG Yifan }
146*d4726bddSHONG Yifan 
raw_rlocation(mode: &Mode, path: impl AsRef<Path>) -> PathBuf147*d4726bddSHONG Yifan fn raw_rlocation(mode: &Mode, path: impl AsRef<Path>) -> PathBuf {
148*d4726bddSHONG Yifan     let path = path.as_ref();
149*d4726bddSHONG Yifan     match mode {
150*d4726bddSHONG Yifan         Mode::DirectoryBased(runfiles_dir) => runfiles_dir.join(path),
151*d4726bddSHONG Yifan         Mode::ManifestBased(path_mapping) => path_mapping
152*d4726bddSHONG Yifan             .get(path)
153*d4726bddSHONG Yifan             .unwrap_or_else(|| panic!("Path {} not found among runfiles.", path.to_string_lossy()))
154*d4726bddSHONG Yifan             .clone(),
155*d4726bddSHONG Yifan     }
156*d4726bddSHONG Yifan }
157*d4726bddSHONG Yifan 
parse_repo_mapping(path: PathBuf) -> io::Result<RepoMapping>158*d4726bddSHONG Yifan fn parse_repo_mapping(path: PathBuf) -> io::Result<RepoMapping> {
159*d4726bddSHONG Yifan     let mut repo_mapping = RepoMapping::new();
160*d4726bddSHONG Yifan 
161*d4726bddSHONG Yifan     for line in std::fs::read_to_string(path)?.lines() {
162*d4726bddSHONG Yifan         let parts: Vec<&str> = line.splitn(3, ',').collect();
163*d4726bddSHONG Yifan         if parts.len() < 3 {
164*d4726bddSHONG Yifan             return Err(make_io_error("Malformed repo_mapping file"));
165*d4726bddSHONG Yifan         }
166*d4726bddSHONG Yifan         repo_mapping.insert((parts[0].into(), parts[1].into()), parts[2].into());
167*d4726bddSHONG Yifan     }
168*d4726bddSHONG Yifan 
169*d4726bddSHONG Yifan     Ok(repo_mapping)
170*d4726bddSHONG Yifan }
171*d4726bddSHONG Yifan 
172*d4726bddSHONG Yifan /// Returns the .runfiles directory for the currently executing binary.
find_runfiles_dir() -> io::Result<PathBuf>173*d4726bddSHONG Yifan pub fn find_runfiles_dir() -> io::Result<PathBuf> {
174*d4726bddSHONG Yifan     assert!(std::env::var_os(MANIFEST_FILE_ENV_VAR).is_none());
175*d4726bddSHONG Yifan 
176*d4726bddSHONG Yifan     // If bazel told us about the runfiles dir, use that without looking further.
177*d4726bddSHONG Yifan     if let Some(runfiles_dir) = std::env::var_os(RUNFILES_DIR_ENV_VAR).map(PathBuf::from) {
178*d4726bddSHONG Yifan         if runfiles_dir.is_dir() {
179*d4726bddSHONG Yifan             return Ok(runfiles_dir);
180*d4726bddSHONG Yifan         }
181*d4726bddSHONG Yifan     }
182*d4726bddSHONG Yifan     if let Some(test_srcdir) = std::env::var_os(TEST_SRCDIR_ENV_VAR).map(PathBuf::from) {
183*d4726bddSHONG Yifan         if test_srcdir.is_dir() {
184*d4726bddSHONG Yifan             return Ok(test_srcdir);
185*d4726bddSHONG Yifan         }
186*d4726bddSHONG Yifan     }
187*d4726bddSHONG Yifan 
188*d4726bddSHONG Yifan     // Consume the first argument (argv[0])
189*d4726bddSHONG Yifan     let exec_path = std::env::args().next().expect("arg 0 was not set");
190*d4726bddSHONG Yifan 
191*d4726bddSHONG Yifan     let mut binary_path = PathBuf::from(&exec_path);
192*d4726bddSHONG Yifan     loop {
193*d4726bddSHONG Yifan         // Check for our neighboring $binary.runfiles directory.
194*d4726bddSHONG Yifan         let mut runfiles_name = binary_path.file_name().unwrap().to_owned();
195*d4726bddSHONG Yifan         runfiles_name.push(".runfiles");
196*d4726bddSHONG Yifan 
197*d4726bddSHONG Yifan         let runfiles_path = binary_path.with_file_name(&runfiles_name);
198*d4726bddSHONG Yifan         if runfiles_path.is_dir() {
199*d4726bddSHONG Yifan             return Ok(runfiles_path);
200*d4726bddSHONG Yifan         }
201*d4726bddSHONG Yifan 
202*d4726bddSHONG Yifan         // Check if we're already under a *.runfiles directory.
203*d4726bddSHONG Yifan         {
204*d4726bddSHONG Yifan             // TODO: 1.28 adds Path::ancestors() which is a little simpler.
205*d4726bddSHONG Yifan             let mut next = binary_path.parent();
206*d4726bddSHONG Yifan             while let Some(ancestor) = next {
207*d4726bddSHONG Yifan                 if ancestor
208*d4726bddSHONG Yifan                     .file_name()
209*d4726bddSHONG Yifan                     .map_or(false, |f| f.to_string_lossy().ends_with(".runfiles"))
210*d4726bddSHONG Yifan                 {
211*d4726bddSHONG Yifan                     return Ok(ancestor.to_path_buf());
212*d4726bddSHONG Yifan                 }
213*d4726bddSHONG Yifan                 next = ancestor.parent();
214*d4726bddSHONG Yifan             }
215*d4726bddSHONG Yifan         }
216*d4726bddSHONG Yifan 
217*d4726bddSHONG Yifan         if !fs::symlink_metadata(&binary_path)?.file_type().is_symlink() {
218*d4726bddSHONG Yifan             break;
219*d4726bddSHONG Yifan         }
220*d4726bddSHONG Yifan         // Follow symlinks and keep looking.
221*d4726bddSHONG Yifan         let link_target = binary_path.read_link()?;
222*d4726bddSHONG Yifan         binary_path = if link_target.is_absolute() {
223*d4726bddSHONG Yifan             link_target
224*d4726bddSHONG Yifan         } else {
225*d4726bddSHONG Yifan             let link_dir = binary_path.parent().unwrap();
226*d4726bddSHONG Yifan             env::current_dir()?.join(link_dir).join(link_target)
227*d4726bddSHONG Yifan         }
228*d4726bddSHONG Yifan     }
229*d4726bddSHONG Yifan 
230*d4726bddSHONG Yifan     Err(make_io_error("failed to find .runfiles directory"))
231*d4726bddSHONG Yifan }
232*d4726bddSHONG Yifan 
make_io_error(msg: &str) -> io::Error233*d4726bddSHONG Yifan fn make_io_error(msg: &str) -> io::Error {
234*d4726bddSHONG Yifan     io::Error::new(io::ErrorKind::Other, msg)
235*d4726bddSHONG Yifan }
236*d4726bddSHONG Yifan 
237*d4726bddSHONG Yifan #[cfg(test)]
238*d4726bddSHONG Yifan mod test {
239*d4726bddSHONG Yifan     use super::*;
240*d4726bddSHONG Yifan 
241*d4726bddSHONG Yifan     use std::fs::File;
242*d4726bddSHONG Yifan     use std::io::prelude::*;
243*d4726bddSHONG Yifan 
244*d4726bddSHONG Yifan     #[test]
test_can_read_data_from_runfiles()245*d4726bddSHONG Yifan     fn test_can_read_data_from_runfiles() {
246*d4726bddSHONG Yifan         // We want to run multiple test cases with different environment variables set. Since
247*d4726bddSHONG Yifan         // environment variables are global state, we need to ensure the test cases do not run
248*d4726bddSHONG Yifan         // concurrently. Rust runs tests in parallel and does not provide an easy way to synchronise
249*d4726bddSHONG Yifan         // them, so we run all test cases in the same #[test] function.
250*d4726bddSHONG Yifan 
251*d4726bddSHONG Yifan         let test_srcdir =
252*d4726bddSHONG Yifan             env::var_os(TEST_SRCDIR_ENV_VAR).expect("bazel did not provide TEST_SRCDIR");
253*d4726bddSHONG Yifan         let runfiles_dir =
254*d4726bddSHONG Yifan             env::var_os(RUNFILES_DIR_ENV_VAR).expect("bazel did not provide RUNFILES_DIR");
255*d4726bddSHONG Yifan         let runfiles_manifest_file = env::var_os(MANIFEST_FILE_ENV_VAR).unwrap_or("".into());
256*d4726bddSHONG Yifan 
257*d4726bddSHONG Yifan         // Test case 1: Only $RUNFILES_DIR is set.
258*d4726bddSHONG Yifan         {
259*d4726bddSHONG Yifan             env::remove_var(TEST_SRCDIR_ENV_VAR);
260*d4726bddSHONG Yifan             env::remove_var(MANIFEST_FILE_ENV_VAR);
261*d4726bddSHONG Yifan             let r = Runfiles::create().unwrap();
262*d4726bddSHONG Yifan 
263*d4726bddSHONG Yifan             let d = rlocation!(r, "rules_rust");
264*d4726bddSHONG Yifan             let f = rlocation!(r, "rules_rust/tools/runfiles/data/sample.txt");
265*d4726bddSHONG Yifan             assert_eq!(d.join("tools/runfiles/data/sample.txt"), f);
266*d4726bddSHONG Yifan 
267*d4726bddSHONG Yifan             let mut f = File::open(f).unwrap();
268*d4726bddSHONG Yifan 
269*d4726bddSHONG Yifan             let mut buffer = String::new();
270*d4726bddSHONG Yifan             f.read_to_string(&mut buffer).unwrap();
271*d4726bddSHONG Yifan 
272*d4726bddSHONG Yifan             assert_eq!("Example Text!", buffer);
273*d4726bddSHONG Yifan             env::set_var(TEST_SRCDIR_ENV_VAR, &test_srcdir);
274*d4726bddSHONG Yifan             env::set_var(MANIFEST_FILE_ENV_VAR, &runfiles_manifest_file);
275*d4726bddSHONG Yifan         }
276*d4726bddSHONG Yifan         // Test case 2: Only $TEST_SRCDIR is set.
277*d4726bddSHONG Yifan         {
278*d4726bddSHONG Yifan             env::remove_var(RUNFILES_DIR_ENV_VAR);
279*d4726bddSHONG Yifan             env::remove_var(MANIFEST_FILE_ENV_VAR);
280*d4726bddSHONG Yifan             let r = Runfiles::create().unwrap();
281*d4726bddSHONG Yifan 
282*d4726bddSHONG Yifan             let mut f =
283*d4726bddSHONG Yifan                 File::open(rlocation!(r, "rules_rust/tools/runfiles/data/sample.txt")).unwrap();
284*d4726bddSHONG Yifan 
285*d4726bddSHONG Yifan             let mut buffer = String::new();
286*d4726bddSHONG Yifan             f.read_to_string(&mut buffer).unwrap();
287*d4726bddSHONG Yifan 
288*d4726bddSHONG Yifan             assert_eq!("Example Text!", buffer);
289*d4726bddSHONG Yifan             env::set_var(RUNFILES_DIR_ENV_VAR, &runfiles_dir);
290*d4726bddSHONG Yifan             env::set_var(MANIFEST_FILE_ENV_VAR, &runfiles_manifest_file);
291*d4726bddSHONG Yifan         }
292*d4726bddSHONG Yifan 
293*d4726bddSHONG Yifan         // Test case 3: Neither are set
294*d4726bddSHONG Yifan         {
295*d4726bddSHONG Yifan             env::remove_var(RUNFILES_DIR_ENV_VAR);
296*d4726bddSHONG Yifan             env::remove_var(TEST_SRCDIR_ENV_VAR);
297*d4726bddSHONG Yifan             env::remove_var(MANIFEST_FILE_ENV_VAR);
298*d4726bddSHONG Yifan 
299*d4726bddSHONG Yifan             let r = Runfiles::create().unwrap();
300*d4726bddSHONG Yifan 
301*d4726bddSHONG Yifan             let mut f =
302*d4726bddSHONG Yifan                 File::open(rlocation!(r, "rules_rust/tools/runfiles/data/sample.txt")).unwrap();
303*d4726bddSHONG Yifan 
304*d4726bddSHONG Yifan             let mut buffer = String::new();
305*d4726bddSHONG Yifan             f.read_to_string(&mut buffer).unwrap();
306*d4726bddSHONG Yifan 
307*d4726bddSHONG Yifan             assert_eq!("Example Text!", buffer);
308*d4726bddSHONG Yifan 
309*d4726bddSHONG Yifan             env::set_var(TEST_SRCDIR_ENV_VAR, &test_srcdir);
310*d4726bddSHONG Yifan             env::set_var(RUNFILES_DIR_ENV_VAR, &runfiles_dir);
311*d4726bddSHONG Yifan             env::set_var(MANIFEST_FILE_ENV_VAR, &runfiles_manifest_file);
312*d4726bddSHONG Yifan         }
313*d4726bddSHONG Yifan     }
314*d4726bddSHONG Yifan 
315*d4726bddSHONG Yifan     #[test]
test_manifest_based_can_read_data_from_runfiles()316*d4726bddSHONG Yifan     fn test_manifest_based_can_read_data_from_runfiles() {
317*d4726bddSHONG Yifan         let mut path_mapping = HashMap::new();
318*d4726bddSHONG Yifan         path_mapping.insert("a/b".into(), "c/d".into());
319*d4726bddSHONG Yifan         let r = Runfiles {
320*d4726bddSHONG Yifan             mode: Mode::ManifestBased(path_mapping),
321*d4726bddSHONG Yifan             repo_mapping: RepoMapping::new(),
322*d4726bddSHONG Yifan         };
323*d4726bddSHONG Yifan 
324*d4726bddSHONG Yifan         assert_eq!(r.rlocation("a/b"), PathBuf::from("c/d"));
325*d4726bddSHONG Yifan     }
326*d4726bddSHONG Yifan }
327