fix(rust): address removal review feedback

rust/README.md

@@ -109,10 +109,16 @@ not guarantee the original pixels or erase arbitrary third-party watermarks.
 The CLI writes PNG or JPEG (`--quality 90` by default for JPEG); the library
 accepts decoded `DynamicImage` values. Missing or incompatible remover models
 return errors. A shared remover serializes inference independently of encoders
-and decoders. Existing callers incur no removal-model startup or memory cost.
+and decoders. Encode/decode-only callers that do not construct a `WatermarkRemover`
+incur no additional removal-model startup or memory cost. Constructing a remover
+loads a separate model session; removal also allocates image and inference buffers.
 The library returns a floating-point RGB/RGBA image; convert it to 8-bit pixels
 before saving PNG/JPEG. The CLI performs this conversion and replaces the output
 only after encoding succeeds. `--quality` is accepted only for JPEG output.
+PNG output retains RGB/RGBA channel selection. When replacing an existing removal
+output, the CLI preserves its ordinary file permissions (not ownership or ACLs).
+New removal outputs use private temporary-file permissions (owner read/write on
+Unix); on Unix, replacement does not retain special permission bits.
 
 ## Running the benchmarks
 

rust/crates/trustmark-cli/src/main.rs

@@ -169,7 +169,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     encoded.to_rgb8().write_with_encoder(encoder).unwrap();
                 }
                 _ => {
-                    encoded.to_rgba8().save(&output).unwrap();
+                    output_pixels(&encoded).save(&output).unwrap();
                 }
             }
         }
@@ -215,6 +215,15 @@ fn remove_image(
     save_removed_image(&recovered, output, format, quality)
 }
 
+// Preserve the presence of alpha when converting float model output for encoding.
+fn output_pixels(image: &image::DynamicImage) -> image::DynamicImage {
+    if image.color().has_alpha() {
+        image.to_rgba8().into()
+    } else {
+        image.to_rgb8().into()
+    }
+}
+
 fn save_removed_image(
     recovered: &image::DynamicImage,
     output: &std::path::Path,
@@ -225,6 +234,11 @@ fn save_removed_image(
         .parent()
         .filter(|path| !path.as_os_str().is_empty())
         .unwrap_or_else(|| std::path::Path::new("."));
+    let permissions = match std::fs::metadata(output) {
+        Ok(metadata) => Some(metadata.permissions()),
+        Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
+        Err(error) => return Err(error.into()),
+    };
     // Encode beside the destination, then atomically replace it. A failure
     // leaves the old output intact and removes the temporary file on drop.
     let mut file = tempfile::NamedTempFile::new_in(parent)?;
@@ -236,8 +250,18 @@ fn save_removed_image(
                 quality.unwrap_or(90),
             ))?;
     } else {
-        image::DynamicImage::ImageRgba8(recovered.to_rgba8())
-            .write_to(file.as_file_mut(), format)?;
+        output_pixels(recovered).write_to(file.as_file_mut(), format)?;
+    }
+    if let Some(permissions) = permissions {
+        // Preserve ordinary access bits, not setuid/setgid/sticky bits on a new file.
+        #[cfg(unix)]
+        let permissions = {
+            use std::os::unix::fs::PermissionsExt as _;
+            let mut permissions = permissions;
+            permissions.set_mode(permissions.mode() & 0o777);
+            permissions
+        };
+        file.as_file().set_permissions(permissions)?;
     }
     file.as_file().sync_all()?;
     file.persist(output)?;
@@ -359,15 +383,84 @@ mod tests {
         .into()
     }
 
