Take ownership of inherited FDs using the rustutils crate

Bug: 3259955
Test: watch TH
Merged-In: I316a56142401b5f9bcb8ed350280019d7ddab123
Change-Id: I316a56142401b5f9bcb8ed350280019d7ddab123
diff --git a/android/fd_server/Android.bp b/android/fd_server/Android.bp
index b02c104..748c660 100644
--- a/android/fd_server/Android.bp
+++ b/android/fd_server/Android.bp
@@ -18,6 +18,7 @@
         "liblog_rust",
         "libnix",
         "librpcbinder_rs",
+        "librustutils",
     ],
     prefer_rlib: true,
     apex_available: ["com.android.virt"],
@@ -39,6 +40,7 @@
         "liblog_rust",
         "libnix",
         "librpcbinder_rs",
+        "librustutils",
     ],
     prefer_rlib: true,
     test_suites: ["general-tests"],
diff --git a/android/fd_server/src/main.rs b/android/fd_server/src/main.rs
index 26315bf..07f0896 100644
--- a/android/fd_server/src/main.rs
+++ b/android/fd_server/src/main.rs
@@ -29,9 +29,10 @@
 use log::debug;
 use nix::sys::stat::{umask, Mode};
 use rpcbinder::RpcServer;
+use rustutils::inherited_fd::take_fd_ownership;
 use std::collections::BTreeMap;
 use std::fs::File;
-use std::os::unix::io::{FromRawFd, OwnedFd};
+use std::os::unix::io::OwnedFd;
 
 use aidl::{FdConfig, FdService};
 use authfs_fsverity_metadata::parse_fsverity_metadata;
@@ -39,20 +40,6 @@
 // TODO(b/259920193): support dynamic port for multiple fd_server instances
 const RPC_SERVICE_PORT: u32 = 3264;
 
-fn is_fd_valid(fd: i32) -> bool {
-    // SAFETY: a query-only syscall
-    let retval = unsafe { libc::fcntl(fd, libc::F_GETFD) };
-    retval >= 0
-}
-
-fn fd_to_owned<T: FromRawFd>(fd: i32) -> Result<T> {
-    if !is_fd_valid(fd) {
-        bail!("Bad FD: {}", fd);
-    }
-    // SAFETY: The caller is supposed to provide valid FDs to this process.
-    Ok(unsafe { T::from_raw_fd(fd) })
-}
-
 fn parse_arg_ro_fds(arg: &str) -> Result<(i32, FdConfig)> {
     let result: Result<Vec<i32>, _> = arg.split(':').map(|x| x.parse::<i32>()).collect();
     let fds = result?;
@@ -62,13 +49,13 @@
     Ok((
         fds[0],
         FdConfig::Readonly {
-            file: fd_to_owned(fds[0])?,
+            file: take_fd_ownership(fds[0])?.into(),
             // Alternative metadata source, if provided
             alt_metadata: fds
                 .get(1)
-                .map(|fd| fd_to_owned(*fd))
+                .map(|fd| take_fd_ownership(*fd))
                 .transpose()?
-                .and_then(|f| parse_fsverity_metadata(f).ok()),
+                .and_then(|f| parse_fsverity_metadata(f.into()).ok()),
         },
     ))
 }
@@ -105,23 +92,26 @@
         fd_pool.insert(fd, config);
     }
     for fd in args.rw_fds {
-        let file = fd_to_owned::<File>(fd)?;
+        let file: File = take_fd_ownership(fd)?.into();
         if file.metadata()?.len() > 0 {
             bail!("File is expected to be empty");
         }
         fd_pool.insert(fd, FdConfig::ReadWrite(file));
     }
     for fd in args.ro_dirs {
-        fd_pool.insert(fd, FdConfig::InputDir(fd_to_owned(fd)?));
+        fd_pool.insert(fd, FdConfig::InputDir(take_fd_ownership(fd)?));
     }
     for fd in args.rw_dirs {
-        fd_pool.insert(fd, FdConfig::OutputDir(fd_to_owned(fd)?));
+        fd_pool.insert(fd, FdConfig::OutputDir(take_fd_ownership(fd)?));
     }
