rust-apk 1.0.0

//! Compiling a [`Document`] to binary XML.
//!
//! Layout follows aapt2: a UTF-16 string pool whose first entries are the
//! names of attributes with resource IDs (in ID order, mirrored by the
//! resource map), then namespace and element chunks. Attributes within an
//! element are sorted by resource ID, which Android's attribute lookup
//! requires.

use std::collections::BTreeMap;

use binrw::binrw;

use super::{ANDROID_NAMESPACE, AttrName, Attribute, Document, Element, MAX_DEPTH, Step, XmlError, walk};
use crate::android::attr;
use crate::res::chunk::{ChunkType, NoHeader, Put, write_chunk};
use crate::res::{ChunkError, Encoding, RawValue, ResourceId, ResourceLookup, StringPoolBuilder, data_type, resolve};

/// `ResXMLTree_node` (after the base chunk header).
#[binrw]
#[brw(little)]
pub(super) struct NodeHeader {
    pub line: u32,
    pub comment: u32,
}

/// `ResXMLTree_attrExt`.
#[binrw]
#[brw(little)]
pub(super) struct ElementStart {
    pub namespace: u32,
    pub name: u32,
    pub attribute_start: u16,
    pub attribute_size: u16,
    pub attribute_count: u16,
    /// 1-based positions of the `android:id`, `class`, and `style` attributes.
    pub id_index: u16,
    pub class_index: u16,
    pub style_index: u16,
}

/// `ResXMLTree_attribute`.
#[binrw]
#[brw(little)]
pub(super) struct AttributeRecord {
    pub namespace: u32,
    pub name: u32,
    pub raw_value: u32,
    pub typed_value: RawValue,
}

pub(super) const NONE: u32 = u32::MAX;
pub(super) const ATTRIBUTE_SIZE: u16 = 20;

impl Document {
    /// Compiles this document to binary XML, resolving `@type/name`
    /// references to application resources through `resources`.
    pub fn encode(&self, resources: &impl ResourceLookup) -> Result<Vec<u8>, XmlError> {
        let pool = Pool::collect(&self.root)?;
        let mut line = 0;
        let mut node = || {
            line += 1;
            NodeHeader { line, comment: NONE }
        };
        let mut out = Vec::new();
        write_chunk(&mut out, ChunkType::XML, &NoHeader, |out| {
            pool.strings.write(out, Encoding::Utf16)?;
            write_chunk(out, ChunkType::XML_RESOURCE_MAP, &NoHeader, |out| {
                pool.ids.iter().for_each(|id| out.put_u32(id.0));
                Ok::<_, XmlError>(())
            })?;
            for (_, step) in walk(&self.root) {
                match step {
                    Step::Open(element) => {
                        for namespace in &element.namespaces {
                            let words = [pool.optional(namespace.prefix.as_deref()), pool.other(&namespace.uri)];
                            write_chunk(out, ChunkType::XML_START_NAMESPACE, &node(), |out| put_words(out, words))?;
                        }
                        let start = pool.element(element, resources)?;
                        write_chunk(out, ChunkType::XML_START_ELEMENT, &node(), |out| {
                            out.put(&start.header);
                            start.attributes.iter().for_each(|record| out.put(record));
                            Ok::<_, XmlError>(())
                        })?;
                    }
                    Step::Text(text) => {
                        write_chunk(out, ChunkType::XML_CDATA, &node(), |out| {
                            out.put_u32(pool.other(text));
                            out.put(&RawValue::new(data_type::NULL, 0));
                            Ok::<_, XmlError>(())
                        })?;
                    }
                    Step::Close(element) => {
                        let words = [pool.optional(element.namespace.as_deref()), pool.other(&element.name)];
                        write_chunk(out, ChunkType::XML_END_ELEMENT, &node(), |out| put_words(out, words))?;
                        for namespace in element.namespaces.iter().rev() {
                            let words = [pool.optional(namespace.prefix.as_deref()), pool.other(&namespace.uri)];
                            write_chunk(out, ChunkType::XML_END_NAMESPACE, &node(), |out| put_words(out, words))?;
                        }
                    }
                }
            }
            Ok::<_, XmlError>(())
        })?;
        Ok(out)
    }

