Skip to content

Commit

Permalink
Add string_sanitization_mode to writer
Browse files Browse the repository at this point in the history
  • Loading branch information
sorokya committed Oct 13, 2024
1 parent be029f5 commit 7d1127c
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 2 deletions.
7 changes: 7 additions & 0 deletions build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 => {
Expand All @@ -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 => {
Expand Down Expand Up @@ -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");
}

Expand Down
29 changes: 27 additions & 2 deletions src/data/eo_writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ impl From<String> for EoWriterError {
/// ````
pub struct EoWriter {
data: BytesMut,
string_sanitization_mode: bool,
}

impl EoWriter {
Expand All @@ -59,6 +60,7 @@ impl EoWriter {
pub fn with_capacity(size: usize) -> Self {
Self {
data: BytesMut::with_capacity(size),
..Default::default()
}
}

Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 7d1127c

Please sign in to comment.