-    let ready_fd = args.ready_fd.map(fd_to_owned).transpose()?;
+    let ready_fd = args.ready_fd.map(take_fd_ownership).transpose()?;
     Ok((fd_pool, ready_fd))
 }
 
 fn main() -> Result<()> {
+    // SAFETY: nobody has taken ownership of the inherited FDs yet.
+    unsafe { rustutils::inherited_fd::init_once()? };
+
     android_logger::init_once(
         android_logger::Config::default()
             .with_tag("fd_server")
diff --git a/android/virtmgr/src/main.rs b/android/virtmgr/src/main.rs
index 67e7282..1625009 100644
--- a/android/virtmgr/src/main.rs
+++ b/android/virtmgr/src/main.rs
@@ -25,15 +25,15 @@
 
 use crate::aidl::{GLOBAL_SERVICE, VirtualizationService};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::BnVirtualizationService;
-use anyhow::{bail, Context, Result};
+use anyhow::{bail, Result};
 use binder::{BinderFeatures, ProcessState};
 use log::{info, LevelFilter};
 use rpcbinder::{FileDescriptorTransportMode, RpcServer};
-use std::os::unix::io::{AsFd, FromRawFd, OwnedFd, RawFd};
+use std::os::unix::io::{AsFd, RawFd};
 use std::sync::LazyLock;
 use clap::Parser;
-use nix::fcntl::{fcntl, F_GETFD, F_SETFD, FdFlag};
 use nix::unistd::{write, Pid, Uid};
+use rustutils::inherited_fd::take_fd_ownership;
 use std::os::unix::raw::{pid_t, uid_t};
 
 const LOG_TAG: &str = "virtmgr";
@@ -71,32 +71,6 @@
     ready_fd: RawFd,
 }
 
-fn take_fd_ownership(raw_fd: RawFd, owned_fds: &mut Vec<RawFd>) -> Result<OwnedFd, anyhow::Error> {
-    // Basic check that the integer value does correspond to a file descriptor.
-    fcntl(raw_fd, F_GETFD).with_context(|| format!("Invalid file descriptor {raw_fd}"))?;
-
-    // The file descriptor had CLOEXEC disabled to be inherited from the parent.
-    // Re-enable it to make sure it is not accidentally inherited further.
-    fcntl(raw_fd, F_SETFD(FdFlag::FD_CLOEXEC))
-        .with_context(|| format!("Could not set CLOEXEC on file descriptor {raw_fd}"))?;
-
-    // Creating OwnedFd for stdio FDs is not safe.
-    if [libc::STDIN_FILENO, libc::STDOUT_FILENO, libc::STDERR_FILENO].contains(&raw_fd) {
-        bail!("File descriptor {raw_fd} is standard I/O descriptor");
-    }
-
-    // Reject RawFds that already have a corresponding OwnedFd.
-    if owned_fds.contains(&raw_fd) {
-        bail!("File descriptor {raw_fd} already owned");
-    }
-    owned_fds.push(raw_fd);
-
-    // SAFETY: Initializing OwnedFd for a RawFd provided in cmdline arguments.
-    // We checked that the integer value corresponds to a valid FD and that this
-    // is the first argument to claim its ownership.
-    Ok(unsafe { OwnedFd::from_raw_fd(raw_fd) })
-}
-
 fn check_vm_support() -> Result<()> {
     if hypervisor_props::is_any_vm_supported()? {
         Ok(())
@@ -109,6 +83,10 @@
 }
 
 fn main() {
+    // SAFETY: nobody has taken ownership of the inherited FDs yet.
+    unsafe { rustutils::inherited_fd::init_once() }
+        .expect("Failed to take ownership of inherited FDs");
+
     android_logger::init_once(
         android_logger::Config::default()
             .with_tag(LOG_TAG)
@@ -120,11 +98,9 @@
 
     let args = Args::parse();
 
-    let mut owned_fds = vec![];
-    let rpc_server_fd = take_fd_ownership(args.rpc_server_fd, &mut owned_fds)
-        .expect("Failed to take ownership of rpc_server_fd");
-    let ready_fd = take_fd_ownership(args.ready_fd, &mut owned_fds)
-        .expect("Failed to take ownership of ready_fd");
+    let rpc_server_fd =
+        take_fd_ownership(args.rpc_server_fd).expect("Failed to take ownership of rpc_server_fd");
+    let ready_fd = take_fd_ownership(args.ready_fd).expect("Failed to take ownership of ready_fd");
 
     // Start thread pool for kernel Binder connection to VirtualizationServiceInternal.
     ProcessState::start_thread_pool();
diff --git a/guest/authfs_service/src/main.rs b/guest/authfs_service/src/main.rs
index 97e684d..be0f1b2 100644
--- a/guest/authfs_service/src/main.rs
+++ b/guest/authfs_service/src/main.rs
@@ -28,7 +28,6 @@
 use rustutils::sockets::android_get_control_socket;
 use std::ffi::OsString;
 use std::fs::{create_dir, read_dir, remove_dir_all, remove_file};
-use std::os::unix::io::{FromRawFd, OwnedFd};
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 use authfs_aidl_interface::aidl::com::android::virt::fs::AuthFsConfig::AuthFsConfig;
@@ -108,27 +107,11 @@
     Ok(())
 }
 
-/// Prepares a socket file descriptor for the authfs service.
-///
-/// # Safety requirement
-///
-/// The caller must ensure that this function is the only place that claims ownership
-/// of the file descriptor and it is called only once.
-unsafe fn prepare_authfs_service_socket() -> Result<OwnedFd> {
-    let raw_fd = android_get_control_socket(AUTHFS_SERVICE_SOCKET_NAME)?;
-
-    // Creating OwnedFd for stdio FDs is not safe.
-    if [libc::STDIN_FILENO, libc::STDOUT_FILENO, libc::STDERR_FILENO].contains(&raw_fd) {
-        bail!("File descriptor {raw_fd} is standard I/O descriptor");
-    }
-    // SAFETY: Initializing OwnedFd for a RawFd created by the init.
-    // We checked that the integer value corresponds to a valid FD and that the caller
-    // ensures that this is the only place to claim its ownership.
-    Ok(unsafe { OwnedFd::from_raw_fd(raw_fd) })
-}
-
 #[allow(clippy::eq_op)]
 fn try_main() -> Result<()> {
+    // SAFETY: nobody has taken ownership of the inherited FDs yet.
+    unsafe { rustutils::inherited_fd::init_once()? };
+
     let debuggable = env!("TARGET_BUILD_VARIANT") != "user";
     let log_level = if debuggable { log::LevelFilter::Trace } else { log::LevelFilter::Info };
     android_logger::init_once(
@@ -137,8 +120,7 @@
 
     clean_up_working_directory()?;
 
-    // SAFETY: This is the only place we take the ownership of the fd of the authfs service.
-    let socket_fd = unsafe { prepare_authfs_service_socket()? };
+    let socket_fd = android_get_control_socket(AUTHFS_SERVICE_SOCKET_NAME)?;
     let service = AuthFsService::new_binder(debuggable).as_binder();
     debug!("{} is starting as a rpc service.", AUTHFS_SERVICE_SOCKET_NAME);
     let server = RpcServer::new_bound_socket(service, socket_fd)?;
diff --git a/guest/microdroid_manager/src/main.rs b/guest/microdroid_manager/src/main.rs
index 8186e9d..fa089fa 100644
--- a/guest/microdroid_manager/src/main.rs
+++ b/guest/microdroid_manager/src/main.rs
@@ -56,7 +56,7 @@
 use std::ffi::CString;
 use std::fs::{self, create_dir, File, OpenOptions};
 use std::io::{Read, Write};
-use std::os::unix::io::{FromRawFd, OwnedFd};
+use std::os::unix::io::OwnedFd;
 use std::os::unix::process::CommandExt;
 use std::os::unix::process::ExitStatusExt;
 use std::path::Path;
@@ -170,6 +170,9 @@
 }
 
 fn main() -> Result<()> {
+    // SAFETY: nobody has taken ownership of the inherited FDs yet.
+    unsafe { rustutils::inherited_fd::init_once()? };
+
     // If debuggable, print full backtrace to console log with stdio_to_kmsg
     if is_debuggable()? {
         env::set_var("RUST_BACKTRACE", "full");
@@ -199,13 +202,7 @@
     );
     info!("started.");
 
-    // SAFETY: This is the only place we take the ownership of the fd of the vm payload service.
-    //
-    // To ensure that the CLOEXEC flag is set on the file descriptor as early as possible,
-    // it is necessary to fetch the socket corresponding to vm_payload_service at the
-    // very beginning, as android_get_control_socket() sets the CLOEXEC flag on the file
-    // descriptor.
-    let vm_payload_service_fd = unsafe { prepare_vm_payload_service_socket()? };
+    let vm_payload_service_fd = android_get_control_socket(VM_PAYLOAD_SERVICE_SOCKET_NAME)?;
 
     load_crashkernel_if_supported().context("Failed to load crashkernel")?;
 
@@ -486,25 +483,6 @@
         .context("Could not connect to IVirtualMachineService")
 }
 
-/// Prepares a socket file descriptor for the vm payload service.
-///
-/// # Safety
-///
-/// The caller must ensure that this function is the only place that claims ownership
-/// of the file descriptor and it is called only once.
-unsafe fn prepare_vm_payload_service_socket() -> Result<OwnedFd> {
-    let raw_fd = android_get_control_socket(VM_PAYLOAD_SERVICE_SOCKET_NAME)?;
-
-    // Creating OwnedFd for stdio FDs is not safe.
-    if [libc::STDIN_FILENO, libc::STDOUT_FILENO, libc::STDERR_FILENO].contains(&raw_fd) {
-        bail!("File descriptor {raw_fd} is standard I/O descriptor");
-    }
-    // SAFETY: Initializing OwnedFd for a RawFd created by the init.
-    // We checked that the integer value corresponds to a valid FD and that the caller
-    // ensures that this is the only place to claim its ownership.
-    Ok(unsafe { OwnedFd::from_raw_fd(raw_fd) })
-}
-
 fn is_strict_boot() -> bool {
     Path::new(AVF_STRICT_BOOT).exists()
 }