From 7d1127c49ec04d89e3e89170573c59961a71e53c Mon Sep 17 00:00:00 2001
From: Richard Leek <richard@richardleek.com>
Date: Sun, 13 Oct 2024 19:49:20 +0200
Subject: [PATCH] Add string_sanitization_mode to writer

---
 build.rs              |  7 +++++++
 src/data/eo_writer.rs | 29 +++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/build.rs b/build.rs
index 57c52ef..13a367d 100644
--- a/build.rs
+++ b/build.rs
@@ -645,6 +645,9 @@ fn write_struct_serialize(
     enums: &[Enum],
     structs: &[Struct],
 ) {
+    code.push_str(
+        "        let current_string_sanitization_mode = writer.get_string_sanitization_mode();\n",
+    );
     for element in elements {
         match element {
             StructElement::Break => {
@@ -666,6 +669,7 @@ fn write_struct_serialize(
                 generate_serialize_switch(code, name, switch);
             }
             StructElement::Chunked(chunked) => {
+                code.push_str("        writer.set_string_sanitization_mode(true);\n");
                 for element in &chunked.elements {
                     match element {
                         StructElement::Break => {
@@ -695,6 +699,9 @@ fn write_struct_serialize(
             _ => {}
         }
     }
+    code.push_str(
+        "        writer.set_string_sanitization_mode(current_string_sanitization_mode);\n",
+    );
     code.push_str("        Ok(())\n");
 }
 
diff --git a/src/data/eo_writer.rs b/src/data/eo_writer.rs
index 4fce74b..1771533 100644
--- a/src/data/eo_writer.rs
+++ b/src/data/eo_writer.rs
@@ -47,6 +47,7 @@ impl From<String> for EoWriterError {
 /// ````
 pub struct EoWriter {
     data: BytesMut,
+    string_sanitization_mode: bool,
 }
 
 impl EoWriter {
@@ -59,6 +60,7 @@ impl EoWriter {
     pub fn with_capacity(size: usize) -> Self {
         Self {
             data: BytesMut::with_capacity(size),
+            ..Default::default()
         }
     }
 
@@ -112,20 +114,43 @@ impl EoWriter {
         Ok(())
     }
 
+    fn sanitize_string(&self, string: &str) -> String {
+        if self.string_sanitization_mode {
+            string
+                .chars()
+                .map(|c| if c as i32 == 0xff { 0x79 as char } else { c })
+                .collect()
+        } else {
+            string.to_owned()
+        }
+    }
+
     /// adds a string to the data stream
     pub fn add_string(&mut self, string: &str) {
-        let (string, _, _) = WINDOWS_1252.encode(string);
+        let string = self.sanitize_string(string);
+        let (string, _, _) = WINDOWS_1252.encode(&string);
         self.data.put_slice(&string);
     }
 
     /// encodes a string and adds it to the data stream
     pub fn add_encoded_string(&mut self, string: &str) {
-        let (mut string, _, _) = WINDOWS_1252.encode(string);
+        let string = self.sanitize_string(string);
+        let (mut string, _, _) = WINDOWS_1252.encode(&string);
         let string = string.to_mut();
         encode_string(&mut *string);
         self.data.put_slice(string);
     }
 
+    /// gets the string sanitization mode
+    pub fn get_string_sanitization_mode(&self) -> bool {
+        self.string_sanitization_mode
+    }
+
+    /// sets the string sanitization mode
+    pub fn set_string_sanitization_mode(&mut self, mode: bool) {
+        self.string_sanitization_mode = mode;
+    }
+
     /// freezes the data and returns a [Bytes] object that can be freely cloned
     pub fn to_byte_array(self) -> Bytes {
         self.data.freeze()