feat: Implementation offsets in paths.json

crates/rattler/src/install/link.rs

@@ -165,6 +165,7 @@ pub fn link_file(
     let link_method = if let Some(PrefixPlaceholder {
         file_mode,
         placeholder,
+        offsets
     }) = path_json_entry.prefix_placeholder.as_ref()
     {
         // Memory map the source file. This provides us with easy access to a continuous stream of
@@ -180,7 +181,8 @@ pub fn link_file(
             fs::File::create(&destination_path)
                 .map_err(LinkFileError::FailedToOpenDestinationFile)?,
         );
-        let mut destination_writer = HashingWriter::<_, rattler_digest::Sha256>::new(destination);
+        let mut destination_writer = 
+            HashingWriter::<_, rattler_digest::Sha256>::new(destination);
 
         // Convert back-slashes (\) on windows with forward-slashes (/) to avoid problems with
         // string escaping. For instance if we replace the prefix in the following text
@@ -203,16 +205,33 @@ pub fn link_file(
             Cow::Borrowed(target_prefix)
         };
 
-        // Replace the prefix placeholder in the file with the new placeholder
-        copy_and_replace_placeholders(
-            source.as_ref(),
-            &mut destination_writer,
-            placeholder,
-            &target_prefix,
-            &target_platform,
-            *file_mode,
-        )
-        .map_err(|err| LinkFileError::IoError(String::from("replacing placeholders"), err))?;
+        // depending on the availability of the offsets 
+        match offsets {
+            Some(offsets) => {
+                copy_and_replace_placeholders_with_offsets(
+                    source.as_ref(),
+                    &mut destination_writer,
+                    placeholder,
+                    &target_prefix,
+                    &target_platform,
+                    *file_mode,
+                    offsets.to_vec()
+                )
+                .map_err(|err| LinkFileError::IoError(String::from("replacing placeholders"), err))?;
+            },
+            None => {
+                // Replace the prefix placeholder in the file with the new placeholder
+                copy_and_replace_placeholders(
+                    source.as_ref(),
+                    &mut destination_writer,
+                    placeholder,
+                    &target_prefix,
+                    &target_platform,
+                    *file_mode,
+                )
+                .map_err(|err| LinkFileError::IoError(String::from("replacing placeholders"), err))?;
+            },
+        }
 
         let (mut file, current_hash) = destination_writer.finalize();
 
@@ -566,6 +585,51 @@ pub fn copy_and_replace_placeholders(
     Ok(())
 }
 
+/// Given the contents of a file copy it to the `destination` and in the process replace the
+/// `prefix_placeholder` text with the `target_prefix` text.
+///
+/// This switches to more specialized functions that handle the replacement of either
+/// textual and binary placeholders, the [`FileMode`] enum switches between the two functions.
+/// See both [`copy_and_replace_cstring_placeholder`] and [`copy_and_replace_textual_placeholder`]
+pub fn copy_and_replace_placeholders_with_offsets(
+    source_bytes: &[u8],
+    mut destination: impl Write,
+    prefix_placeholder: &str,
+    target_prefix: &str,
+    target_platform: &Platform,
+    file_mode: FileMode,
+    offsets: Vec<usize>,
+) -> Result<(), std::io::Error> {
+    match file_mode {
+        FileMode::Text => {
+            copy_and_replace_textual_placeholder_offsets(
+                source_bytes,
+                destination,
+                prefix_placeholder,
+                target_prefix,
+                target_platform,
+                offsets,
+            )?;
+        }
+        FileMode::Binary => {
+            // conda does not replace the prefix in the binary files on windows
+            // DLLs are loaded quite differently anyways (there is no rpath, for example).
+            if target_platform.is_windows() {
+                destination.write_all(source_bytes)?;
+            } else {
+                copy_and_replace_cstring_placeholder_offsets(
+                    source_bytes,
+                    destination,
+                    prefix_placeholder,
+                    target_prefix,
+                    offsets,
+                )?;
+            }
+        }
+    }
+    Ok(())
+}
+
 static SHEBANG_REGEX: Lazy<Regex> = Lazy::new(|| {
     // ^(#!      pretty much the whole match string
     // (?:[ ]*)  allow spaces between #! and beginning of
@@ -712,6 +776,57 @@ pub fn copy_and_replace_textual_placeholder(
     Ok(())
 }
 
+/// Given the contents of a file copy it to the `destination` and in the process replace the
+/// `prefix_placeholder` text with the `target_prefix` text using the offsets from the paths.json
+///
+/// This is a text based version where the complete string is replaced. This works fine for text
+/// files but will not work correctly for binary files where the length of the string is often
+/// important. See [`copy_and_replace_cstring_placeholder`] when you are dealing with binary
+/// content.
+pub fn copy_and_replace_textual_placeholder_offsets(
+    mut source_bytes: &[u8],
+    mut destination: impl Write,
+    prefix_placeholder: &str,
+    target_prefix: &str,
+    target_platform: &Platform,
+    offsets: Vec<usize>,
+) -> Result<(), std::io::Error> {
+    let old_prefix = prefix_placeholder.as_bytes();
+    let new_prefix = target_prefix.as_bytes();
+
+    // check if we have a shebang. We need to handle it differently because it has a maximum length
+    // that can be exceeded in very long target prefix's.
+    if target_platform.is_unix() && source_bytes.starts_with(b"#!") {
+        // extract first line
+        let (first, rest) =
+            source_bytes.split_at(source_bytes.iter().position(|&c| c == b'\n').unwrap_or(0));
+        let first_line = String::from_utf8_lossy(first);
+        let new_shebang = replace_shebang(
+            first_line,
+            (prefix_placeholder, target_prefix),
+            target_platform,
+        );
+        // let replaced = first_line.replace(prefix_placeholder, target_prefix);
+        destination.write_all(new_shebang.as_bytes())?;
+        source_bytes = rest;
+    }
+
+    let mut last_match = 0;
+
+    for offset in offsets {
+            destination.write_all(&source_bytes[last_match..offset])?;
+            destination.write_all(new_prefix)?;
+            last_match = offset + old_prefix.len();
+    }
+
+    // Write remaining bytes
+    if last_match < source_bytes.len() {
+        destination.write_all(&source_bytes[last_match..])?;
+    }
+
+    Ok(())
+}
+
 /// Given the contents of a file, copies it to the `destination` and in the process replace any
 /// binary c-style string that contains the text `prefix_placeholder` with a binary compatible
 /// c-string where the `prefix_placeholder` text is replaced with the `target_prefix` text.
@@ -787,6 +902,80 @@ pub fn copy_and_replace_cstring_placeholder(
     }
 }
 
+/// Given the contents & offsets of a file, copies it to the `destination` and in the process replace 
+/// any binary c-style string that contains the text `prefix_placeholder` with a binary compatible
+/// c-string where the `prefix_placeholder` text is replaced with the `target_prefix` text.
+///
+/// The length of the input will match the output.
+///
+/// This function replaces binary c-style strings using pre-computed offsets for better performance.
+pub fn copy_and_replace_cstring_placeholder_offsets(
+    source_bytes: &[u8],
+    mut destination: impl Write,
+    prefix_placeholder: &str,
+    target_prefix: &str,
+    offsets: Vec<usize>
+) -> Result<(), std::io::Error> {
+    let old_prefix = prefix_placeholder.as_bytes();
+    let new_prefix = target_prefix.as_bytes();
+
+    if new_prefix.len() > old_prefix.len() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "target prefix cannot be longer than the placeholder prefix",
+        ));
+    }
+
+    if offsets.is_empty() {
+        return destination.write_all(source_bytes);
+    }
+
+    let mut last_pos = 0;
+
+    for &offset in &offsets {
+        destination.write_all(&source_bytes[last_pos..offset])?;
+
+        let mut end = offset + old_prefix.len();
+        while end < source_bytes.len() && source_bytes[end] != b'\0' {
+            end += 1;
+        }
+
+        let segment = &source_bytes[offset..end];
+        let old_len = segment.len();
+
+        let mut out = Vec::with_capacity(old_len);
+        let mut remaining = segment;
+        let finder = memchr::memmem::Finder::new(old_prefix);
+
+        while let Some(index) = finder.find(remaining) {
+            out.write_all(&remaining[..index])?;
+            out.write_all(new_prefix)?;
+            remaining = &remaining[index + old_prefix.len()..];
+        }
+        out.write_all(remaining)?;
+
+        if out.len() > old_len {
+            destination.write_all(&out[..old_len])?;
+        } else {
+            destination.write_all(&out)?;
+        }
+
+        let padding = old_len.saturating_sub(out.len());
+        if padding > 0 {
+            destination.write_all(&vec![0; padding])?;
+        }
+
+        last_pos = end;
+    }
+
+    // Write any remaining bytes after the last replacement
+    if last_pos < source_bytes.len() {
+        destination.write_all(&source_bytes[last_pos..])?;
+    }
+
+    Ok(())
+}
+
 fn symlink(source_path: &Path, destination_path: &Path) -> std::io::Result<()> {
     #[cfg(windows)]
     return fs_err::os::windows::fs::symlink_file(source_path, destination_path);
@@ -1094,4 +1283,135 @@ mod test {
         let empty: [u8; 0] = [];
         assert_eq!(FileType::detect(&empty), None);
     }
-}
+
+
+    #[rstest]
+    #[case("Hello, cruel world!", [7].to_vec(), "cruel", "fabulous", "Hello, fabulous world!")]
+    #[case(
+        "prefix_placeholder",
+        [0].to_vec(),
+        "prefix_placeholder",
+        "target_prefix",
+        "target_prefix"
+    )]
+    pub fn test_copy_and_replace_textual_placeholder_with_offsets(
+        #[case] input: &str,
+        #[case] offsets: Vec<usize>,
+        #[case] prefix_placeholder: &str,
+        #[case] target_prefix: &str,
+        #[case] expected_output: &str,
+    ) {
+        let mut output = Cursor::new(Vec::new());
+        super::copy_and_replace_textual_placeholder_offsets(
+            input.as_bytes(),
+            &mut output,
+            prefix_placeholder,
+            target_prefix,
+            &Platform::Linux64,
+            offsets
+        )
+        .unwrap();
+        assert_eq!(
+            &String::from_utf8_lossy(&output.into_inner()),
+            expected_output
+        );
+    }
+
+    #[rstest]
+    #[case(
+        b"12345Hello, fabulous world!\x006789",
+        [12].to_vec(),
+        "fabulous",
+        "cruel",
+        b"12345Hello, cruel world!\x00\x00\x00\x006789"
+    )]
+    pub fn test_copy_and_replace_binary_placeholder_offsets(
+        #[case] input: &[u8],
+        #[case] offsets: Vec<usize>,
+        #[case] prefix_placeholder: &str,
+        #[case] target_prefix: &str,
+        #[case] expected_output: &[u8],
+    ) {
+        assert_eq!(
+            expected_output.len(),
+            input.len(),
+            "input and expected output must have the same length"
+        );
+        let mut output = Cursor::new(Vec::new());
+        super::copy_and_replace_cstring_placeholder_offsets(
+            input,
+            &mut output,
+            prefix_placeholder,
+            target_prefix,
+            offsets,
+        )
+        .unwrap();
+        assert_eq!(&output.into_inner(), expected_output);
+    }
+
+    #[rstest]
+    #[case(b"short\x00", [0].to_vec(), "short", "verylong")]
+    #[case(b"short1234\x00", [0].to_vec(), "short", "verylong")]
+    pub fn test_shorter_binary_placeholder_offsets(
+        #[case] input: &[u8],
+        #[case] offsets: Vec<usize>,
+        #[case] prefix_placeholder: &str,
+        #[case] target_prefix: &str,
+    ) {
+        assert!(target_prefix.len() > prefix_placeholder.len());
+
+        let mut output = Cursor::new(Vec::new());
+        let result = super::copy_and_replace_cstring_placeholder_offsets(
+            input,
+            &mut output,
+            prefix_placeholder,
+            target_prefix,
+            offsets,
+        );
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn replace_binary_path_var_offsets() {
+        let input =
+            b"beginrandomdataPATH=/placeholder/etc/share:/placeholder/bin/:\x00somemoretext";
+        let mut output = Cursor::new(Vec::new());
+        let offsets: Vec<usize> = [20].to_vec();
+        super::copy_and_replace_cstring_placeholder_offsets(input, &mut output, "/placeholder", "/target", offsets)
+            .unwrap();
+        let out = &output.into_inner();
+        assert_eq!(out, b"beginrandomdataPATH=/target/etc/share:/target/bin/:\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00somemoretext");
+        assert_eq!(out.len(), input.len());
+    }
+
+    #[test]
+    fn test_replace_long_prefix_in_text_file_offsets() {
+        let test_data_dir =
+            std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../test-data");
+        let test_file = test_data_dir.join("shebang_test.txt");
+        let prefix_placeholder = "/this/is/placeholder";
+        let mut target_prefix = "/super/long/".to_string();
+        for _ in 0..15 {
+            target_prefix.push_str("verylongstring/");
+        }
+        let input = fs::read(test_file).unwrap();
+
+        let offsets: Vec<usize>= Vec::new();        
+
+        let mut output = Cursor::new(Vec::new());
+        super::copy_and_replace_textual_placeholder_offsets(
+            &input,
+            &mut output,
+            prefix_placeholder,
+            &target_prefix,
+            &Platform::Linux64,
+            offsets,
+        )
+        .unwrap();
+
+        let output = output.into_inner();
+        let replaced = String::from_utf8_lossy(&output);
+        insta::assert_snapshot!(replaced);
+    }
+
+}