aconfig: add flag table offset query function

Bug: b/321077378
Test: atest aconfig.test; atest aconfig_storage_file.test
Change-Id: Ib0ec1ec809c65d8f9f1284e4214cfbb683812f1d
diff --git a/tools/aconfig/aconfig/src/storage/flag_table.rs b/tools/aconfig/aconfig/src/storage/flag_table.rs
index 70878a8..bebac890 100644
--- a/tools/aconfig/aconfig/src/storage/flag_table.rs
+++ b/tools/aconfig/aconfig/src/storage/flag_table.rs
@@ -17,7 +17,7 @@
 use crate::commands::assign_flag_ids;
 use crate::storage::FlagPackage;
 use aconfig_storage_file::{
-    get_bucket_index, get_table_size, FlagTable, FlagTableHeader, FlagTableNode, FILE_VERSION,
+    get_table_size, FlagTable, FlagTableHeader, FlagTableNode, FILE_VERSION,
 };
 use anyhow::{anyhow, Result};
 
@@ -39,8 +39,7 @@
     flag_id: u16,
     num_buckets: u32,
 ) -> FlagTableNode {
-    let full_flag_name = package_id.to_string() + "/" + flag_name;
-    let bucket_index = get_bucket_index(&full_flag_name, num_buckets);
+    let bucket_index = FlagTableNode::find_bucket_index(package_id, flag_name, num_buckets);
     FlagTableNode {
         package_id,
         flag_name: flag_name.to_string(),
diff --git a/tools/aconfig/aconfig_storage_file/src/flag_table.rs b/tools/aconfig/aconfig_storage_file/src/flag_table.rs
index dab752a..421f847 100644
--- a/tools/aconfig/aconfig_storage_file/src/flag_table.rs
+++ b/tools/aconfig/aconfig_storage_file/src/flag_table.rs
@@ -17,8 +17,8 @@
 //! flag table module defines the flag table file format and methods for serialization
 //! and deserialization
 
-use crate::{read_str_from_bytes, read_u16_from_bytes, read_u32_from_bytes};
-use anyhow::Result;
+use crate::{read_str_from_bytes, read_u16_from_bytes, read_u32_from_bytes, get_bucket_index};
+use anyhow::{anyhow, Result};
 
 /// Flag table header struct
 #[derive(PartialEq, Debug)]
@@ -86,9 +86,9 @@
     }
 
     /// Deserialize from bytes
-    pub fn from_bytes(bytes: &[u8], num_buckets: u32) -> Result<Self> {
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
         let mut head = 0;
-        let mut node = Self {
+        let node = Self {
             package_id: read_u32_from_bytes(bytes, &mut head)?,
             flag_name: read_str_from_bytes(bytes, &mut head)?,
             flag_type: read_u16_from_bytes(bytes, &mut head)?,
@@ -99,10 +99,14 @@
             },
             bucket_index: 0,
         };
-        let full_flag_name = node.package_id.to_string() + "/" + &node.flag_name;
-        node.bucket_index = crate::get_bucket_index(&full_flag_name, num_buckets);
         Ok(node)
     }
+
+    /// Calculate node bucket index
+    pub fn find_bucket_index(package_id: u32, flag_name: &str, num_buckets: u32) -> u32 {
+        let full_flag_name = package_id.to_string() + "/" + flag_name;
+        get_bucket_index(&full_flag_name, num_buckets)
+    }
 }
 
 #[derive(PartialEq, Debug)]
@@ -138,8 +142,10 @@
             .collect();
         let nodes = (0..num_flags)
             .map(|_| {
-                let node = FlagTableNode::from_bytes(&bytes[head..], num_buckets).unwrap();
+                let mut node = FlagTableNode::from_bytes(&bytes[head..]).unwrap();
                 head += node.as_bytes().len();
+                node.bucket_index = FlagTableNode::find_bucket_index(
+                    node.package_id, &node.flag_name, num_buckets);
                 node
             })
             .collect();
@@ -149,6 +155,42 @@
     }
 }
 
