nucleide/
module.rs

1// Copyright © 2022-2023 The Nucleide Contributors.
2//
3// Licensed under any of:
4// - Apache License, Version 2.0 (https://www.apache.org/licenses/LICENSE-2.0)
5// - Boost Software License, Version 1.0 (https://www.boost.org/LICENSE_1_0.txt)
6// - MIT License (https://mit-license.org/)
7// At your choosing (See accompanying files LICENSE_APACHE_2_0.txt,
8// LICENSE_MIT.txt and LICENSE_BOOST_1_0.txt).
9
10use alloc::{string::String, vec::Vec};
11use core::mem;
12
13use parity_wasm::elements::{self, Serialize};
14
15use crate::{section::SectionKind, Error, Result, Section};
16
17/// Represents WebAssembly module. Use new to build from buffer.
18#[derive(Debug)]
19pub struct Module(elements::Module);
20
21impl Module {
22    /// Creates a Module from buffer.
23    pub fn new(buf: &[u8]) -> Result<Self> {
24        Ok(Module(elements::Module::from_bytes(buf).map_err(Error)?))
25    }
26
27    /// Returns an iterator over the module’s custom sections.
28    ///
29    /// [`Section`]s are always yielded as the `Any` variant (borrowed).  They
30    /// can be parsed with [`Section::to()`].
31    pub fn sections(&self) -> Result<impl Iterator<Item = Section<'_>>> {
32        const ERROR_MESSAGE: Error = Error::with_msg("Incorrect Section Order");
33
34        let mut kind = SectionKind::Name;
35        let iter = self.0.custom_sections();
36
37        for section in self.0.custom_sections() {
38            match section.name() {
39                "name" if kind <= SectionKind::Name => {
40                    kind = SectionKind::Producers
41                }
42                "producers" if kind <= SectionKind::Producers => {
43                    kind = SectionKind::Daku
44                }
45                "daku" if kind <= SectionKind::Daku => {
46                    kind = SectionKind::Unknown
47                }
48                "name" | "producers" | "daku" => return Err(ERROR_MESSAGE),
49                _ => {}
50            }
51        }
52
53        Ok(iter.map(|section| Section::Any {
54            name: section.name().into(),
55            data: section.payload().into(),
56        }))
57    }
58
59    /// Sets the payload associated with the given custom section, or adds a new
60    /// custom section, as appropriate.
61    pub fn set_section(&mut self, mut section: Section<'_>) -> Option<()> {
62        let (name, data) = section.to_any()?;
63
64        self.0.set_custom_section(name, data.to_vec());
65
66        Some(())
67    }
68
69    /// Removes the given custom section, if it exists. Returns the removed
70    /// section if it existed, or None otherwise.
71    pub fn clear_section(
72        &mut self,
73        name: impl AsRef<str>,
74    ) -> Option<Section<'static>> {
75        let mut section = self.0.clear_custom_section(&name)?;
76        let (mut name, mut data) = (String::new(), Vec::new());
77
78        mem::swap(&mut name, section.name_mut());
79        mem::swap(&mut data, section.payload_mut());
80
81        Some(Section::Any {
82            name: name.into(),
83            data: data.into(),
84        })
85    }
86
87    /// Write out module to a `Vec` of bytes.
88    pub fn into_buffer(self) -> Result<Vec<u8>> {
89        let mut v = Vec::new();
90        self.0.serialize(&mut v).map_err(Error)?;
91        Ok(v)
92    }
93}