aflags: read from both system and mainline sockets

There are now two aconfigd daemons: one for the vendor/product/system partitions, and one for mainline partitions. Read from both of them, and combine their results.

Change-Id: I5a5e49661fbcf5c561bffbd226d9fa5aad34d053
Test: aflags list
Bug: 369808805
diff --git a/tools/aconfig/aflags/src/aconfig_storage_source.rs b/tools/aconfig/aflags/src/aconfig_storage_source.rs
index aef7d7e..766807a 100644
--- a/tools/aconfig/aflags/src/aconfig_storage_source.rs
+++ b/tools/aconfig/aflags/src/aconfig_storage_source.rs
@@ -17,6 +17,23 @@
 
 pub struct AconfigStorageSource {}
 
+static ACONFIGD_SYSTEM_SOCKET_NAME: &str = "/dev/socket/aconfigd_system";
+static ACONFIGD_MAINLINE_SOCKET_NAME: &str = "/dev/socket/aconfigd_mainline";
+
+enum AconfigdSocket {
+    System,
+    Mainline,
+}
+
+impl AconfigdSocket {
+    pub fn name(&self) -> &str {
+        match self {
+            AconfigdSocket::System => ACONFIGD_SYSTEM_SOCKET_NAME,
+            AconfigdSocket::Mainline => ACONFIGD_MAINLINE_SOCKET_NAME,
+        }
+    }
+}
+
 fn load_flag_to_container() -> Result<HashMap<String, String>> {
     Ok(load_protos::load()?.into_iter().map(|p| (p.qualified_name(), p.container)).collect())
 }
@@ -81,7 +98,7 @@
     })
 }
 
-fn read_from_socket() -> Result<Vec<ProtoFlagQueryReturnMessage>> {
+fn read_from_socket(socket: AconfigdSocket) -> Result<Vec<ProtoFlagQueryReturnMessage>> {
     let messages = ProtoStorageRequestMessages {
         msgs: vec![ProtoStorageRequestMessage {
             msg: Some(ProtoStorageRequestMessageMsg::ListStorageMessage(ProtoListStorageMessage {
@@ -93,8 +110,7 @@
         special_fields: SpecialFields::new(),
     };
 
-    let socket_name = "/dev/socket/aconfigd_system";
-    let mut socket = UnixStream::connect(socket_name)?;
+    let mut socket = UnixStream::connect(socket.name())?;
 
     let message_buffer = messages.write_to_bytes()?;
     let mut message_length_buffer: [u8; 4] = [0; 4];
@@ -128,14 +144,20 @@
 impl FlagSource for AconfigStorageSource {
     fn list_flags() -> Result<Vec<Flag>> {
         let containers = load_flag_to_container()?;
-        read_from_socket()
-            .map(|query_messages| {
-                query_messages
-                    .iter()
-                    .map(|message| convert(message.clone(), &containers))
-                    .collect::<Vec<_>>()
-            })?
+        let system_messages = read_from_socket(AconfigdSocket::System);
+        let mainline_messages = read_from_socket(AconfigdSocket::Mainline);
+
+        let mut all_messages = vec![];
+        if let Ok(system_messages) = system_messages {
+            all_messages.extend_from_slice(&system_messages);
+        }
+        if let Ok(mainline_messages) = mainline_messages {
+            all_messages.extend_from_slice(&mainline_messages);
+        }
+
+        all_messages
             .into_iter()
+            .map(|query_message| convert(query_message.clone(), &containers))
             .collect()
     }