1use core::convert::{TryFrom, TryInto};
2use std::collections::BTreeMap;
3
4use bytes::Bytes;
5use libipld_core::cid::Cid;
6use libipld_core::error::{Result, TypeError, TypeErrorType};
7use libipld_core::ipld::Ipld;
8use quick_protobuf::sizeofs::{sizeof_len, sizeof_varint};
9use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer, WriterBackend};
10
11#[derive(Debug, PartialEq, Eq, Clone)]
13pub struct PbLink {
14    pub cid: Cid,
16    pub name: Option<String>,
18    pub size: Option<u64>,
20}
21
22#[derive(Debug, PartialEq, Eq, Clone, Default)]
24pub struct PbNode {
25    pub links: Vec<PbLink>,
27    pub data: Option<Bytes>,
29}
30
31#[derive(Debug, PartialEq, Eq, Clone, Default)]
32pub(crate) struct PbNodeRef<'a> {
33    links: Vec<PbLink>,
34    data: Option<&'a [u8]>,
35}
36
37impl PbNode {
38    pub(crate) fn links(bytes: Bytes, links: &mut impl Extend<Cid>) -> Result<()> {
39        let node = PbNode::from_bytes(bytes)?;
40        for link in node.links {
41            links.extend(Some(link.cid));
42        }
43        Ok(())
44    }
45
46    pub fn from_bytes(buf: Bytes) -> Result<Self> {
48        let mut reader = BytesReader::from_bytes(&buf);
49        let node = PbNodeRef::from_reader(&mut reader, &buf)?;
50        let data = node.data.map(|d| buf.slice_ref(d));
51
52        Ok(PbNode {
53            links: node.links,
54            data,
55        })
56    }
57
58    pub fn into_bytes(mut self) -> Box<[u8]> {
60        self.links.sort_by(|a, b| {
63            let a = a.name.as_ref().map(|s| s.as_bytes()).unwrap_or(&[][..]);
64            let b = b.name.as_ref().map(|s| s.as_bytes()).unwrap_or(&[][..]);
65            a.cmp(b)
66        });
67
68        let mut buf = Vec::with_capacity(self.get_size());
69        let mut writer = Writer::new(&mut buf);
70        self.write_message(&mut writer)
71            .expect("protobuf to be valid");
72        buf.into_boxed_slice()
73    }
74}
75
76impl PbNodeRef<'_> {
77    pub fn into_bytes(mut self) -> Box<[u8]> {
79        self.links.sort_by(|a, b| {
82            let a = a.name.as_ref().map(|s| s.as_bytes()).unwrap_or(&[][..]);
83            let b = b.name.as_ref().map(|s| s.as_bytes()).unwrap_or(&[][..]);
84            a.cmp(b)
85        });
86
87        let mut buf = Vec::with_capacity(self.get_size());
88        let mut writer = Writer::new(&mut buf);
89        self.write_message(&mut writer)
90            .expect("protobuf to be valid");
91        buf.into_boxed_slice()
92    }
93}
94
95impl From<PbNode> for Ipld {
96    fn from(node: PbNode) -> Self {
97        let mut map = BTreeMap::<String, Ipld>::new();
98        let links = node
99            .links
100            .into_iter()
101            .map(|link| link.into())
102            .collect::<Vec<Ipld>>();
103        map.insert("Links".to_string(), links.into());
104        if let Some(data) = node.data {
105            map.insert("Data".to_string(), Ipld::Bytes(data.to_vec()));
106        }
107        map.into()
108    }
109}
110
111impl From<PbLink> for Ipld {
112    fn from(link: PbLink) -> Self {
113        let mut map = BTreeMap::<String, Ipld>::new();
114        map.insert("Hash".to_string(), link.cid.into());
115
116        if let Some(name) = link.name {
117            map.insert("Name".to_string(), name.into());
118        }
119        if let Some(size) = link.size {
120            map.insert("Tsize".to_string(), size.into());
121        }
122        map.into()
123    }
124}
125
126impl<'a> TryFrom<&'a Ipld> for PbNodeRef<'a> {
127    type Error = TypeError;
128
129    fn try_from(ipld: &'a Ipld) -> core::result::Result<Self, Self::Error> {
130        let mut node = PbNodeRef::default();
131
132        match ipld.get("Links")? {
133            Ipld::List(links) => {
134                let mut prev_name = "".to_string();
135                for link in links.iter() {
136                    match link {
137                        Ipld::Map(_) => {
138                            let pb_link: PbLink = link.try_into()?;
139                            if let Some(ref name) = pb_link.name {
141                                if name.as_bytes() < prev_name.as_bytes() {
142                                    return Err(TypeError::new(TypeErrorType::Link, ipld));
145                                }
146                                prev_name = name.clone()
147                            }
148                            node.links.push(pb_link)
149                        }
150                        ipld => return Err(TypeError::new(TypeErrorType::Link, ipld)),
151                    }
152                }
153            }
154            ipld => return Err(TypeError::new(TypeErrorType::List, ipld)),
155        }
156
157        match ipld.get("Data") {
158            Ok(Ipld::Bytes(data)) => node.data = Some(&data[..]),
159            Ok(ipld) => return Err(TypeError::new(TypeErrorType::Bytes, ipld)),
160            _ => (),
161        }
162
163        Ok(node)
164    }
165}
166
167impl TryFrom<&Ipld> for PbLink {
168    type Error = TypeError;
169
170    fn try_from(ipld: &Ipld) -> core::result::Result<PbLink, Self::Error> {
171        if let Ipld::Map(map) = ipld {
172            let mut cid = None;
173            let mut name = None;
174            let mut size = None;
175            for (key, value) in map {
176                match key.as_str() {
177                    "Hash" => {
178                        cid = if let Ipld::Link(cid) = value {
179                            Some(*cid)
180                        } else {
181                            return Err(TypeError::new(TypeErrorType::Link, ipld));
182                        };
183                    }
184                    "Name" => {
185                        name = if let Ipld::String(name) = value {
186                            Some(name.clone())
187                        } else {
188                            return Err(TypeError::new(TypeErrorType::String, ipld));
189                        }
190                    }
191                    "Tsize" => {
192                        size = if let Ipld::Integer(size) = value {
193                            Some(
194                                u64::try_from(*size)
195                                    .map_err(|_| TypeError::new(TypeErrorType::Integer, value))?,
196                            )
197                        } else {
198                            return Err(TypeError::new(TypeErrorType::Integer, ipld));
199                        }
200                    }
201                    _ => {
202                        return Err(TypeError::new(
203                            TypeErrorType::Key("Hash, Name or Tsize".to_string()),
204                            TypeErrorType::Key(key.to_string()),
205                        ));
206                    }
207                }
208            }
209
210            match cid {
212                Some(cid) => Ok(PbLink { cid, name, size }),
213                None => Err(TypeError::new(TypeErrorType::Key("Hash".to_string()), ipld)),
214            }
215        } else {
216            Err(TypeError::new(TypeErrorType::Map, ipld))
217        }
218    }
219}
220
221impl<'a> MessageRead<'a> for PbLink {
222    fn from_reader(r: &mut BytesReader, bytes: &'a [u8]) -> quick_protobuf::Result<Self> {
223        let mut cid = None;
224        let mut name = None;
225        let mut size = None;
226
227        while !r.is_eof() {
228            match r.next_tag(bytes) {
229                Ok(10) => {
230                    let bytes = r.read_bytes(bytes)?;
231                    cid = Some(
232                        Cid::try_from(bytes)
233                            .map_err(|e| quick_protobuf::Error::Message(e.to_string()))?,
234                    );
235                }
236                Ok(18) => name = Some(r.read_string(bytes)?.to_string()),
237                Ok(24) => size = Some(r.read_uint64(bytes)?),
238                Ok(_) => {
239                    return Err(quick_protobuf::Error::Message(
240                        "unexpected bytes".to_string(),
241                    ))
242                }
243                Err(e) => return Err(e),
244            }
245        }
246        Ok(PbLink {
247            cid: cid.ok_or_else(|| quick_protobuf::Error::Message("missing Hash".into()))?,
248            name,
249            size,
250        })
251    }
252}
253
254impl MessageWrite for PbLink {
255    fn get_size(&self) -> usize {
256        let mut size = 0;
257        let l = self.cid.encoded_len();
258        size += 1 + sizeof_len(l);
259
260        if let Some(ref name) = self.name {
261            size += 1 + sizeof_len(name.as_bytes().len());
262        }
263
264        if let Some(tsize) = self.size {
265            size += 1 + sizeof_varint(tsize);
266        }
267        size
268    }
269
270    fn write_message<W: WriterBackend>(&self, w: &mut Writer<W>) -> quick_protobuf::Result<()> {
271        let bytes = self.cid.to_bytes();
272        w.write_with_tag(10, |w| w.write_bytes(&bytes))?;
273
274        if let Some(ref name) = self.name {
275            w.write_with_tag(18, |w| w.write_string(name))?;
276        }
277        if let Some(size) = self.size {
278            w.write_with_tag(24, |w| w.write_uint64(size))?;
279        }
280        Ok(())
281    }
282}
283
284impl<'a> MessageRead<'a> for PbNodeRef<'a> {
285    fn from_reader(r: &mut BytesReader, bytes: &'a [u8]) -> quick_protobuf::Result<Self> {
286        let mut msg = Self::default();
287        let mut links_before_data = false;
288        while !r.is_eof() {
289            match r.next_tag(bytes) {
290                Ok(18) => {
291                    if links_before_data {
293                        return Err(quick_protobuf::Error::Message(
294                            "duplicate Links section".to_string(),
295                        ));
296                    }
297                    msg.links.push(r.read_message::<PbLink>(bytes)?)
298                }
299                Ok(10) => {
300                    msg.data = Some(r.read_bytes(bytes)?);
301                    if !msg.links.is_empty() {
302                        links_before_data = true
303                    }
304                }
305                Ok(_) => {
306                    return Err(quick_protobuf::Error::Message(
307                        "unexpected bytes".to_string(),
308                    ))
309                }
310                Err(e) => return Err(e),
311            }
312        }
313        Ok(msg)
314    }
315}
316
317impl MessageWrite for PbNode {
318    fn get_size(&self) -> usize {
319        let mut size = 0;
320        if let Some(ref data) = self.data {
321            size += 1 + sizeof_len(data.len());
322        }
323
324        size += self
325            .links
326            .iter()
327            .map(|s| 1 + sizeof_len((s).get_size()))
328            .sum::<usize>();
329
330        size
331    }
332
333    fn write_message<W: WriterBackend>(&self, w: &mut Writer<W>) -> quick_protobuf::Result<()> {
334        for s in &self.links {
335            w.write_with_tag(18, |w| w.write_message(s))?;
336        }
337
338        if let Some(ref data) = self.data {
339            w.write_with_tag(10, |w| w.write_bytes(data))?;
340        }
341
342        Ok(())
343    }
344}
345
346impl MessageWrite for PbNodeRef<'_> {
347    fn get_size(&self) -> usize {
348        let mut size = 0;
349        if let Some(data) = self.data {
350            size += 1 + sizeof_len(data.len());
351        }
352
353        size += self
354            .links
355            .iter()
356            .map(|s| 1 + sizeof_len((s).get_size()))
357            .sum::<usize>();
358
359        size
360    }
361
362    fn write_message<W: WriterBackend>(&self, w: &mut Writer<W>) -> quick_protobuf::Result<()> {
363        for s in &self.links {
364            w.write_with_tag(18, |w| w.write_message(s))?;
365        }
366
367        if let Some(data) = self.data {
368            w.write_with_tag(10, |w| w.write_bytes(data))?;
369        }
370
371        Ok(())
372    }
373}