    /// Checks what [`encode`](Self::encode) checks before it resolves
    /// references: nesting depth, and every attribute name and value.
    pub(crate) fn check(&self) -> Result<(), XmlError> {
        Pool::collect(&self.root).map(drop)
    }
}

fn put_words(out: &mut Vec<u8>, words: [u32; 2]) -> Result<(), XmlError> {
    words.into_iter().for_each(|word| out.put_u32(word));
    Ok(())
}

/// A document's string pool, borrowing every string from the document.
struct Pool<'a> {
    strings: StringPoolBuilder<'a>,
    /// The resource map: the ID of each leading attribute name.
    ids: Vec<ResourceId>,
    attr_index: BTreeMap<ResourceId, u32>,
}

/// An encoded element start.
struct Start {
    header: ElementStart,
    attributes: Vec<AttributeRecord>,
}

impl<'a> Pool<'a> {
    fn collect(root: &'a Element) -> Result<Self, XmlError> {
        let mut attrs: BTreeMap<ResourceId, &'a str> = BTreeMap::new();
        let mut others: Vec<&'a str> = Vec::new();
        for (depth, step) in walk(root) {
            match step {
                Step::Open(_) if depth >= MAX_DEPTH => return Err(XmlError::TooDeep),
                Step::Open(element) => {
                    for namespace in &element.namespaces {
                        others.extend(namespace.prefix.as_deref());
                        others.push(&namespace.uri);
                    }
                    others.extend(element.namespace.as_deref());
                    others.push(&element.name);
                    for attribute in &element.attributes {
                        match attribute_id(attribute)? {
                            Some(id) => _ = attrs.insert(id, &attribute.name.name),
                            None => others.push(&attribute.name.name),
                        }
                        others.extend(attribute.name.namespace.as_deref());
                        others.extend(attribute.value.as_str());
                    }
                }
                Step::Text(text) => others.push(text),
                Step::Close(_) => {}
            }
        }
        let mut strings = StringPoolBuilder::default();
        let ids = attrs.keys().copied().collect();
        let attr_index = attrs.into_iter().map(|(id, name)| (id, strings.push(name))).collect();
        // Names without IDs must not alias the resource-mapped prefix of the
        // pool, which was pushed without interning, so they are added again
        // even if the text repeats.
        others.into_iter().for_each(|s| _ = strings.intern(s));
        Ok(Self { strings, ids, attr_index })
    }

    fn other(&self, s: &str) -> u32 {
        self.strings.get(s).expect("collected every string of the document")
    }

    fn optional(&self, s: Option<&str>) -> u32 {
        s.map_or(NONE, |s| self.other(s))
    }

    fn element(&self, element: &Element, resources: &impl ResourceLookup) -> Result<Start, XmlError> {
        let mut attributes: Vec<(&AttrName, AttributeRecord)> = Vec::with_capacity(element.attributes.len());
        for attribute in &element.attributes {
            let name = match attribute_id(attribute)? {
                Some(id) => self.attr_index[&id],
                None => self.other(&attribute.name.name),
            };
            let typed_value = attribute.value.to_raw(
                |s| self.other(s),
                |reference| {
                    resolve(reference, resources)
                        .ok_or_else(|| XmlError::UnresolvedReference { reference: reference.clone() })
                },
            )?;
            let raw_value = if typed_value.data_type == data_type::STRING { typed_value.data } else { NONE };
            let namespace = self.optional(attribute.name.namespace.as_deref());
            attributes.push((&attribute.name, AttributeRecord { namespace, name, raw_value, typed_value }));
        }
        attributes.sort_by(|a, b| a.0.sort_key().cmp(&b.0.sort_key()));
        let index_of = |wanted: &dyn Fn(&AttrName) -> bool| {
            attributes.iter().position(|(name, _)| wanted(name)).map_or(0, |i| i as u16 + 1)
        };
        let header = ElementStart {
            namespace: self.optional(element.namespace.as_deref()),
            name: self.other(&element.name),
            attribute_start: ATTRIBUTE_SIZE,
            attribute_size: ATTRIBUTE_SIZE,
            attribute_count: u16::try_from(attributes.len()).map_err(|_| ChunkError::TooLarge)?,
            id_index: index_of(&|name| name.android_def() == Some(attr::ID)),
            class_index: index_of(&|name| name.namespace.is_none() && name.name == "class"),
            style_index: index_of(&|name| name.namespace.is_none() && name.name == "style"),
        };
        Ok(Start { header, attributes: attributes.into_iter().map(|(_, record)| record).collect() })
    }
}

