add export sources

Author: zsqw123

Cargo.lock

@@ -43,5 +43,6 @@ parking_lot = { version = "0.12.3", features = ["deadlock_detection"] }
 
 fontdb = "0.23.0"
 rfd = "0.15.0"
+open = "5.3.2"
 
 log = { version = "0.4", features = ["std"] }

Cargo.toml

@@ -2,10 +2,10 @@ pub use raw::*;
 pub mod element;
 
 use crate::dex::element::{AsElement, ClassContentElement};
-use crate::impls::jvms::r::{ReadContext, U32BasedSize};
+use crate::impls::jvms::r::{ReadContext, ReadFrom, U32BasedSize};
 use crate::impls::ToArc;
 use crate::smali::SmaliNode;
-use crate::{AsmErr, AsmResult};
+use crate::{AsmErr, AsmResult, StrRef};
 pub use constant::*;
 use std::io::Read;
 pub use util::*;
@@ -29,20 +29,21 @@ impl DexFile {
     }
     pub fn resolve_from_bytes(bytes: &[u8]) -> AsmResult<Self> {
         let mut context = ReadContext::little_endian(bytes);
-        context.read()
+        DexFile::read_from(&mut context)
     }
 }
 
 pub struct DexFileAccessor {
     pub file: DexFile,
+    pub file_name: StrRef,
     pub bytes: Vec<u8>,
     pub endian: bool,
     pub call_site_ids: Vec<CallSiteId>,
     pub method_handles: Vec<MethodHandle>,
 }
 
 impl DexFileAccessor {
-    pub fn new(file: DexFile, bytes: Vec<u8>) -> Self {
+    pub fn new(file: DexFile, bytes: Vec<u8>, file_name: StrRef) -> Self {
         let endian = file.header.endian_tag == Header::BIG_ENDIAN_TAG;
         let map_list = Self::get_map_list(&bytes, &file.header, endian)
             .unwrap_or_default();
@@ -65,7 +66,7 @@ impl DexFileAccessor {
         }
         let call_site_ids = Self::get_call_site_ids(&bytes, call_site_off, call_site_size, endian);
         let method_handles = Self::get_method_handles(&bytes, method_handle_off, method_handle_size, endian);
-        Self { file, bytes, endian, call_site_ids, method_handles }
+        Self { file, bytes, endian, call_site_ids, method_handles, file_name }
     }
 
     pub fn get_class_element(&self, class_data_off: DUInt) -> AsmResult<ClassContentElement> {

asm/src/dex/mod.rs

@@ -166,8 +166,8 @@ impl ReadFrom for EncodedCatchHandler {
     }
 }
 
-impl ReadFrom for DexFile {
-    fn read_from(context: &mut ReadContext) -> AsmResult<Self> {
+impl DexFile {
+    pub(crate) fn read_from(context: &mut ReadContext) -> AsmResult<Self> {
         let header: Header = context.read()?;
         let string_ids = context.read_vec(header.string_ids_size)?;
         let type_ids = context.read_vec(header.type_ids_size)?;

asm/src/impls/dex/r/raw.rs

@@ -12,6 +12,8 @@ pub struct SmaliNode {
 
 #[derive(Debug, Clone, Hash, Eq, PartialEq)]
 pub enum SmaliToken {
+    SourceInfo(StrRef),
+    
     Raw(ConstStr),
     Op(ConstStr),
 
@@ -125,6 +127,7 @@ pub fn tokens_to_raw(tokens: &[SmaliToken]) -> String {
 impl SmaliToken {
     pub fn raw(&self) -> String {
         match self {
+            Self::SourceInfo(source) => format!("# source: {source}"),
             Self::Raw(tag) => tag.to_string(),
             Self::Op(op) => op.to_string(),
             Self::LineStartOffsetMarker { raw, .. } => raw.clone(),

asm/src/smali.rs

@@ -1,6 +1,7 @@
 use java_asm::dex::{DexFile, DexFileAccessor};
 use std::sync::Arc;
 use std::time::Instant;
+use java_asm::StrRef;
 
 #[test]
 fn read_dex_test() {
@@ -30,5 +31,5 @@ fn read_test_dex_file() -> DexFileAccessor {
     let dex_file_bytes = include_bytes!("../res/dex/classes14.dex");
     let dex_file = DexFile::resolve_from_bytes(dex_file_bytes).unwrap();
     println!("Dex file resolved in {:?}", start.elapsed());
-    DexFileAccessor::new(dex_file, dex_file_bytes.to_vec())
+    DexFileAccessor::new(dex_file, dex_file_bytes.to_vec(), StrRef::from("classes14.dex"))
 }

asm/tests/dex/read_test.rs

@@ -2,6 +2,7 @@ use eframe::epaint::Color32;
 use egui::text::LayoutJob;
 use egui::{CursorIcon, FontId, Response, ScrollArea, TextStyle, Ui, Vec2};
 use java_asm::smali::SmaliToken;
+use java_asm::StrRef;
 use java_asm_server::ui::{AppContainer, Content, Top};
 use java_asm_server::AsmServer;
 use std::ops::{Deref, DerefMut};
@@ -104,6 +105,20 @@ impl<'a> RenderContext<'a> {
         } = self;
         let dft_color = *dft_color;
         match token {
+            SmaliToken::SourceInfo(source_info) => {
+                let text = format!("#from: {source_info}");
+                let text_ui = simple_text(ui, text, font, smali_style.literal)
+                    .on_hover_cursor(CursorIcon::PointingHand);
+                let text_ui = text_ui
+                    .on_hover_ui(|ui| {
+                        ui.style_mut().interaction.selectable_labels = true;
+                        self.source_file_info_menu(ui, source_info);
+                    });
+                text_ui.context_menu(|ui| {
+                    self.source_file_info_menu(ui, source_info);
+                });
+                text_ui
+            },
             SmaliToken::Raw(s) => simple_text(ui, s.to_string(), font, dft_color),
             SmaliToken::Op(s) => simple_text(ui, s.to_string(), font, smali_style.op),
             SmaliToken::LineStartOffsetMarker { raw, .. } => {
@@ -139,6 +154,17 @@ impl<'a> RenderContext<'a> {
         }
     }
 
+    fn source_file_info_menu(
+        &mut self, ui: &mut Ui, source_file_info: &StrRef,
+    ) {
+        ui.horizontal(|ui| {
+            let link = ui.link(format!("Export Source: {source_file_info}"));
+            if link.clicked() {
+                self.server.dialog_to_save_file(source_file_info);
+            }
+        });
+    }
+
     fn descriptor_menu(
         &mut self, ui: &mut Ui, descriptor: &str,
     ) {

asm_egui/src/smali.rs

@@ -17,6 +17,7 @@ log = { workspace = true }
 
 fontdb = { workspace = true }
 rfd = { workspace = true }
+open = { workspace = true }
 
 enum_dispatch = { workspace = true }
 

asm_server/Cargo.toml

@@ -1,8 +1,8 @@
 use crate::impls::server::{ProgressMessage, ServerMessage};
 use crate::server::OpenFileError;
-use crate::Accessor;
+use crate::{Accessor, ExportableSource};
 use java_asm::dex::{ClassDef, DexFile, DexFileAccessor};
-use java_asm::smali::SmaliNode;
+use java_asm::smali::{stb, SmaliNode, SmaliToken};
 use java_asm::{DescriptorRef, StrRef};
 use log::{error, warn};
 use std::collections::HashMap;
@@ -13,6 +13,7 @@ use zip::ZipArchive;
 
 pub struct ApkAccessor {
     pub map: HashMap<DescriptorRef, ClassPosition>,
+    pub dex_sources: HashMap<StrRef, Arc<DexFileAccessor>>,
 }
 
 type ClassPosition = (Arc<DexFileAccessor>, ClassDef);
@@ -37,12 +38,18 @@ pub async fn read_apk(
 
     // put dex files
     let dex_file_count = dex_files.len();
-    let dex_files = dex_files.iter().map(|name| {
-        let mut file = zip_archive.by_index(*name).map_err(OpenFileError::LoadZip)?;
+    let mut dex_sources = HashMap::new();
+    let dex_files = dex_files.iter().map(|entry_index| {
+        let mut file = zip_archive.by_index(*entry_index).map_err(OpenFileError::LoadZip)?;
+        let file_name = StrRef::from(file.name());
+        let file_name_for_dex_sources = file_name.clone();
         let mut bytes = Vec::new();
         file.read_to_end(&mut bytes).map_err(OpenFileError::Io)?;
         let dex_file = DexFile::resolve_from_bytes(&bytes).map_err(OpenFileError::ResolveError)?;
-        Ok(Arc::new(DexFileAccessor::new(dex_file, bytes)))
+        let dex_file_accessor = Arc::new(DexFileAccessor::new(dex_file, bytes, file_name));
+
+        dex_sources.insert(file_name_for_dex_sources, dex_file_accessor.clone());
+        Ok(dex_file_accessor)
     }).map(|res: Result<Arc<DexFileAccessor>, OpenFileError>| {
         match res {
             Ok(dex_file) => Some(dex_file),
@@ -73,7 +80,7 @@ pub async fn read_apk(
     };
     map.shrink_to_fit();
     send_loaded(&sender).await;
-    Ok(ApkAccessor { map })
+    Ok(ApkAccessor { map, dex_sources })
 }
 
 async fn send_progress(
@@ -118,10 +125,34 @@ impl Accessor for ApkAccessor {
     fn read_content(&self, class_key: &str) -> Option<SmaliNode> {
         let class_position = self.map.get(class_key);
         if let Some((accessor, class_def)) = class_position {
-            accessor.get_class_smali(*class_def).ok()
+            let dex_file_name = accessor.file_name.clone();
+            let smali_node = accessor.get_class_smali(*class_def).ok();
+            let Some(smali_node) = smali_node else {
+                warn!("No class content found for: {}", class_key);
+                return None;
+            };
+            let mut smali_node = smali_node;
+            let source_tag_smali = stb().push(SmaliToken::SourceInfo(dex_file_name)).s();
+            smali_node.children.insert(0, source_tag_smali);
+            Some(smali_node)
         } else {
             warn!("No class content found for: {}", class_key);
             None
         }
     }
+
+    // source key is the dex file name.
+    fn peek_source(&self, source_key: &str) -> Option<ExportableSource> {
+        let dex_source = self.dex_sources.get(source_key);
+        let Some(dex_source) = dex_source else {
+            warn!("No source found for: {source_key} when trying peek source.");
+            return None;
+        };
+        let file_name = dex_source.file_name.clone();
+        let source = dex_source.bytes.clone();
+        Some(ExportableSource {
+            exportable_name: file_name,
+            source,
+        })
+    }
 }

asm_server/src/impls/apk_load.rs

@@ -48,6 +48,13 @@ pub trait Accessor {
     /// the format of class_key is [DescriptorRef]
     fn exist_class(&self, class_key: &str) -> bool;
     fn read_content(&self, class_key: &str) -> Option<SmaliNode>;
+
+    fn peek_source(&self, source_key: &str) -> Option<ExportableSource>;
+}
+
+pub struct ExportableSource {
+    pub exportable_name: StrRef,
+    pub source: Vec<u8>,
 }
 
 pub struct MethodMeta {