From c225511f951bf3f6480e16efa3d03fb50737d58b Mon Sep 17 00:00:00 2001 From: Andrej Mihajlov Date: Thu, 29 May 2025 18:14:55 +0200 Subject: [PATCH] Add Windows impl --- Cargo.lock | 109 +++++++++++++++++++-- sqlx-pool-guard/Cargo.toml | 19 +++- sqlx-pool-guard/src/lib.rs | 141 +++++++------------------- sqlx-pool-guard/src/linux.rs | 34 +++++++ sqlx-pool-guard/src/macos.rs | 41 ++++++++ sqlx-pool-guard/src/windows.rs | 174 +++++++++++++++++++++++++++++++++ 6 files changed, 403 insertions(+), 115 deletions(-) create mode 100644 sqlx-pool-guard/src/linux.rs create mode 100644 sqlx-pool-guard/src/macos.rs create mode 100644 sqlx-pool-guard/src/windows.rs diff --git a/Cargo.lock b/Cargo.lock index 3a13a133f4..01820eb102 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9543,6 +9543,7 @@ dependencies = [ "sqlx", "tempfile", "tokio", + "windows 0.61.1", ] [[package]] @@ -11559,6 +11560,28 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5ee8f3d025738cb02bad7868bbb5f8a6327501e870bf51f1b455b0a2454a419" +dependencies = [ + "windows-collections", + "windows-core 0.61.2", + "windows-future", + "windows-link", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + [[package]] name = "windows-core" version = "0.52.0" @@ -11593,6 +11616,30 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement 0.60.0", + "windows-interface 0.59.1", + "windows-link", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link", + "windows-threading", +] + [[package]] name = "windows-implement" version = "0.57.0" @@ -11615,6 +11662,17 @@ dependencies = [ "syn 2.0.98", ] +[[package]] +name = "windows-implement" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a47fddd13af08290e67f4acabf4b459f647552718f683a7b415d290ac744a836" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.98", +] + [[package]] name = "windows-interface" version = "0.57.0" @@ -11638,10 +11696,31 @@ dependencies = [ ] [[package]] -name = "windows-link" -version = "0.1.0" +name = "windows-interface" +version = "0.59.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dccfd733ce2b1753b03b6d3c65edf020262ea35e20ccdf3e288043e6dd620e3" +checksum = "bd9211b69f8dcdfa817bfd14bf1c97c9188afa36f4750130fcdf3f400eca9fa8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.98", +] + +[[package]] +name = "windows-link" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76840935b766e1b0a05c0066835fb9ec80071d4c09a16f6bd5f7e655e3c14c38" + +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link", +] [[package]] name = "windows-registry" @@ -11649,7 +11728,7 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4286ad90ddb45071efd1a66dfa43eb02dd0dfbae1545ad6cc3c51cf34d7e8ba3" dependencies = [ - "windows-result 0.3.1", + "windows-result 0.3.4", "windows-strings 0.3.1", "windows-targets 0.53.0", ] @@ -11674,9 +11753,9 @@ dependencies = [ [[package]] name = "windows-result" -version = "0.3.1" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06374efe858fab7e4f881500e6e86ec8bc28f9462c47e5a9941a0142ad86b189" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" dependencies = [ "windows-link", ] @@ -11700,6 +11779,15 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.45.0" @@ -11798,6 +11886,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.0", ] +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.42.2" diff --git a/sqlx-pool-guard/Cargo.toml b/sqlx-pool-guard/Cargo.toml index e0b4452d3c..a77ee80024 100644 --- a/sqlx-pool-guard/Cargo.toml +++ b/sqlx-pool-guard/Cargo.toml @@ -8,12 +8,27 @@ license.workspace = true workspace = true [dependencies] -tokio = { workspace = true, features = ["rt-multi-thread", "macros", "time", "fs"] } +tokio = { workspace = true, features = [ + "rt-multi-thread", + "macros", + "time", + "fs", +] } sqlx = { workspace = true, features = ["runtime-tokio-rustls", "sqlite"] } log.workspace = true [target.'cfg(target_os = "macos")'.dependencies] proc_pidinfo.workspace = true +[target.'cfg(windows)'.dependencies] +windows = { version = "0.61", features = [ + "Win32", + "Win32_System", + "Win32_System_Memory", + "Win32_System_Threading", + "Win32_Storage_FileSystem", + "Wdk_System_SystemInformation", +] } + [dev-dependencies] -tempfile = { workspace = true } \ No newline at end of file +tempfile = { workspace = true } diff --git a/sqlx-pool-guard/src/lib.rs b/sqlx-pool-guard/src/lib.rs index b3e0f672bc..8b120ee261 100644 --- a/sqlx-pool-guard/src/lib.rs +++ b/sqlx-pool-guard/src/lib.rs @@ -8,13 +8,23 @@ use std::{ time::Duration, }; -#[cfg(target_os = "macos")] -use proc_pidinfo::{ - ProcFDInfo, ProcFDType, VnodeFdInfoWithPath, proc_pidfdinfo_self, proc_pidinfo_list_self, -}; +#[cfg(windows)] +#[path = "windows.rs"] +mod imp; -const SQL_CLOSE_MAX_ATTEMPTS: u8 = 10; -const SQL_CLOSE_RETRY_DELAY: Duration = Duration::from_millis(100); +#[cfg(target_os = "macos")] +#[path = "macos.rs"] +mod imp; + +#[cfg(any(target_os = "linux", target_os = "android"))] +#[path = "linux.rs"] +mod imp; + +/// Max number of retry attempts +const CHECK_FILES_CLOSED_MAX_ATTEMPTS: u8 = 10; + +/// Delay between file checks +const CHECK_FILES_CLOSED_RETRY_DELAY: Duration = Duration::from_millis(100); pub struct SqlitePoolGuard { database_path: PathBuf, @@ -42,7 +52,7 @@ impl SqlitePoolGuard { } } - /// Close udnerlying sqlite pool + /// Close udnerlying sqlite pool and wait for files to be closed before returning. pub async fn close_pool(&self) { _ = self.close_pool_inner(); } @@ -50,7 +60,7 @@ impl SqlitePoolGuard { async fn close_pool_inner(&self) -> std::io::Result<()> { self.connection_pool.close().await; - if let Err(e) = self.wait_io_close().await { + if let Err(e) = self.wait_for_db_files_close().await { log::error!("Failed to wait for file to close: {e}"); } @@ -82,25 +92,17 @@ impl SqlitePoolGuard { database_files } - /// Wait for I/O close to the database files - /// - /// - macOS: uses `proc_pidinfo` (`sys/proc_info.h`) - /// See: http://blog.palominolabs.com/2012/06/19/getting-the-files-being-used-by-a-process-on-mac-os-x/ - /// - /// - Linux, Android: uses `/proc/self/fd/` to list open file descriptors - /// See: https://stackoverflow.com/a/59797198/351305 - /// - /// - Windows: attempts to open files to detect whether they are still open. - async fn wait_io_close(&self) -> std::io::Result<()> { + /// Wait for database files to be closed before returning. + async fn wait_for_db_files_close(&self) -> std::io::Result<()> { let database_files = self.all_database_files(); let paths: Vec<&Path> = database_files.iter().map(PathBuf::as_path).collect(); - for _ in 0..SQL_CLOSE_MAX_ATTEMPTS { - match Self::check_io_close(&paths) + for _ in 0..CHECK_FILES_CLOSED_MAX_ATTEMPTS { + match imp::check_files_closed(&paths) .await - .inspect_err(|e| log::error!("check_io_close() failure: {e}")) + .inspect_err(|e| log::error!("imp::check_files_closed() failure: {e}")) { - Ok(false) | Err(_) => tokio::time::sleep(SQL_CLOSE_RETRY_DELAY).await, + Ok(false) | Err(_) => tokio::time::sleep(CHECK_FILES_CLOSED_RETRY_DELAY).await, Ok(true) => return Ok(()), } } @@ -110,95 +112,12 @@ impl SqlitePoolGuard { "timed out waiting for sqlite files to be closed", )) } - - /// Check if no more open file descriptors exist for the given files. - #[cfg(target_os = "macos")] - async fn check_io_close(file_paths: &[&Path]) -> io::Result { - let fd_list = proc_pidinfo_list_self::()?; - - for fd in fd_list - .iter() - .filter(|s| s.fd_type() == Ok(ProcFDType::VNODE)) - { - let Some(vnode) = proc_pidfdinfo_self::(fd.proc_fd) - .inspect_err(|e| { - log::warn!("proc_pidfdinfo_self::() failure: {e}"); - }) - .ok() - .flatten() - else { - continue; - }; - - if let Ok(true) = vnode - .path() - .map(|vnode_path| file_paths.contains(&vnode_path)) - .inspect_err(|e| { - log::warn!("vnode.path() failure: {e:?}"); - }) - { - return Ok(false); - } - } - - Ok(true) - } - - /// Check if no more open file descriptors exist for the given files. - #[cfg(any(target_os = "linux", target_os = "android"))] - async fn check_io_close(file_paths: &[&Path]) -> io::Result { - let mut dir = tokio::fs::read_dir("/proc/self/fd/").await?; - - while let Ok(Some(entry)) = dir.next_entry().await { - if entry - .file_type() - .await - .inspect_err(|e| log::warn!("entry.file_type() failure: {e}")) - .is_ok_and(|entry_type| entry_type.is_symlink()) - { - match tokio::fs::read_link(entry.path()).await { - Ok(resolved_path) => { - if file_paths.contains(&resolved_path.as_ref()) { - return Ok(false); - } - } - Err(e) => { - log::error!("Failed to read symlink: {e}"); - } - } - } - } - - Ok(true) - } - - #[cfg(windows)] - async fn check_io_close(file_paths: &[&Path]) -> io::Result { - // Error code returned when file is still in use. - const FILE_IN_USE_ERR: i32 = 32; - - for file_path in file_paths { - if let Err(e) = tokio::fs::OpenOptions::new() - .read(true) - .open(file_path) - .await - { - if e.raw_os_error() == Some(FILE_IN_USE_ERR) { - return Ok(false); - } else if e.kind() != io::ErrorKind::NotFound { - log::error!("Failed to open file: {}", file_path.display()); - } - } - } - - Ok(true) - } } #[cfg(test)] mod tests { use sqlx::{ - ConnectOptions, + ConnectOptions, Executor, sqlite::{SqliteAutoVacuum, SqliteSynchronous}, }; @@ -208,6 +127,9 @@ mod tests { async fn test_wait_close() { let temp_dir = tempfile::tempdir().unwrap(); let database_path = temp_dir.path().join("storage.sqlite"); + + println!("Database path: {}", database_path.display()); + let opts = sqlx::sqlite::SqliteConnectOptions::new() .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) .synchronous(SqliteSynchronous::Normal) @@ -217,10 +139,15 @@ mod tests { .disable_statement_logging(); let connection_pool = sqlx::SqlitePool::connect_with(opts).await.unwrap(); + connection_pool + .execute("create table test (col int)") + .await + .unwrap(); + let guard = SqlitePoolGuard::new(database_path, connection_pool); assert!( guard - .wait_io_close() + .wait_for_db_files_close() .await .err() .is_some_and(|e| e.kind() == io::ErrorKind::TimedOut) diff --git a/sqlx-pool-guard/src/linux.rs b/sqlx-pool-guard/src/linux.rs new file mode 100644 index 0000000000..5a1fc94bd4 --- /dev/null +++ b/sqlx-pool-guard/src/linux.rs @@ -0,0 +1,34 @@ +// Copyright 2025 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +static PROC_SELF_FD_DIR: &str = "/proc/self/fd/"; + +/// Check if there are no open file descriptors for the given files. +/// +/// Linux, Android: uses `/proc/self/fd/` to list open file descriptors +/// See: https://stackoverflow.com/a/59797198/351305 +pub async fn check_files_closed(file_paths: &[&Path]) -> io::Result { + let mut dir = tokio::fs::read_dir(PROC_SELF_FD_DIR).await?; + + while let Ok(Some(entry)) = dir.next_entry().await { + if entry + .file_type() + .await + .inspect_err(|e| log::warn!("entry.file_type() failure: {e}")) + .is_ok_and(|entry_type| entry_type.is_symlink()) + { + match tokio::fs::read_link(entry.path()).await { + Ok(resolved_path) => { + if file_paths.contains(&resolved_path.as_ref()) { + return Ok(false); + } + } + Err(e) => { + log::error!("Failed to read symlink: {e}"); + } + } + } + } + + Ok(true) +} diff --git a/sqlx-pool-guard/src/macos.rs b/sqlx-pool-guard/src/macos.rs new file mode 100644 index 0000000000..1142ffd11b --- /dev/null +++ b/sqlx-pool-guard/src/macos.rs @@ -0,0 +1,41 @@ +// Copyright 2025 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +use proc_pidinfo::{ + ProcFDInfo, ProcFDType, VnodeFdInfoWithPath, proc_pidfdinfo_self, proc_pidinfo_list_self, +}; + +/// Check if there are no open file descriptors for the given files. +/// +/// Uses `proc_pidinfo` (`sys/proc_info.h`) +/// See: http://blog.palominolabs.com/2012/06/19/getting-the-files-being-used-by-a-process-on-mac-os-x/ +pub async fn check_files_closed(file_paths: &[&Path]) -> io::Result { + let fd_list = proc_pidinfo_list_self::()?; + + for fd in fd_list + .iter() + .filter(|s| s.fd_type() == Ok(ProcFDType::VNODE)) + { + let Some(vnode) = proc_pidfdinfo_self::(fd.proc_fd) + .inspect_err(|e| { + log::warn!("proc_pidfdinfo_self::() failure: {e}"); + }) + .ok() + .flatten() + else { + continue; + }; + + if let Ok(true) = vnode + .path() + .map(|vnode_path| file_paths.contains(&vnode_path)) + .inspect_err(|e| { + log::warn!("vnode.path() failure: {e:?}"); + }) + { + return Ok(false); + } + } + + Ok(true) +} diff --git a/sqlx-pool-guard/src/windows.rs b/sqlx-pool-guard/src/windows.rs new file mode 100644 index 0000000000..43eee8f465 --- /dev/null +++ b/sqlx-pool-guard/src/windows.rs @@ -0,0 +1,174 @@ +// Copyright 2025 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + ffi::{OsString, c_uchar, c_ulong, c_ushort, c_void}, + io, + os::windows::ffi::OsStringExt, + path::{Path, PathBuf}, +}; + +use windows::{ + Wdk::System::SystemInformation::{NtQuerySystemInformation, SYSTEM_INFORMATION_CLASS}, + Win32::{ + Foundation::{HANDLE, MAX_PATH, NTSTATUS, STATUS_INFO_LENGTH_MISMATCH}, + Storage::FileSystem::{ + FILE_NAME_NORMALIZED, FILE_TYPE_DISK, GetFileType, GetFinalPathNameByHandleW, + }, + System::{ + Memory::{ + GetProcessHeap, HEAP_FLAGS, HEAP_ZERO_MEMORY, HeapAlloc, HeapFree, HeapReAlloc, + }, + Threading::GetCurrentProcessId, + }, + }, +}; + +/// Private information class used to retrieve open file handles +const SYSTEM_HANDLE_INFORMATION_CLASS: SYSTEM_INFORMATION_CLASS = SYSTEM_INFORMATION_CLASS(0x10); + +/// Initial buffer size holding the handle info +/// The number is based on what I observe on a pretty standard Windows 11 +const SYSTEM_HANDLE_INFORMATION_INITIAL_SIZE: usize = 2_500_000; + +/// Check if there are no open handles to the given files. +/// +/// Uses undocumented NT API to obtain open handles on the system. +/// See: https://www.ired.team/miscellaneous-reversing-forensics/windows-kernel-internals/get-all-open-handles-and-kernel-object-address-from-userland +pub async fn check_files_closed(file_paths: &[&Path]) -> io::Result { + let current_pid = unsafe { GetCurrentProcessId() }; + + // Allocate info struct on heap with some initial value + let mut reserved_memory = SYSTEM_HANDLE_INFORMATION_INITIAL_SIZE; + let mut handle_table_info = HeapGuard::::new(reserved_memory)?; + + let mut status: NTSTATUS = NTSTATUS::default(); + let mut return_len = reserved_memory as u32; + for _ in 0..2 { + status = unsafe { + NtQuerySystemInformation( + SYSTEM_HANDLE_INFORMATION_CLASS, + handle_table_info.inner as _, + return_len, + &mut return_len, + ) + }; + + // Buffer is too small, resize memory and retry again. + if status == STATUS_INFO_LENGTH_MISMATCH { + log::trace!("Buffer is too small ({reserved_memory}), resizing to {return_len}"); + println!("Buffer is too small ({reserved_memory}), resizing to {return_len}"); + reserved_memory = return_len as usize; + handle_table_info.reallocate(reserved_memory)?; + } else { + break; + } + } + status.ok()?; + + let num_handles = unsafe { (*handle_table_info.inner).number_of_handles }; + let proc_entries = unsafe { + std::slice::from_raw_parts( + (*handle_table_info.inner).handles.as_ptr(), + num_handles as usize, + ) + }; + + for entry in proc_entries { + if entry.unique_process_id == current_pid { + let file_handle = HANDLE(entry.handle_value as _); + + if unsafe { GetFileType(file_handle) } == FILE_TYPE_DISK { + let mut file_handle_path = vec![0u16; MAX_PATH as usize]; + let num_chars_without_nul = unsafe { + GetFinalPathNameByHandleW( + file_handle, + &mut file_handle_path, + FILE_NAME_NORMALIZED, + ) as usize + }; + + if num_chars_without_nul > 0 { + let path_str = OsString::from_wide(&file_handle_path[0..num_chars_without_nul]); + let file_handle_pathbuf = PathBuf::from(path_str); + + if file_paths.contains(&file_handle_pathbuf.as_path()) { + return Ok(false); + } + } + } + } + } + + Ok(true) +} + +#[repr(C)] +#[derive(Copy, Clone)] +struct SystemHandleInformation { + pub number_of_handles: c_ulong, + pub handles: [SystemHandleTableEntryInfo; 1], +} + +#[repr(C)] +#[derive(Copy, Clone)] +struct SystemHandleTableEntryInfo { + pub unique_process_id: c_ulong, + pub object_type_index: c_uchar, + pub handle_attributes: c_uchar, + pub handle_value: c_ushort, + pub object: *mut c_void, + pub granted_access: c_ulong, +} + +struct HeapGuard { + pub inner: *mut T, + process_heap: HANDLE, +} + +impl HeapGuard { + fn new(length: usize) -> io::Result { + let process_heap = unsafe { GetProcessHeap()? }; + let inner: *mut T = unsafe { HeapAlloc(process_heap, HEAP_ZERO_MEMORY, length) as _ }; + + if inner.is_null() { + Err(io::Error::other("Failed to allocate memory")) + } else { + Ok(Self { + inner, + process_heap, + }) + } + } + + fn reallocate(&mut self, new_length: usize) -> io::Result<()> { + let new_ptr: *mut T = unsafe { + HeapReAlloc( + self.process_heap, + HEAP_ZERO_MEMORY, + Some(self.inner as _), + new_length, + ) as _ + }; + + if new_ptr.is_null() { + Err(io::Error::other("Failed to reallocate memory")) + } else { + self.inner = new_ptr; + Ok(()) + } + } +} + +impl Drop for HeapGuard { + fn drop(&mut self) { + unsafe { + HeapFree( + self.process_heap, + HEAP_FLAGS(0), + Some(self.inner as *mut c_void), + ) + } + .expect("HeapFree failure"); + } +}