/// The namespace of attributes an app defines for itself, and the prefix
/// of the per-package form older tools used.
const RES_AUTO_NAMESPACE: &str = "http://schemas.android.com/apk/res-auto";
const PACKAGE_NAMESPACE_PREFIX: &str = "http://schemas.android.com/apk/res/";

/// The resource ID `attribute` compiles to, if any.
///
/// A framework attribute must exist and hold a value of a format it
/// accepts. An app-defined attribute must already carry its ID, as in a
/// decoded document: this crate does not compile `<attr>` resources, and
/// without an ID Android would never see the attribute.
fn attribute_id(attribute: &Attribute) -> Result<Option<ResourceId>, XmlError> {
    let name = &attribute.name;
    if let Some(def) = name.android_def() {
        if !def.accepts(&attribute.value) {
            let value = attribute.value.clone();
            return Err(XmlError::MistypedValue { attribute: def.name(), value, format: def.format() });
        }
        return Ok(Some(def.id()));
    }
    match name.namespace.as_deref() {
        // A decoded document may carry IDs of private framework attributes.
        _ if name.id.is_some() => Ok(name.id),
        Some(ANDROID_NAMESPACE) => Err(XmlError::UnknownAttribute { name: name.name.clone() }),
        Some(namespace) if namespace == RES_AUTO_NAMESPACE || namespace.starts_with(PACKAGE_NAMESPACE_PREFIX) => {
            Err(XmlError::AppAttribute { namespace: namespace.to_owned(), name: name.name.clone() })
        }
        _ => Ok(None),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::res::chunk::Chunk;
    use crate::res::{ResourceRef, StringPool, Value};
    use crate::xml::Document;

    /// The string pool and resource map of an encoded document.
    fn pool_and_map(bytes: &[u8]) -> (Vec<String>, Vec<u32>) {
        let root = Chunk::parse_root(bytes, ChunkType::XML).unwrap();
        let mut children = root.children();
        let pool = StringPool::parse(&children.next().unwrap().unwrap()).unwrap();
        let strings = pool.iter().map(|s| s.unwrap().into_owned()).collect();
        let map = children.next().unwrap().unwrap();
        (strings, map.body.as_chunks::<4>().0.iter().map(|&b| u32::from_le_bytes(b)).collect())
    }

    /// The `ResXMLTree_attrExt` of the root element.
    fn root_start(bytes: &[u8]) -> ElementStart {
        let root = Chunk::parse_root(bytes, ChunkType::XML).unwrap();
        let start = root.children().map(Result::unwrap).find(|c| c.ty == ChunkType::XML_START_ELEMENT).unwrap();
        crate::res::chunk::read_padded(start.body)
    }

    #[test]
    fn resource_mapped_names_lead_the_pool_in_id_order() {
        let element = Element::new("activity")
            .with_android_namespace()
            .with_attr(attr::EXPORTED, true)
            .with_attr("name", "plain")
            .with_attr(attr::NAME, "android.app.NativeActivity")
            .with_attr(attr::LABEL, "Hello");
        let (strings, ids) = pool_and_map(&Document::new(element).encode(&()).unwrap());
        assert_eq!(ids, [attr::LABEL.id().0, attr::NAME.id().0, attr::EXPORTED.id().0]);
        assert_eq!(strings[..3], ["label", "name", "exported"]);
        // The plain `name` attribute gets its own, unmapped pool entry.
        assert!(strings[3..].iter().any(|s| s == "name"));
    }

    #[test]
    fn records_the_id_class_and_style_attributes() {
        let style: ResourceRef = "@android:style/Theme.Material".parse().unwrap();
        let element = Element::new("View")
            .with_android_namespace()
            .with_attr(attr::ID, ResourceRef::Id(ResourceId(0x7f08_0001)))
            .with_attr(attr::LAYOUT_WIDTH, -1)
            .with_attr("style", style)
            .with_attr("class", "org.example.View");
        let start = root_start(&Document::new(element).encode(&()).unwrap());
        // Attributes are ordered by ID (id, then layout_width), then by name (class, style).
        assert_eq!((start.id_index, start.class_index, start.style_index), (1, 3, 4));
    }

    #[test]
    fn unknown_framework_attributes_are_rejected() {
        let mut element = Element::new("application");
        element.set_attr(
            AttrName { namespace: Some(ANDROID_NAMESPACE.into()), name: "bogusAttribute".into(), id: None },
            true,
        );
        assert!(matches!(
            Document::new(element).encode(&()),
            Err(XmlError::UnknownAttribute { name }) if name == "bogusAttribute"
        ));
    }

    #[test]
    fn values_must_fit_the_attribute_format() {
        // `TypedArray.getBoolean` would read these as `false`, behind the
        // back of every check that reads the manifest.
        for value in [Value::String("false".into()), Value::Int(0)] {
            let element = Element::new("application").with_attr(attr::HAS_CODE, value);
            assert!(matches!(
                Document::new(element).encode(&()),
                Err(XmlError::MistypedValue { attribute: "hasCode", .. })
            ));
        }
        let reference = Element::new("application").with_attr(attr::HAS_CODE, ResourceRef::Id(ResourceId(0x7f01_0000)));
        assert!(Document::new(reference).encode(&()).is_ok());
    }

    #[test]
    fn app_attributes_without_ids_are_rejected() {
        for namespace in ["http://schemas.android.com/apk/res-auto", "http://schemas.android.com/apk/res/org.example"] {
            let name = AttrName { namespace: Some(namespace.into()), name: "tint".into(), id: None };
            let element = Element::new("ImageView").with_attr(name.clone(), "#fff");
            assert!(matches!(Document::new(element).encode(&()), Err(XmlError::AppAttribute { .. })));
            let compiled = AttrName { id: Some(ResourceId(0x7f01_0000)), ..name };
            assert!(Document::new(Element::new("ImageView").with_attr(compiled, "#fff")).encode(&()).is_ok());
        }
    }

    #[test]
    fn attributes_with_the_same_id_are_one_attribute() {
        let renamed =
            AttrName { namespace: Some(ANDROID_NAMESPACE.into()), name: "a".into(), id: Some(attr::LABEL.id()) };
        let element = Element::new("application").with_attr(renamed, "old").with_attr(attr::LABEL, "new");
        assert_eq!(element.attributes.len(), 1);
        assert_eq!(element.android_attr(attr::LABEL), Some(&Value::String("new".into())));
    }

    #[test]
    fn unresolved_app_references_are_rejected() {
        let element =
            Element::new("application").with_attr(attr::LABEL, "@string/missing".parse::<ResourceRef>().unwrap());
        assert!(matches!(Document::new(element).encode(&()), Err(XmlError::UnresolvedReference { .. })));
    }

    #[test]
    fn depth_is_limited_without_recursion() {
        let mut element = Element::new("leaf");
        for _ in 0..MAX_DEPTH {
            element = Element::new("node").with_child(element);
        }
        assert!(matches!(Document::new(element).encode(&()), Err(XmlError::TooDeep)));
    }
}