From 83581af281a8360e4704d5795670cd09027a4d4a Mon Sep 17 00:00:00 2001
From: cry-inc <cry-inc@users.noreply.github.com>
Date: Wed, 27 Mar 2024 22:33:18 +0100
Subject: [PATCH] Add finalize variant that allows XML customization for E57
 extensions

---
 src/e57_writer.rs     | 16 ++++++++++++++++
 tests/writer_tests.rs | 29 +++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/src/e57_writer.rs b/src/e57_writer.rs
index 8c11a9a..e769cc8 100644
--- a/src/e57_writer.rs
+++ b/src/e57_writer.rs
@@ -103,12 +103,28 @@ impl<T: Write + Read + Seek> E57Writer<T> {
     /// This will generate and write the XML metadata to finalize and complete the E57 file.
     /// Without calling this method before dropping the E57 file will be incomplete and invalid!
     pub fn finalize(&mut self) -> Result<()> {
+        self.finalize_customized_xml(Ok)
+    }
+
+    /// Same as `finalize()` but with additional XML transformation step.
+    ///
+    /// Allows customizing the XML data before its written into the E57 file.
+    /// This is required for adding E57 extension data to the XML.
+    /// The transformer receives an XML string and must return an XML string.
+    /// The client is responsible for parsing, modifying and serializing th XML again in a non-destructive way.
+    /// The E57 library will not validate the XML string before writing it into the E57 file!
+    /// If the transformer fails, the finalization is aborted and any error is forwarded.
+    pub fn finalize_customized_xml(
+        &mut self,
+        transformer: impl Fn(String) -> Result<String>,
+    ) -> Result<()> {
         let xml = serialize_root(
             &self.root,
             &self.pointclouds,
             &self.images,
             &self.extensions,
         )?;
+        let xml = transformer(xml)?;
         let xml_bytes = xml.as_bytes();
         let xml_length = xml_bytes.len();
         let xml_offset = self.writer.physical_position()?;
diff --git a/tests/writer_tests.rs b/tests/writer_tests.rs
index 3cda20b..0f84154 100644
--- a/tests/writer_tests.rs
+++ b/tests/writer_tests.rs
@@ -916,3 +916,32 @@ fn extensions_write_read_blobs() {
 
     remove_file(path).unwrap();
 }
+
+#[test]
+fn custom_xml_test() {
+    let path = Path::new("custom_xml_test.e57");
+    let inserted_xml = "<myext:mytag type=\"Structure\"></myext:mytag>";
+    let extension = Extension::new("myext", "https://mycompany.com/myext");
+
+    {
+        let mut e57_writer = E57Writer::from_file(path, "guid_file").unwrap();
+        e57_writer.register_extesion(extension).unwrap();
+        e57_writer
+            .finalize_customized_xml(|xml| {
+                assert!(!xml.contains(inserted_xml));
+                let from = "</e57Root>";
+                let to = format!("{}\n</e57Root>", inserted_xml);
+                Ok(xml.replace(from, &to))
+            })
+            .unwrap();
+    }
+
+    {
+        let e57 = E57Reader::from_file(path).unwrap();
+        assert_eq!(e57.guid(), "guid_file");
+        let xml = e57.xml();
+        assert!(xml.contains(inserted_xml));
+    }
+
+    remove_file(path).unwrap();
+}