diff --git a/tools/aconfig/aconfig/src/storage/flag_table.rs b/tools/aconfig/aconfig/src/storage/flag_table.rs index 70878a8b58..bebac890a0 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 @@ fn new_node( 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 dab752a006..421f84705d 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 @@ impl FlagTableNode { } /// Deserialize from bytes - pub fn from_bytes(bytes: &[u8], num_buckets: u32) -> Result { + pub fn from_bytes(bytes: &[u8]) -> Result { 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 @@ impl FlagTableNode { }, 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 @@ impl FlagTable { .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 @@ impl FlagTable { } } +/// Query flag within package offset +pub fn find_flag_offset(buf: &[u8], package_id: u32, flag: &str) -> Result> { + 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 @@ mod tests { let nodes: &Vec = &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 + ) + ); + } }