+/// Query flag within package offset
+pub fn find_flag_offset(buf: &[u8], package_id: u32, flag: &str) -> Result<Option<u16>> {
+    let interpreted_header = FlagTableHeader::from_bytes(buf)?;
+    if interpreted_header.version > crate::FILE_VERSION {
+        return Err(anyhow!(
+            "Cannot read storage file with a higher version of {} with lib version {}",
+            interpreted_header.version,
+            crate::FILE_VERSION
+        ));
+    }
+
+    let num_buckets = (interpreted_header.node_offset - interpreted_header.bucket_offset) / 4;
+    let bucket_index = FlagTableNode::find_bucket_index(package_id, flag, num_buckets);
+
+    let mut pos = (interpreted_header.bucket_offset + 4 * bucket_index) as usize;
+    let mut flag_node_offset = read_u32_from_bytes(buf, &mut pos)? as usize;
+    if flag_node_offset < interpreted_header.node_offset as usize
+        || flag_node_offset >= interpreted_header.file_size as usize
+    {
+        return Ok(None);
+    }
+
+    loop {
+        let interpreted_node = FlagTableNode::from_bytes(&buf[flag_node_offset..])?;
+        if interpreted_node.package_id == package_id &&
+            interpreted_node.flag_name == flag {
+            return Ok(Some(interpreted_node.flag_id));
+        }
+        match interpreted_node.next_offset {
+            Some(offset) => flag_node_offset = offset as usize,
+            None => return Ok(None),
+        }
+    }
+
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -228,13 +270,70 @@
         let nodes: &Vec<FlagTableNode> = &flag_table.nodes;
         let num_buckets = crate::get_table_size(header.num_flags).unwrap();
         for node in nodes.iter() {
-            let reinterpreted_node = FlagTableNode::from_bytes(&node.as_bytes(), num_buckets);
-            assert!(reinterpreted_node.is_ok());
-            assert_eq!(node, &reinterpreted_node.unwrap());
+            let mut reinterpreted_node = FlagTableNode::from_bytes(&node.as_bytes()).unwrap();
+            reinterpreted_node.bucket_index = FlagTableNode::find_bucket_index(
+                reinterpreted_node.package_id,
+                &reinterpreted_node.flag_name,
+                num_buckets
+            );
+            assert_eq!(node, &reinterpreted_node);
         }
 
         let reinterpreted_table = FlagTable::from_bytes(&flag_table.as_bytes());
         assert!(reinterpreted_table.is_ok());
         assert_eq!(&flag_table, &reinterpreted_table.unwrap());
     }
+
+    #[test]
+    // this test point locks down table query
+    fn test_flag_query() {
+        let flag_table = create_test_flag_table().unwrap().as_bytes();
+        let baseline = vec![
+            (0, "enabled_ro", 1u16),
+            (0, "enabled_rw", 2u16),
+            (1, "disabled_ro", 0u16),
+            (2, "enabled_ro", 1u16),
+            (1, "enabled_fixed_ro", 1u16),
+            (1, "enabled_ro", 2u16),
+            (2, "enabled_fixed_ro", 0u16),
+            (0, "disabled_rw", 0u16),
+        ];
+        for (package_id, flag_name, expected_offset) in baseline.into_iter() {
+            let flag_offset =
+                find_flag_offset(&flag_table[..], package_id, flag_name)
+                .unwrap()
+                .unwrap();
+            assert_eq!(flag_offset, expected_offset);
+        }
+    }
+
+    #[test]
+    // this test point locks down table query of a non exist flag
+    fn test_not_existed_flag_query() {
+        let flag_table = create_test_flag_table().unwrap().as_bytes();
+        let flag_offset =
+            find_flag_offset(&flag_table[..], 1, "disabled_fixed_ro").unwrap();
+        assert_eq!(flag_offset, None);
+        let flag_offset =
+            find_flag_offset(&flag_table[..], 2, "disabled_rw").unwrap();
+        assert_eq!(flag_offset, None);
+    }
+
+    #[test]
+    // this test point locks down query error when file has a higher version
+    fn test_higher_version_storage_file() {
+        let mut table = create_test_flag_table().unwrap();
+        table.header.version = crate::FILE_VERSION + 1;
+        let flag_table = table.as_bytes();
+        let error = find_flag_offset(&flag_table[..], 0, "enabled_ro")
+            .unwrap_err();
+        assert_eq!(
+            format!("{:?}", error),
+            format!(
+                "Cannot read storage file with a higher version of {} with lib version {}",
+                crate::FILE_VERSION + 1,
+                crate::FILE_VERSION
+            )
+        );
+    }
 }