1use bstr::{BStr, BString, ByteSlice};
2use config::any_err;
3use mailparsing::{DecodedBody, MimePart, PartPointer};
4use mlua::{MetaMethod, UserData, UserDataFields, UserDataMethods, UserDataRef};
5use parking_lot::Mutex;
6use std::sync::Arc;
7
8#[derive(Clone)]
9pub struct PartRef {
10 root_part: Arc<Mutex<MimePart<'static>>>,
11 ptr: PartPointer,
12}
13
14impl PartRef {
15 pub fn new(part: MimePart<'static>) -> Self {
16 Self {
17 root_part: Arc::new(Mutex::new(part)),
18 ptr: PartPointer::root(),
19 }
20 }
21
22 pub fn resolve(&self) -> anyhow::Result<MimePart<'_>> {
23 let root = self.root_part.lock();
24 root.resolve_ptr(self.ptr.clone())
25 .ok_or_else(|| anyhow::anyhow!("failed to resolve PartRef to MimePart"))
26 .cloned()
27 }
28
29 pub fn mutate<F: FnOnce(&mut MimePart) -> anyhow::Result<R>, R>(
30 &self,
31 f: F,
32 ) -> anyhow::Result<R> {
33 let mut root = self.root_part.lock();
34 let part = root.resolve_ptr_mut(self.ptr.clone());
35 match part {
36 Some(p) => (f)(p),
37 None => anyhow::bail!("failed to resolve PartRef to MimePart"),
38 }
39 }
40
41 pub fn make_ref(&self, ptr: PartPointer) -> Self {
42 Self {
43 root_part: self.root_part.clone(),
44 ptr,
45 }
46 }
47
48 pub fn get_simple_structure(&self) -> anyhow::Result<SimpleStructure> {
49 let part = self.resolve()?;
50
51 let parts = part.simplified_structure_pointers().map_err(any_err)?;
52
53 let mut attachments = vec![];
54 for ptr in parts.attachments {
55 attachments.push(self.make_ref(ptr));
56 }
57
58 Ok(SimpleStructure {
59 text_part: parts.text_part.map(|ptr| self.make_ref(ptr)),
60 html_part: parts.html_part.map(|ptr| self.make_ref(ptr)),
61 header_part: self.make_ref(parts.header_part),
62 attachments,
63 })
64 }
65
66 pub fn replace_body(
67 &self,
68 body: mlua::String,
69 mut content_type: Option<BString>,
70 ) -> anyhow::Result<()> {
71 self.mutate(|part| {
72 if content_type.is_none() {
73 if let Ok(Some(params)) = part.headers().content_type() {
74 content_type.replace(params.value);
75 }
76 }
77
78 match body.to_str() {
79 Ok(s) => {
80 part.replace_text_body(
81 content_type
82 .as_ref()
83 .map(|b| b.as_bstr())
84 .unwrap_or_else(|| BStr::new("text/plain")),
85 &*s,
86 )?;
87 }
88 _ => {
89 part.replace_binary_body(
90 content_type
91 .as_ref()
92 .map(|b| b.as_bstr())
93 .unwrap_or_else(|| BStr::new("application/octet-stream")),
94 &body.as_bytes(),
95 )?;
96 }
97 }
98 Ok(())
99 })
100 }
101}
102
103impl UserData for PartRef {
104 fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
105 methods.add_method(
106 "append_part",
107 move |_lua, this, part: UserDataRef<PartRef>| {
108 let part_to_append = part.resolve().map_err(any_err)?.to_owned();
109
110 this.mutate(|this_part| {
111 this_part.child_parts_mut().push(part_to_append);
112 Ok(())
113 })
114 .map_err(any_err)?;
115
116 Ok(())
117 },
118 );
119
120 methods.add_method("get_simple_structure", move |lua, this, ()| {
121 let s = this.get_simple_structure().map_err(any_err)?;
122
123 let result = lua.create_table()?;
124 result.set("text_part", s.text_part)?;
125 result.set("html_part", s.html_part)?;
126 result.set("header_part", s.header_part)?;
127
128 let attachments = lua.create_table()?;
129 for a in s.attachments {
130 let a_part = a.resolve().map_err(any_err)?;
131 let attach = lua.create_table()?;
132
133 let mut file_name = None;
134 let mut inline = true;
135 let mut content_id = None;
136 let mut content_type = None;
137
138 let info = a_part.rfc2045_info();
139 if let Some(mut opts) = info.attachment_options {
140 if let Some(name) = opts.file_name.take() {
141 file_name.replace(name);
142 }
143 inline = opts.inline;
144 content_id = opts.content_id;
145 }
146
147 if let Some(ct) = info.content_type {
148 content_type.replace(ct.value);
149 }
150
151 attach.set("file_name", file_name)?;
152 attach.set("inline", inline)?;
153 attach.set("content_id", content_id)?;
154 match content_type {
155 Some(ct) => {
156 attach.set("content_type", ct)?;
157 }
158 None => {
159 attach.set("content_type", "application/octet-stream")?;
160 }
161 }
162
163 attach.set("part", a)?;
164 attachments.push(attach)?;
165 }
166
167 result.set("attachments", attachments)?;
168
169 Ok(result)
170 });
171
172 methods.add_meta_method(MetaMethod::ToString, move |lua, this, ()| {
173 let root = this.root_part.lock();
174 lua.create_string(root.to_message_bytes())
175 });
176
177 methods.add_method("rebuild", move |_lua, this, ()| {
178 let root = this.root_part.lock();
179 let part = root.rebuild(None).map_err(any_err)?;
180 Ok(PartRef::new(part))
181 });
182
183 methods.add_method(
184 "replace_body",
185 move |_lua, this, (body, content_type): (mlua::String, Option<String>)| {
186 this.replace_body(body, content_type.map(Into::into))
187 .map_err(any_err)
188 },
189 );
190 }
191
192 fn add_fields<F: UserDataFields<Self>>(fields: &mut F) {
193 fields.add_field_method_get("body", |lua, this| {
194 let part = this.resolve().map_err(any_err)?;
195 let body = part.body().map_err(any_err)?;
196 match body {
197 DecodedBody::Text(s) => lua.create_string(s.as_bytes()),
198 DecodedBody::Binary(b) => lua.create_string(b),
199 }
200 });
201 fields.add_field_method_set("body", |_lua, this, body: mlua::String| {
202 this.replace_body(body, None).map_err(any_err)
203 });
204 fields.add_field_method_get("headers", |_lua, this| {
205 let _part = this.resolve().map_err(any_err)?;
206 Ok(crate::headers::HeaderMapRef::new(this.clone()))
207 });
208 fields.add_field_method_get("id", |_lua, this| Ok(this.ptr.id_string()));
209 }
210}
211
212pub struct SimpleStructure {
213 pub text_part: Option<PartRef>,
214 pub html_part: Option<PartRef>,
215 pub header_part: PartRef,
216 pub attachments: Vec<PartRef>,
217}