aflags: read protos from all containers

Create one library for reading protos from all containers, instead of
having numerous libraries perform the same logic. For Java, we will
create a similar library reusing the same
partition_aconfig_flags_paths.txt.

Bug: 324436145
Test: adb shell aflags list # Confirm that various containers appear
Change-Id: I924e281a50f9a609e1c07c03267eebe3dce52752
diff --git a/tools/aconfig/Cargo.toml b/tools/aconfig/Cargo.toml
index 6bd0d06..bf5e1a9 100644
--- a/tools/aconfig/Cargo.toml
+++ b/tools/aconfig/Cargo.toml
@@ -2,6 +2,7 @@
 
 members = [
     "aconfig",
+    "aconfig_device_paths",
     "aconfig_protos",
     "aconfig_storage_file",
     "aconfig_storage_read_api",
diff --git a/tools/aconfig/aconfig_device_paths/Android.bp b/tools/aconfig/aconfig_device_paths/Android.bp
new file mode 100644
index 0000000..21aa9a9
--- /dev/null
+++ b/tools/aconfig/aconfig_device_paths/Android.bp
@@ -0,0 +1,38 @@
+// Copyright (C) 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_defaults {
+    name: "libaconfig_device_paths.defaults",
+    edition: "2021",
+    clippy_lints: "android",
+    lints: "android",
+    srcs: ["src/lib.rs"],
+    rustlibs: [
+        "libaconfig_protos",
+        "libanyhow",
+        "libprotobuf",
+        "libregex",
+    ],
+}
+
+rust_library {
+    name: "libaconfig_device_paths",
+    crate_name: "aconfig_device_paths",
+    host_supported: true,
+    defaults: ["libaconfig_device_paths.defaults"],
+}
diff --git a/tools/aconfig/aconfig_device_paths/Cargo.toml b/tools/aconfig/aconfig_device_paths/Cargo.toml
new file mode 100644
index 0000000..dbe9b3a
--- /dev/null
+++ b/tools/aconfig/aconfig_device_paths/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "aconfig_device_paths"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.82"
diff --git a/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt b/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt
new file mode 100644
index 0000000..3d2deb2
--- /dev/null
+++ b/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt
@@ -0,0 +1,6 @@
+[
+    "/system/etc/aconfig_flags.pb",
+    "/system_ext/etc/aconfig_flags.pb",
+    "/product/etc/aconfig_flags.pb",
+    "/vendor/etc/aconfig_flags.pb",
+]
diff --git a/tools/aconfig/aconfig_device_paths/src/lib.rs b/tools/aconfig/aconfig_device_paths/src/lib.rs
new file mode 100644
index 0000000..7bb62f4
--- /dev/null
+++ b/tools/aconfig/aconfig_device_paths/src/lib.rs
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! Library for finding all aconfig on-device protobuf file paths.
+
+use anyhow::Result;
+use std::path::PathBuf;
+
+use std::fs;
+
+/// Determine all paths that contain an aconfig protobuf file.
+pub fn parsed_flags_proto_paths() -> Result<Vec<PathBuf>> {
+    let mut result: Vec<PathBuf> = include!("../partition_aconfig_flags_paths.txt")
+        .map(|s| PathBuf::from(s.to_string()))
+        .to_vec();
+    for dir in fs::read_dir("/apex")? {
+        let dir = dir?;
+
+        // Only scan the currently active version of each mainline module; skip the @version dirs.
+        if dir.file_name().as_encoded_bytes().iter().any(|&b| b == b'@') {
+            continue;
+        }
+
+        let mut path = PathBuf::from("/apex");
+        path.push(dir.path());
+        path.push("etc");
+        path.push("aconfig_flags.pb");
+        if path.exists() {
+            result.push(path);
+        }
+    }
+
+    Ok(result)
+}
diff --git a/tools/aconfig/aconfig_storage_write_api/src/lib.rs b/tools/aconfig/aconfig_storage_write_api/src/lib.rs
index 8b7e459..7148d06 100644
--- a/tools/aconfig/aconfig_storage_write_api/src/lib.rs
+++ b/tools/aconfig/aconfig_storage_write_api/src/lib.rs
@@ -443,10 +443,12 @@
             .unwrap();
             for i in 0..8 {
                 set_flag_is_sticky(&mut file, FlagValueType::Boolean, i, true).unwrap();
-                let attribute = get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+                let attribute =
+                    get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
                 assert!((attribute & (FlagInfoBit::IsSticky as u8)) != 0);
                 set_flag_is_sticky(&mut file, FlagValueType::Boolean, i, false).unwrap();
-                let attribute = get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+                let attribute =
+                    get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
                 assert!((attribute & (FlagInfoBit::IsSticky as u8)) == 0);
             }
         }
