rust_docx/
transport.rs

1//! Out-of-band payloads for large image fields during model deserialization.
2//!
3//! `DocNode` and `ParagraphChildNode` are untagged, so serde buffers each node
4//! into a generic tree before matching a variant. A JSON number array then costs
5//! one buffered value per image byte, repeated for every untagged level. Callers
6//! that already hold image bytes and data URIs can instead place them in
7//! [`TransportPayloads`] and write a short token in the JSON; the image field
8//! deserializers resolve the token, so only the token is buffered.
9//!
10//! Image `data` also accepts a base64 string or the legacy number array, and
11//! `src` accepts any ordinary string.
12
13use std::cell::RefCell;
14use std::fmt;
15
16use base64::Engine;
17use serde::de::{self, Deserializer, SeqAccess, Visitor};
18
19/// Prefix of a token that refers to [`TransportPayloads::bytes`].
20pub const BYTES_TOKEN_PREFIX: &str = "\u{0}rust-docx-bytes:";
21/// Prefix of a token that refers to [`TransportPayloads::text`].
22pub const TEXT_TOKEN_PREFIX: &str = "\u{0}rust-docx-text:";
23
24/// Large values referenced by tokens in model JSON.
25#[derive(Debug, Default)]
26pub struct TransportPayloads {
27    pub bytes: Vec<Vec<u8>>,
28    pub text: Vec<String>,
29}
30
31thread_local! {
32    static PAYLOADS: RefCell<Option<TransportPayloads>> = const { RefCell::new(None) };
33}
34
35struct Installed(Option<TransportPayloads>);
36
37impl Drop for Installed {
38    fn drop(&mut self) {
39        let previous = self.0.take();
40        PAYLOADS.with(|slot| *slot.borrow_mut() = previous);
41    }
42}
43
44/// Runs `parse` with `payloads` available to token-aware image fields.
45///
46/// Tokens are only resolved inside `parse`; the payloads are released when it
47/// returns, including on panic.
48pub fn with_transport_payloads<T>(payloads: TransportPayloads, parse: impl FnOnce() -> T) -> T {
49    let previous = PAYLOADS.with(|slot| slot.borrow_mut().replace(payloads));
50    let _installed = Installed(previous);
51    parse()
52}
53
54fn token_index(value: &str, prefix: &str) -> Option<Result<usize, String>> {
55    let index = value.strip_prefix(prefix)?;
56    Some(index.parse().map_err(|_| format!("invalid transport token {value:?}")))
57}
58
59fn resolve<T: Clone>(index: usize, select: impl Fn(&TransportPayloads) -> &Vec<T>) -> Result<T, String> {
60    PAYLOADS.with(|slot| {
61        let payloads = slot.borrow();
62        let payloads = payloads
63            .as_ref()
64            .ok_or_else(|| "transport token used outside a transport scope".to_string())?;
65        // Untagged variant attempts may resolve the same field more than once,
66        // so the payload is cloned rather than moved out.
67        select(payloads)
68            .get(index)
69            .cloned()
70            .ok_or_else(|| format!("transport payload {index} is missing"))
71    })
72}
73
74struct ImageDataVisitor;
75
76impl<'de> Visitor<'de> for ImageDataVisitor {
77    type Value = Option<Vec<u8>>;
78
79    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
80        formatter.write_str("image bytes as a byte array, base64 string or transport token")
81    }
82
83    fn visit_none<E: de::Error>(self) -> Result<Self::Value, E> {
84        Ok(None)
85    }
86
87    fn visit_unit<E: de::Error>(self) -> Result<Self::Value, E> {
88        Ok(None)
89    }
90
91    fn visit_some<D: Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
92        deserializer.deserialize_any(self)
93    }
94
95    fn visit_str<E: de::Error>(self, value: &str) -> Result<Self::Value, E> {
96        if let Some(index) = token_index(value, BYTES_TOKEN_PREFIX) {
97            return index.and_then(|index| resolve(index, |p| &p.bytes)).map(Some).map_err(E::custom);
98        }
99        base64::engine::general_purpose::STANDARD
100            .decode(value)
101            .map(Some)
102            .map_err(|error| E::custom(format!("invalid base64 image data: {error}")))
103    }
104
105    fn visit_bytes<E: de::Error>(self, value: &[u8]) -> Result<Self::Value, E> {
106        Ok(Some(value.to_vec()))
107    }
108
109    fn visit_byte_buf<E: de::Error>(self, value: Vec<u8>) -> Result<Self::Value, E> {
110        Ok(Some(value))
111    }
112
113    fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
114        let mut bytes = Vec::with_capacity(seq.size_hint().unwrap_or(0).min(1 << 20));
115        while let Some(byte) = seq.next_element::<u8>()? {
116            bytes.push(byte);
117        }
118        Ok(Some(bytes))
119    }
120}
121
122/// Deserializes optional image bytes from a byte array, base64 string or bytes token.
123pub fn deserialize_image_data<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error> {
124    deserializer.deserialize_option(ImageDataVisitor)
125}
126
127struct TextVisitor;
128
129impl<'de> Visitor<'de> for TextVisitor {
130    type Value = Option<String>;
131
132    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
133        formatter.write_str("a string or transport token")
134    }
135
136    fn visit_none<E: de::Error>(self) -> Result<Self::Value, E> {
137        Ok(None)
138    }
139
140    fn visit_unit<E: de::Error>(self) -> Result<Self::Value, E> {
141        Ok(None)
142    }
143
144    fn visit_some<D: Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
145        deserializer.deserialize_str(self)
146    }
147
148    fn visit_str<E: de::Error>(self, value: &str) -> Result<Self::Value, E> {
149        match token_index(value, TEXT_TOKEN_PREFIX) {
150            Some(index) => index.and_then(|index| resolve(index, |p| &p.text)).map(Some).map_err(E::custom),
151            None => Ok(Some(value.to_owned())),
152        }
153    }
154
155    fn visit_string<E: de::Error>(self, value: String) -> Result<Self::Value, E> {
156        if value.starts_with(TEXT_TOKEN_PREFIX) {
157            return self.visit_str(&value);
158        }
159        Ok(Some(value))
160    }
161}
162
163/// Deserializes an optional string that may be a text token.
164pub fn deserialize_optional_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<String>, D::Error> {
165    deserializer.deserialize_option(TextVisitor)
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use crate::ImageRunNode;
172
173    fn image(json: &str) -> Result<ImageRunNode, serde_json::Error> {
174        serde_json::from_str(json)
175    }
176
177    #[test]
178    fn image_data_accepts_number_arrays_base64_and_null() {
179        assert_eq!(image(r#"{"data":[1,2,255]}"#).unwrap().data, Some(vec![1, 2, 255]));
180        assert_eq!(image(r#"{"data":"AQL/"}"#).unwrap().data, Some(vec![1, 2, 255]));
181        assert_eq!(image(r#"{"data":null}"#).unwrap().data, None);
182        assert_eq!(image(r#"{}"#).unwrap().data, None);
183        assert!(image(r#"{"data":"not base64!"}"#).is_err());
184    }
185
186    #[test]
187    fn tokens_resolve_inside_untagged_nodes() {
188        let payloads = TransportPayloads {
189            bytes: vec![vec![9, 8, 7]],
190            text: vec!["data:image/png;base64,CQgH".into()],
191        };
192        let json = format!(
193            r#"{{"type":"paragraph","children":[{{"type":"image","src":"{}0","data":"{}0","partName":"word/media/a.png"}}]}}"#,
194            TEXT_TOKEN_PREFIX.replace('\0', "\\u0000"),
195            BYTES_TOKEN_PREFIX.replace('\0', "\\u0000"),
196        );
197        let node: crate::DocNode = with_transport_payloads(payloads, || serde_json::from_str(&json)).unwrap();
198        let crate::DocNode::Paragraph(paragraph) = node else { panic!("expected paragraph") };
199        let crate::ParagraphChildNode::Image(image) = &paragraph.children[0] else { panic!("expected image") };
200        assert_eq!(image.data.as_deref(), Some(&[9, 8, 7][..]));
201        assert_eq!(image.src.as_deref(), Some("data:image/png;base64,CQgH"));
202    }
203
204    #[test]
205    fn tokens_fail_outside_a_scope_or_when_missing() {
206        let token = format!(r#"{{"data":"{}0"}}"#, BYTES_TOKEN_PREFIX.replace('\0', "\\u0000"));
207        assert!(image(&token).is_err());
208        let missing = with_transport_payloads(TransportPayloads::default(), || image(&token));
209        assert!(missing.is_err());
210        assert!(PAYLOADS.with(|slot| slot.borrow().is_none()));
211    }
212}
rust_docx/transport.rs

rust_docx/
transport.rs
