Improve storing of symbols

Signed-off-by: Mr-Kanister <[email protected]>

frontend/client/src/main/java/de/amosproj3/ziofa/client/Client.kt

@@ -74,7 +74,7 @@ interface Client {
 
     suspend fun getOdexFiles(pid: UInt): Flow<StringResponse>
 
-    suspend fun getSymbols(pid: UInt, odexFilePath: String): Flow<Symbol>
+    suspend fun getSymbols(odexFilePath: String): Flow<Symbol>
 
     suspend fun initStream(): Flow<Event>
 }

frontend/client/src/mock/java/de/amosproj3/ziofa/client/RustClient.kt

@@ -136,7 +136,7 @@ object RustClient : Client {
         )
     }
 
-    override suspend fun getSymbols(pid: UInt, odexFilePath: String): Flow<Symbol> = flow {
+    override suspend fun getSymbols(odexFilePath: String): Flow<Symbol> = flow {
         emit(
             Symbol(
                 method =

frontend/client/src/real/java/de.amosproj3.ziofa.client/RustClient.kt

@@ -110,8 +110,8 @@ class RustClient(private val inner: uniffi.client.Client) : Client {
     override suspend fun getOdexFiles(pid: UInt): Flow<StringResponse> =
         inner.getOdexFilesFlow(pid).mapNotNull { it.into() }
 
-    override suspend fun getSymbols(pid: UInt, odexFilePath: String): Flow<Symbol> =
-        inner.getSymbolFlow(pid, odexFilePath).mapNotNull { it.into() }
+    override suspend fun getSymbols(odexFilePath: String): Flow<Symbol> =
+        inner.getSymbolFlow(odexFilePath).mapNotNull { it.into() }
 
     override suspend fun initStream(): Flow<Event> = inner.initStreamFlow().mapNotNull { it.into() }
 }
@@ -161,8 +161,8 @@ fun uniffi.client.Client.getOdexFilesFlow(pid: UInt) = flow {
     }
 }
 
-fun uniffi.client.Client.getSymbolFlow(pid: UInt, odexFilePath: String) = flow {
-    getSymbols(pid, odexFilePath).use { stream ->
+fun uniffi.client.Client.getSymbolFlow(odexFilePath: String) = flow {
+    getSymbols(odexFilePath).use { stream ->
         while (true) {
             stream.next()?.also { symbol -> emit(symbol) } ?: break
         }

rust/backend/daemon/src/server.rs

@@ -132,13 +132,10 @@ impl Ziofa for ZiofaImpl {
 
         let (tx, rx) = mpsc::channel(4);
 
-        let symbol_handler_clone = self.symbol_handler.clone();
-
-        // let symbol_handler_guard = self.symbol_handler.lock().await;
-        // let odex_paths = symbol_handler_guard.get_odex_paths(pid)?;
+        let symbol_handler = self.symbol_handler.clone();
 
         tokio::spawn(async move {
-            let mut symbol_handler_guard = symbol_handler_clone.lock().await;
+            let mut symbol_handler_guard = symbol_handler.lock().await;
             // TODO Error Handling
             let odex_paths = match symbol_handler_guard.get_odex_paths(pid) {
                 Ok(paths) => paths,
@@ -169,17 +166,17 @@ impl Ziofa for ZiofaImpl {
         request: Request<GetSymbolsRequest>,
     ) -> Result<Response<Self::GetSymbolsStream>, Status> {
         let process_request = request.into_inner();
-        let pid = process_request.pid;
         let odex_file_path_string = process_request.odex_file_path;
         let odex_file_path = PathBuf::from(odex_file_path_string);
 
-        let symbol_handler_clone = self.symbol_handler.clone();
-
         let (tx, rx) = mpsc::channel(4);
+
+        let symbol_handler = self.symbol_handler.clone();
+
         tokio::spawn(async move {
-            let mut symbol_handler_guard = symbol_handler_clone.lock().await;
+            let mut symbol_handler_guard = symbol_handler.lock().await;
 
-            let symbol = match symbol_handler_guard.get_symbols(pid, &odex_file_path).await{
+            let symbol = match symbol_handler_guard.get_symbols(&odex_file_path).await{
                 Ok(symbol) => symbol,
                 Err(e) => {
                     tx.send(Err(Status::from(e)))

rust/backend/daemon/src/symbols.rs

@@ -8,7 +8,6 @@ use object::{Object, ObjectSymbol, ReadCache};
 use procfs::process::{MMapPath, Process};
 use procfs::ProcError;
 use serde::{Deserialize, Serialize};
-use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::fs::File;
 use std::io::Error;
@@ -31,8 +30,8 @@ pub enum SymbolError {
     SymbolsNotLoaded { pid: u32, odex_path: PathBuf },
     #[error(transparent)]
     SerdeError(#[from] serde_json::Error),
-    #[error("The desired odex file isn't available")]
-    OdexFileNotAvailable { pid: u32, odex_path: PathBuf },
+    #[error("The desired odex file isn't available. Did you call get_odex_files()?")]
+    OdexFileNotAvailable { odex_path: PathBuf },
 }
 
 impl From<SymbolError> for tonic::Status {
@@ -48,13 +47,16 @@ struct JsonSymbol {
 }
 
 pub struct SymbolHandler {
-    /// maps pid, odex file path and symbol name to offset
-    symbols: HashMap<u32, HashMap<PathBuf, HashMap<String, u64>>>,
+    /// maps pid to odex files
+    odex_files: HashMap<u32, HashSet<PathBuf>>,
+    /// maps odex file path and symbol name to offset
+    symbols: HashMap<PathBuf, HashMap<String, u64>>,
 }
 
 impl SymbolHandler {
     pub fn new() -> Self {
         SymbolHandler {
+            odex_files: HashMap::new(),
             symbols: HashMap::new(),
         }
     }
@@ -64,66 +66,50 @@ impl SymbolHandler {
     fn load_odex_paths(&mut self, pid: u32) -> Result<(), ProcError> {
         // if pid was already crawled, nothing it to do
         // TODO: Check for old/potentially outdated entries and reload them
-        if let Entry::Vacant(e) = self.symbols.entry(pid) {
-            e.insert(HashMap::new());
-        } else {
+        if self.odex_files.contains_key(&pid) {
             return Ok(());
         }
+        let odex_files = self.odex_files.entry(pid).or_default();
 
         let process = Process::new(i32::try_from(pid).unwrap())?;
         let maps = process.maps()?;
 
         // for each .odex file: insert a new hashmap into this pids entry of self.symbols
-        for map in maps
+        for odex in maps
             .iter()
             .filter_map(|mm_map| match mm_map.clone().pathname {
                 MMapPath::Path(path) => Some(path),
                 _ => None,
             })
             .filter(|path: &PathBuf| path.to_str().unwrap().ends_with(".odex"))
         {
-            self.symbols
-                .get_mut(&pid)
-                .unwrap()
-                .insert(map, HashMap::new());
+            odex_files.insert(odex.clone());
+            self.symbols.insert(odex, HashMap::new());
         }
 
         Ok(())
         // TODO: Remove old/long unused paths from cache
     }
 
-    pub fn get_odex_paths(&mut self, pid: u32) -> Result<HashSet<&PathBuf>, SymbolError> {
+    pub fn get_odex_paths(&mut self, pid: u32) -> Result<&HashSet<PathBuf>, SymbolError> {
         self.load_odex_paths(pid)?;
 
-        Ok(self
-            .symbols
+        self.odex_files
             .get(&pid)
-            .ok_or(SymbolError::OdexPathsNotLoaded { pid })?
-            .keys()
-            .collect())
+            .ok_or(SymbolError::OdexPathsNotLoaded { pid })
     }
 
-    async fn load_symbols(&mut self, pid: u32, odex_path: &PathBuf) -> Result<(), SymbolError> {
-        // make sure the needed data structures are there
-        self.load_odex_paths(pid)?;
-
-        // the following is in a code block as the immutable references of self.symbols and
-        // odex_to_symbol_map need to be dropped before getting mutable references to them below
-        {
-            // if the .odex file is not cached, throw error
-            let odex_to_symbol_map = self.symbols.get(&pid).unwrap();
-            if !odex_to_symbol_map.contains_key(odex_path) {
-                return Err(SymbolError::OdexFileNotAvailable {
-                    pid,
-                    odex_path: odex_path.to_path_buf(),
-                });
-            }
+    async fn load_symbols(&mut self, odex_path: &PathBuf) -> Result<(), SymbolError> {
+        // if the .odex file is not cached, throw error
+        if !self.symbols.contains_key(odex_path) {
+            return Err(SymbolError::OdexFileNotAvailable {
+                odex_path: odex_path.to_path_buf(),
+            });
+        }
 
-            // if the map already contains entries, nothing needs to be done
-            let symbol_to_offset = odex_to_symbol_map.get(odex_path).unwrap();
-            if !symbol_to_offset.is_empty() {
-                return Ok(());
-            }
+        // if the symbol map already contains entries, nothing needs to be done
+        if !self.symbols.get(odex_path).unwrap().is_empty() {
+            return Ok(());
         }
 
         self.generate_json_oatdump(odex_path)
@@ -132,9 +118,7 @@ impl SymbolHandler {
 
         let oatdata_section_offset = self.get_oatsection_address(odex_path).await?;
 
-        // get map from odex file path to symbol map
-        let odex_to_symbol_map = self.symbols.get_mut(&pid).unwrap();
-        let symbol_to_offset = odex_to_symbol_map.get_mut(odex_path).unwrap();
+        let symbols = self.symbols.get_mut(odex_path).unwrap();
 
         let json_file = tokio::fs::File::open(OATDUMP_PATH).await?;
         let json_reader = tokio::io::BufReader::new(json_file);
@@ -154,27 +138,19 @@ impl SymbolHandler {
             // the actual symbol offset is build from section offset and relative offset
             let offset = relative_offset + oatdata_section_offset;
 
-            symbol_to_offset.insert(symbol.method, offset);
+            symbols.insert(symbol.method, offset);
         }
 
         Ok(())
     }
 
     pub async fn get_symbols(
         &mut self,
-        pid: u32,
         odex_path: &PathBuf,
     ) -> Result<&HashMap<String, u64>, SymbolError> {
-        self.load_symbols(pid, odex_path).await?;
+        self.load_symbols(odex_path).await?;
 
-        self.symbols
-            .get(&pid)
-            .ok_or(SymbolError::OdexPathsNotLoaded { pid })?
-            .get(odex_path)
-            .ok_or(SymbolError::SymbolsNotLoaded {
-                pid,
-                odex_path: odex_path.to_path_buf(),
-            })
+        Ok(self.symbols.get(odex_path).unwrap())
     }
 
     async fn generate_json_oatdump(&self, path: &Path) -> Result<(), SymbolError> {

rust/client/src/bin/cli.rs

@@ -61,10 +61,6 @@ enum Commands {
 
     /// Get all symbols with their offsets
     Symbols {
-        /// Pid for which to get the symbols
-        #[arg(short, long)]
-        pid: u32,
-
         /// Path to the .odex file which should be crawled
         #[arg(short, long)]
         odex_file: String,
@@ -147,8 +143,8 @@ async fn get_odex_files(client: &mut Client, pid: u32, silent: bool) -> Result<(
     Ok(())
 }
 
-async fn get_symbols(client: &mut Client, pid: u32, odex_file: String, silent: bool) -> Result<()> {
-    let mut stream = client.get_symbols(pid, odex_file).await?;
+async fn get_symbols(client: &mut Client, odex_file: String, silent: bool) -> Result<()> {
+    let mut stream = client.get_symbols(odex_file).await?;
     let mut count: u32 = 0;
 
     while let Some(Ok(next)) = stream.next().await {
@@ -192,8 +188,8 @@ pub async fn main() -> anyhow::Result<()> {
         Commands::Odex { pid, silent } => {
             get_odex_files(&mut client, pid, silent).await?;
         },
-        Commands::Symbols { pid, odex_file, silent } => {
-            get_symbols(&mut client, pid, odex_file, silent).await?;
+        Commands::Symbols { odex_file, silent } => {
+            get_symbols(&mut client, odex_file, silent).await?;
         }
     }
 

rust/client/src/bindings.rs

@@ -161,10 +161,10 @@ impl Client {
         Ok(OdexFileStream(Mutex::new(Box::pin(stream))))
     }
 
-    pub async fn get_symbols(&self, pid: u32, odex_file: String) -> Result<SymbolStream> {
+    pub async fn get_symbols(&self, odex_file: String) -> Result<SymbolStream> {
         let mut guard = self.0.lock().await;
         let stream = guard
-            .get_symbols(pid, odex_file)
+            .get_symbols(odex_file)
             .await?
             .map(|x| x.map_err(ClientError::from));
 

rust/client/src/client.rs

@@ -3,17 +3,19 @@
 //
 // SPDX-License-Identifier: MIT
 
+use shared::ziofa::Event;
 use shared::{
     config::Configuration,
     counter::{counter_client::CounterClient, IfaceMessage},
-    ziofa::{ziofa_client::ZiofaClient, GetSymbolsRequest, PidMessage, Process, StringResponse, Symbol},
+    ziofa::{
+        ziofa_client::ZiofaClient, GetSymbolsRequest, PidMessage, Process, StringResponse, Symbol,
+    },
 };
 use tokio_stream::{Stream, StreamExt};
 use tonic::{
     transport::{Channel, Endpoint},
     Request,
 };
-use shared::ziofa::Event;
 
 #[derive(Clone, Debug)]
 pub struct Client {
@@ -100,18 +102,44 @@ impl Client {
     }
 
     pub async fn set_configuration(&mut self, configuration: Configuration) -> Result<u32> {
-        Ok(self.ziofa.set_configuration(configuration).await?.into_inner().response_type)
+        Ok(self
+            .ziofa
+            .set_configuration(configuration)
+            .await?
+            .into_inner()
+            .response_type)
     }
 
     pub async fn init_stream(&mut self) -> Result<impl Stream<Item = Result<Event>>> {
-        Ok(self.ziofa.init_stream(()).await?.into_inner().map(|s| Ok(s?)))
+        Ok(self
+            .ziofa
+            .init_stream(())
+            .await?
+            .into_inner()
+            .map(|s| Ok(s?)))
     }
 
-    pub async fn get_odex_files(&mut self, pid: u32) -> Result<impl Stream<Item = Result<StringResponse>>> {
-        Ok(self.ziofa.get_odex_files(PidMessage {pid}).await?.into_inner().map(|s| Ok(s?)))
+    pub async fn get_odex_files(
+        &mut self,
+        pid: u32,
+    ) -> Result<impl Stream<Item = Result<StringResponse>>> {
+        Ok(self
+            .ziofa
+            .get_odex_files(PidMessage { pid })
+            .await?
+            .into_inner()
+            .map(|s| Ok(s?)))
     }
 
-    pub async fn get_symbols(&mut self, pid: u32, odex_file_path: String) -> Result<impl Stream<Item = Result<Symbol>>> {
-        Ok(self.ziofa.get_symbols(GetSymbolsRequest {pid, odex_file_path}).await?.into_inner().map(|s| Ok(s?)))
+    pub async fn get_symbols(
+        &mut self,
+        odex_file_path: String,
+    ) -> Result<impl Stream<Item = Result<Symbol>>> {
+        Ok(self
+            .ziofa
+            .get_symbols(GetSymbolsRequest { odex_file_path })
+            .await?
+            .into_inner()
+            .map(|s| Ok(s?)))
     }
 }

rust/shared/proto/ziofa.proto

@@ -29,7 +29,6 @@ message StringResponse {
 }
 
 message GetSymbolsRequest {
-    uint32 pid = 1;
     string odex_file_path = 2;
 }