GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Windows impl

Browse Source
This commit is contained in:
Andrej Mihajlov
2025-05-29 18:14:55 +02:00
parent 4eedbb235a
commit c225511f95
6 changed files with 403 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sqlx-pool-guard/Cargo.toml
+17 -2
View File
@@ -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 }
tempfile = { workspace = true }
sqlx-pool-guard/src/lib.rs
+34 -107
View File
@@ -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<bool> {
let fd_list = proc_pidinfo_list_self::<ProcFDInfo>()?;
for fd in fd_list
.iter()
.filter(|s| s.fd_type() == Ok(ProcFDType::VNODE))
{
let Some(vnode) = proc_pidfdinfo_self::<VnodeFdInfoWithPath>(fd.proc_fd)
.inspect_err(|e| {
log::warn!("proc_pidfdinfo_self::<VnodeFdInfoWithPath>() 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<bool> {
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<bool> {
// 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)
sqlx-pool-guard/src/linux.rs
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// 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<bool> {
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)
}
sqlx-pool-guard/src/macos.rs
+41
View File
@@ -0,0 +1,41 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// 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<bool> {
let fd_list = proc_pidinfo_list_self::<ProcFDInfo>()?;
for fd in fd_list
.iter()
.filter(|s| s.fd_type() == Ok(ProcFDType::VNODE))
{
let Some(vnode) = proc_pidfdinfo_self::<VnodeFdInfoWithPath>(fd.proc_fd)
.inspect_err(|e| {
log::warn!("proc_pidfdinfo_self::<VnodeFdInfoWithPath>() 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)
}
sqlx-pool-guard/src/windows.rs
+174
View File
@@ -0,0 +1,174 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// 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<bool> {
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::<SystemHandleInformation>::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<T> {
pub inner: *mut T,
process_heap: HANDLE,
}
impl<T> HeapGuard<T> {
fn new(length: usize) -> io::Result<Self> {
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<T> Drop for HeapGuard<T> {
fn drop(&mut self) {
unsafe {
HeapFree(
self.process_heap,
HEAP_FLAGS(0),
Some(self.inner as *mut c_void),
)
}
.expect("HeapFree failure");
}
}