@@ -485,10 +487,12 @@
             .unwrap();
             for i in 0..8 {
                 set_flag_has_override(&mut file, FlagValueType::Boolean, i, true).unwrap();
-                let attribute = get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+                let attribute =
+                    get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
                 assert!((attribute & (FlagInfoBit::HasOverride as u8)) != 0);
                 set_flag_has_override(&mut file, FlagValueType::Boolean, i, false).unwrap();
-                let attribute = get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+                let attribute =
+                    get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
                 assert!((attribute & (FlagInfoBit::HasOverride as u8)) == 0);
             }
         }
diff --git a/tools/aconfig/aflags/Android.bp b/tools/aconfig/aflags/Android.bp
index 4920a6f..2a02379 100644
--- a/tools/aconfig/aflags/Android.bp
+++ b/tools/aconfig/aflags/Android.bp
@@ -9,6 +9,7 @@
     lints: "android",
     srcs: ["src/main.rs"],
     rustlibs: [
+        "libaconfig_device_paths",
         "libaconfig_protos",
         "libaconfig_storage_read_api",
         "libaconfig_storage_file",
diff --git a/tools/aconfig/aflags/Cargo.toml b/tools/aconfig/aflags/Cargo.toml
index cce7f9d..eeae295 100644
--- a/tools/aconfig/aflags/Cargo.toml
+++ b/tools/aconfig/aflags/Cargo.toml
@@ -13,3 +13,4 @@
 aconfig_storage_file = { version = "0.1.0", path = "../aconfig_storage_file" }
 aconfig_storage_read_api = { version = "0.1.0", path = "../aconfig_storage_read_api" }
 clap = {version = "4.5.2" }
+aconfig_device_paths = { version = "0.1.0", path = "../aconfig_device_paths" }
diff --git a/tools/aconfig/aflags/src/device_config_source.rs b/tools/aconfig/aflags/src/device_config_source.rs
index 089f33d..cf6ab28 100644
--- a/tools/aconfig/aflags/src/device_config_source.rs
+++ b/tools/aconfig/aflags/src/device_config_source.rs
@@ -14,78 +14,17 @@
  * limitations under the License.
  */
 
-use crate::{Flag, FlagPermission, FlagSource, FlagValue, ValuePickedFrom};
-use aconfig_protos::ProtoFlagPermission as ProtoPermission;
-use aconfig_protos::ProtoFlagState as ProtoState;
-use aconfig_protos::ProtoParsedFlag;
-use aconfig_protos::ProtoParsedFlags;
+use crate::load_protos;
+use crate::{Flag, FlagSource, FlagValue, ValuePickedFrom};
+
 use anyhow::{anyhow, bail, Result};
 use regex::Regex;
-use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::process::Command;
-use std::{fs, str};
+use std::str;
 
 pub struct DeviceConfigSource {}
 
-fn convert_parsed_flag(flag: &ProtoParsedFlag) -> Flag {
-    let namespace = flag.namespace().to_string();
-    let package = flag.package().to_string();
-    let name = flag.name().to_string();
-
-    let container = if flag.container().is_empty() {
-        "system".to_string()
-    } else {
-        flag.container().to_string()
-    };
-
-    let value = match flag.state() {
-        ProtoState::ENABLED => FlagValue::Enabled,
-        ProtoState::DISABLED => FlagValue::Disabled,
-    };
-
-    let permission = match flag.permission() {
-        ProtoPermission::READ_ONLY => FlagPermission::ReadOnly,
-        ProtoPermission::READ_WRITE => FlagPermission::ReadWrite,
-    };
-
-    Flag {
-        namespace,
-        package,
-        name,
-        container,
-        value,
-        staged_value: None,
-        permission,
-        value_picked_from: ValuePickedFrom::Default,
-    }
-}
-
-fn read_pb_files() -> Result<Vec<Flag>> {
-    let mut flags: BTreeMap<String, Flag> = BTreeMap::new();
-    for partition in ["system", "system_ext", "product", "vendor"] {
-        let path = format!("/{partition}/etc/aconfig_flags.pb");
-        let Ok(bytes) = fs::read(&path) else {
-            eprintln!("warning: failed to read {}", path);
-            continue;
-        };
-        let parsed_flags: ProtoParsedFlags = protobuf::Message::parse_from_bytes(&bytes)?;
-        for flag in parsed_flags.parsed_flag {
-            let key = format!("{}.{}", flag.package(), flag.name());
-            let container = if flag.container().is_empty() {
-                "system".to_string()
-            } else {
-                flag.container().to_string()
-            };
-
-            if container.eq(partition) {
-                flags.insert(key, convert_parsed_flag(&flag));
-            }
-        }
-    }
-    Ok(flags.values().cloned().collect())
-}
-
 fn parse_device_config(raw: &str) -> Result<HashMap<String, FlagValue>> {
     let mut flags = HashMap::new();
     let regex = Regex::new(r"(?m)^([[[:alnum:]]_]+/[[[:alnum:]]_\.]+)=(true|false)$")?;
@@ -180,7 +119,7 @@
 
 impl FlagSource for DeviceConfigSource {
     fn list_flags() -> Result<Vec<Flag>> {
-        let pb_flags = read_pb_files()?;
+        let pb_flags = load_protos::load()?;
         let dc_flags = read_device_config_flags()?;
         let staged_flags = read_staged_flags()?;
 
diff --git a/tools/aconfig/aflags/src/load_protos.rs b/tools/aconfig/aflags/src/load_protos.rs
new file mode 100644
index 0000000..90d8599
--- /dev/null
+++ b/tools/aconfig/aflags/src/load_protos.rs
@@ -0,0 +1,62 @@
+use crate::{Flag, FlagPermission, FlagValue, ValuePickedFrom};
+use aconfig_protos::ProtoFlagPermission as ProtoPermission;
+use aconfig_protos::ProtoFlagState as ProtoState;
+use aconfig_protos::ProtoParsedFlag;
+use aconfig_protos::ProtoParsedFlags;
+use anyhow::Result;
+use std::fs;
+use std::path::Path;
+
+// TODO(b/329875578): use container field directly instead of inferring.
+fn infer_container(path: &Path) -> String {
+    let path_str = path.to_string_lossy();
+    path_str
+        .strip_prefix("/apex/")
+        .or_else(|| path_str.strip_prefix('/'))
+        .unwrap_or(&path_str)
+        .strip_suffix("/etc/aconfig_flags.pb")
+        .unwrap_or(&path_str)
+        .to_string()
+}
+
+fn convert_parsed_flag(path: &Path, flag: &ProtoParsedFlag) -> Flag {
+    let namespace = flag.namespace().to_string();
+    let package = flag.package().to_string();
+    let name = flag.name().to_string();
+
+    let value = match flag.state() {
+        ProtoState::ENABLED => FlagValue::Enabled,
+        ProtoState::DISABLED => FlagValue::Disabled,
+    };
+
+    let permission = match flag.permission() {
+        ProtoPermission::READ_ONLY => FlagPermission::ReadOnly,
+        ProtoPermission::READ_WRITE => FlagPermission::ReadWrite,
+    };
+
+    Flag {
+        namespace,
+        package,
+        name,
+        container: infer_container(path),
+        value,
+        staged_value: None,
+        permission,
+        value_picked_from: ValuePickedFrom::Default,
+    }
+}
+
+pub(crate) fn load() -> Result<Vec<Flag>> {
+    let mut result = Vec::new();
+
+    let paths = aconfig_device_paths::parsed_flags_proto_paths()?;
+    for path in paths {
+        let bytes = fs::read(path.clone())?;
+        let parsed_flags: ProtoParsedFlags = protobuf::Message::parse_from_bytes(&bytes)?;
+        for flag in parsed_flags.parsed_flag {
+            // TODO(b/334954748): enforce one-container-per-flag invariant.
+            result.push(convert_parsed_flag(&path, &flag));
+        }
+    }
+    Ok(result)
+}
diff --git a/tools/aconfig/aflags/src/main.rs b/tools/aconfig/aflags/src/main.rs
index 1c453c5..4ce0d35 100644
--- a/tools/aconfig/aflags/src/main.rs
+++ b/tools/aconfig/aflags/src/main.rs
@@ -25,6 +25,8 @@
 mod aconfig_storage_source;
 use aconfig_storage_source::AconfigStorageSource;
 
+mod load_protos;
+
 #[derive(Clone, PartialEq, Debug)]
 enum FlagPermission {
     ReadOnly,