+    #[test]
+    fn png_output_preserves_rgb_and_rgba_channels() {
+        let directory = tempfile::tempdir().unwrap();
+        let rgb = test_image().into_rgb32f();
+        let rgba = image::Rgba32FImage::from_pixel(32, 32, image::Rgba([0.25, 0.5, 0.75, 0.4]));
+        for original in [image::DynamicImage::from(rgb), rgba.into()] {
+            // Exercise the shared encode conversion and the removal writer.
+            let encoded = directory.path().join("encoded.png");
+            output_pixels(&original).save(&encoded).unwrap();
+            let removed = directory.path().join("removed.png");
+            save_removed_image(&original, &removed, ImageFormat::Png, None).unwrap();
+            for path in [encoded, removed] {
+                let decoded = image::open(path).unwrap();
+                assert_eq!(decoded.color().has_alpha(), original.color().has_alpha());
+                assert_eq!(decoded.dimensions(), original.dimensions());
+                assert_eq!(decoded.to_rgba8(), original.to_rgba8());
+            }
+        }
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn removal_output_preserves_existing_access_permissions() {
+        use std::os::unix::fs::PermissionsExt as _;
+
+        let directory = tempfile::tempdir().unwrap();
+        let output = directory.path().join("existing.png");
+        for mode in [0o600, 0o640, 0o644, 0o444, 0o2640] {
+            std::fs::write(&output, b"old output").unwrap();
+            std::fs::set_permissions(&output, std::fs::Permissions::from_mode(mode)).unwrap();
+            save_removed_image(&test_image(), &output, ImageFormat::Png, None).unwrap();
+            assert_eq!(
+                std::fs::metadata(&output).unwrap().permissions().mode() & 0o7777,
+                mode & 0o777
+            );
+            assert_eq!(
+                image::open(&output).unwrap().to_rgb8(),
+                test_image().to_rgb8()
+            );
+            // Permit the next iteration to replace the fixture contents.
+            std::fs::set_permissions(&output, std::fs::Permissions::from_mode(0o600)).unwrap();
+        }
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn new_removal_output_is_private() {
+        use std::os::unix::fs::PermissionsExt as _;
+
+        let directory = tempfile::tempdir().unwrap();
+        let output = directory.path().join("new.png");
+        save_removed_image(&test_image(), &output, ImageFormat::Png, None).unwrap();
+        let mode = std::fs::metadata(output).unwrap().permissions().mode();
+        assert_eq!(mode & 0o777, 0o600);
+    }
+
     #[test]
     fn output_encoding_failure_preserves_existing_file() {
         let directory = tempfile::tempdir().unwrap();
         let output = directory.path().join("existing.png");
         std::fs::write(&output, b"original output").unwrap();
+        #[cfg(unix)]
+        {
+            use std::os::unix::fs::PermissionsExt as _;
+            std::fs::set_permissions(&output, std::fs::Permissions::from_mode(0o640)).unwrap();
+        }
         // An 8-bit image cannot be encoded as OpenEXR, forcing an encoder error
         // after the temporary file has been created.
         assert!(save_removed_image(&test_image(), &output, ImageFormat::OpenExr, None).is_err());
         assert_eq!(std::fs::read(&output).unwrap(), b"original output");
+        #[cfg(unix)]
+        {
+            use std::os::unix::fs::PermissionsExt as _;
+            assert_eq!(
+                std::fs::metadata(&output).unwrap().permissions().mode() & 0o777,
+                0o640
+            );
+        }
         assert_eq!(std::fs::read_dir(directory.path()).unwrap().count(), 1);
     }
 

rust/src/image_processing.rs

@@ -204,9 +204,23 @@ pub(super) fn apply_removal_residual(
     residual: ArrayD<f32>,
     variant: Variant,
 ) -> DynamicImage {
-    let (w, h, left, top) = center_crop_size_and_offset(variant, &input);
-    let has_alpha = input.color().has_alpha();
-    let mut output = input.to_rgba32f();
+    let crop = center_crop_size_and_offset(variant, &input);
+    if input.color().has_alpha() {
+        let mut output = input.into_rgba32f();
+        add_removal_residual(&mut output, &residual, crop);
+        output.into()
+    } else {
+        let mut output = input.into_rgb32f();
+        add_removal_residual(&mut output, &residual, crop);
+        output.into()
+    }
+}
+
+fn add_removal_residual<P: Pixel<Subpixel = f32>>(
+    output: &mut ImageBuffer<P, Vec<f32>>,
+    residual: &ArrayD<f32>,
+    (w, h, left, top): (u32, u32, u32, u32),
+) {
     let feather = ((w.min(h) as f32 * 0.01) as u32).clamp(1, 50);
 
     for y in 0..h {
@@ -222,23 +236,18 @@ pub(super) fn apply_removal_residual(
             } else {
                 1.
             };
-            let target = output.get_pixel_mut(left + x, top + y);
-            for channel in 0..3 {
-                let delta = sample_removal_residual(&residual, channel, x, y, w, h);
+            let target = output.get_pixel_mut(left + x, top + y).channels_mut();
+            for (channel, value) in target.iter_mut().take(COLOR_CHANNELS).enumerate() {
+                let delta = sample_removal_residual(residual, channel, x, y, w, h);
                 // Python operates on 8-bit RGB and truncates the final patch, including
                 // feathered pixels. Rounding instead can leave a recoverable payload.
-                let original = (target[channel] as f64 * 255.).round();
+                let original = (*value as f64 * 255.).round();
                 let recovered = (original + delta as f64 * 127.5).clamp(0., 255.);
                 let blended = alpha as f64 * recovered + (1. - alpha as f64) * original;
-                target[channel] = blended.floor() as f32 / 255.;
+                *value = blended.floor() as f32 / 255.;
             }
         }
     }
-    if has_alpha {
-        output.into()
-    } else {
-        DynamicImage::ImageRgba32F(output).to_rgb32f().into()
-    }
 }
 
 // PyTorch interpolate(mode="bilinear", align_corners=False), without antialiasing
@@ -525,6 +534,42 @@ mod tests {
         assert!(output.pixels().all(|p| p[3] == 12345));
     }
 
+    #[test]
+    fn removal_reuses_owned_float_buffers_and_preserves_channel_layout() {
+        let inputs = [
+            DynamicImage::ImageRgb32F(Rgb32FImage::from_pixel(3, 2, image::Rgb([0.5; 3]))),
+            DynamicImage::ImageRgba32F(Rgba32FImage::from_pixel(
+                3,
+                2,
+                image::Rgba([0.5, 0.5, 0.5, 0.25]),
+            )),
+        ];
+        for input in inputs {
+            let color = input.color();
+            let pointer = match &input {
+                DynamicImage::ImageRgb32F(image) => image.as_raw().as_ptr(),
+                DynamicImage::ImageRgba32F(image) => image.as_raw().as_ptr(),
+                _ => unreachable!(),
+            };
+            let residual = Array4::zeros([1, 3, 256, 256]).into_dyn();
+            let output = apply_removal_residual(input, residual, Variant::Q);
+            assert_eq!(output.color(), color);
+            let output_pointer = match &output {
+                DynamicImage::ImageRgb32F(image) => image.as_raw().as_ptr(),
+                DynamicImage::ImageRgba32F(image) => image.as_raw().as_ptr(),
+                _ => unreachable!(),
+            };
+            assert_eq!(output_pointer, pointer);
+            if color.has_alpha() {
+                assert!(output
+                    .as_rgba32f()
+                    .unwrap()
+                    .pixels()
+                    .all(|pixel| pixel[3] == 0.25));
+            }
+        }
+    }
+
     #[test]
     fn removal_crop_and_feather_match_python_for_all_variants() {
         for variant in [Variant::B, Variant::C, Variant::P, Variant::Q] {

rust/src/remover.rs

@@ -21,8 +21,9 @@ const MODEL_SIZE: u32 = 256;
 
 /// An explicitly loaded TrustMark watermark remover.
 ///
-/// Loads only `remover_{variant}.onnx`; existing [`crate::Trustmark`] callers incur
-/// no additional startup, memory, or inference cost. Removal is lossy and does not
+/// Loads only `remover_{variant}.onnx`. Callers using only [`crate::Trustmark`] incur
+/// no additional removal cost. Constructing a remover loads a separate model session,
+/// and removal allocates image and inference buffers. Removal is lossy and does not
 /// guarantee restoration of the original pixels or removal of arbitrary watermarks.
 ///
 /// ```no_run