Kronecker/NotesAgentic
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 4 Wiki Activity

feat: 添加加密凭据恢复与 Windows 会话撤销

Browse Source
This commit is contained in:
admin
2026-09-08 14:46:30 +08:00
parent 0d225f308f
commit 3439d1fe3f
14 changed files with 564 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
frontend/src-tauri/Cargo.lock
Generated
+1
View File
@@ -2943,6 +2943,7 @@ dependencies = [
"tempfile",
"tokio",
"uuid",
"windows-sys 0.61.2",
"zeroize",
]
frontend/src-tauri/Cargo.toml
+3
View File
@@ -38,6 +38,9 @@ argon2 = "0.5"
chacha20poly1305 = "0.10"
fernet = { version = "0.2", default-features = false, features = ["rustcrypto"] }
[target.'cfg(windows)'.dependencies]
windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_System_RemoteDesktop", "Win32_UI_WindowsAndMessaging", "Win32_Graphics_Gdi", "Win32_System_LibraryLoader"] }
[build-dependencies]
tauri-build = { version = "2", optional = true , features = [] }
frontend/src-tauri/build.rs
+2
View File
@@ -21,6 +21,8 @@ fn main() {
"credentials_lock",
"credentials_change_password",
"credentials_import",
"credentials_backup",
"credentials_restore",
"core_request",
"core_request_prepare",
"core_request_cancel",
frontend/src-tauri/capabilities/main.json
+2
View File
@@ -13,6 +13,8 @@
"allow-credentials-lock",
"allow-credentials-change-password",
"allow-credentials-import",
"allow-credentials-backup",
"allow-credentials-restore",
"allow-core-request",
"allow-core-request-prepare",
"allow-core-request-cancel",
frontend/src-tauri/permissions/autogenerated/credentials_backup.toml
+11
View File
@@ -0,0 +1,11 @@
# Automatically generated - DO NOT EDIT!
[[permission]]
identifier = "allow-credentials-backup"
description = "Enables the credentials_backup command without any pre-configured scope."
commands.allow = ["credentials_backup"]
[[permission]]
identifier = "deny-credentials-backup"
description = "Denies the credentials_backup command without any pre-configured scope."
commands.deny = ["credentials_backup"]
frontend/src-tauri/permissions/autogenerated/credentials_restore.toml
+11
View File
@@ -0,0 +1,11 @@
# Automatically generated - DO NOT EDIT!
[[permission]]
identifier = "allow-credentials-restore"
description = "Enables the credentials_restore command without any pre-configured scope."
commands.allow = ["credentials_restore"]
[[permission]]
identifier = "deny-credentials-restore"
description = "Denies the credentials_restore command without any pre-configured scope."
commands.deny = ["credentials_restore"]
frontend/src-tauri/src/credentials.rs
+207 -54
View File
@@ -15,6 +15,10 @@ use std::collections::BTreeMap;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use zeroize::Zeroizing;
type Result<T> = std::result::Result<T, String>;
@@ -204,6 +208,8 @@ pub struct CredentialBroker {
// Separate stable inode: snapshots are atomically replaced, so locking the
// snapshot itself would not protect the next writer after replacement.
ownership: Option<fs::File>,
lock_epoch: Arc<AtomicU64>,
unlocked_epoch: u64,
}
impl CredentialBroker {
@@ -268,7 +274,7 @@ impl CredentialBroker {
.ok_or("MIGRATION_KEY_MISSING")?;
let fernet =
Zeroizing::new(fernet::Fernet::new(key.trim()).ok_or("MIGRATION_KEY_INVALID")?);
let session = self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED")?;
let session = self.session()?;
let mut decoded = Vec::new();
for (id, token) in &tokens {
let key = CredentialId::legacy(id).key()?;
@@ -381,7 +387,7 @@ impl CredentialBroker {
let method = request["rpc"].as_str().ok_or("HOST_REQUEST_INVALID")?;
let params = &request["params"];
if method == "credentials.delete_many" || method == "credentials.move_many" {
let session = self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED")?;
let session = self.session()?;
let result = (|| {
let mut removed = Vec::new();
if method.ends_with("delete_many") {
@@ -483,10 +489,21 @@ impl CredentialBroker {
path,
unlocked: None,
ownership: None,
lock_epoch: Arc::new(AtomicU64::new(0)),
unlocked_epoch: 0,
}
}
pub fn is_locked(&self) -> bool {
self.unlocked.is_none()
self.unlocked.is_none() || self.lock_epoch.load(Ordering::SeqCst) != self.unlocked_epoch
}
pub fn lock_signal(&self) -> Arc<AtomicU64> {
self.lock_epoch.clone()
}
fn session(&self) -> Result<&Unlocked> {
if self.is_locked() {
return Err("CREDENTIALS_LOCKED".into());
}
self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED".into())
}
pub fn lock(&mut self) {
self.unlocked.take();
@@ -494,11 +511,36 @@ impl CredentialBroker {
}
pub fn unlock(&mut self, password: Zeroizing<Vec<u8>>) -> Result<()> {
self.lock();
let epoch = self.lock_epoch.load(Ordering::SeqCst);
let ownership = Self::acquire_ownership(&self.path)?;
let session = if self.path.exists() {
Self::load_snapshot(&self.path, &password)?
} else {
let mut salt = [0u8; 32];
rand::rngs::OsRng
.try_fill_bytes(&mut salt)
.map_err(|_| "CREDENTIAL_ENTROPY_FAILED")?;
let session = Unlocked::derive(&password, salt)?;
session
.stronghold
.create_client(CLIENT)
.map_err(|_| "CREDENTIAL_STORE_FAILED")?;
session.persist(&self.path)?;
session
};
if self.lock_epoch.load(Ordering::SeqCst) != epoch {
return Err("CREDENTIALS_LOCKED".into());
}
self.unlocked_epoch = epoch;
self.unlocked = Some(session);
self.ownership = Some(ownership);
Ok(())
}
fn acquire_ownership(path: &Path) -> Result<fs::File> {
use fs2::FileExt;
let parent = self.path.parent().ok_or("CREDENTIAL_PATH_INVALID")?;
let parent = path.parent().ok_or("CREDENTIAL_PATH_INVALID")?;
fs::create_dir_all(parent).map_err(|_| "CREDENTIAL_IO_FAILED")?;
let mut lock_name = self
.path
let mut lock_name = path
.file_name()
.ok_or("CREDENTIAL_PATH_INVALID")?
.to_os_string();
@@ -534,54 +576,96 @@ impl CredentialBroker {
ownership
.try_lock_exclusive()
.map_err(|_| "CREDENTIALS_BUSY")?;
let session = if self.path.exists() {
Ok(ownership)
}
fn load_snapshot(path: &Path, password: &[u8]) -> Result<Unlocked> {
let metadata = fs::symlink_metadata(path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
if !metadata.is_file() || metadata.len() > MAX_FILE {
return Err("CREDENTIAL_STORE_CORRUPT".into());
}
let data = fs::read(path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
if data.len() < 40 || &data[..8] != MAGIC {
return Err("SCHEMA_INCOMPATIBLE".into());
}
let mut salt = [0u8; 32];
salt.copy_from_slice(&data[8..40]);
let session = Unlocked::derive(password, salt)?;
// Backups can be on read-only media. This temporary file contains ciphertext only.
let mut temp = tempfile::NamedTempFile::new().map_err(|_| "CREDENTIAL_IO_FAILED")?;
temp.write_all(&data[40..])
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
session
.stronghold
.load_client_from_snapshot(
CLIENT,
&session.provider()?,
&SnapshotPath::from_path(temp.path()),
)
.map_err(|_| "CREDENTIAL_UNLOCK_FAILED")?;
Ok(session)
}
/// Native picker selected destination; backup is encrypted and never overwrites.
pub fn backup(&self, destination: &Path) -> Result<()> {
self.session()?;
let parent = destination.parent().ok_or("CREDENTIAL_PATH_INVALID")?;
let mut target =
tempfile::NamedTempFile::new_in(parent).map_err(|_| "CREDENTIAL_IO_FAILED")?;
let bytes = fs::read(&self.path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
target
.write_all(&bytes)
.and_then(|_| target.as_file().sync_all())
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
target
.persist_noclobber(destination)
.map_err(|_| "CREDENTIAL_BACKUP_EXISTS")?;
Ok(())
}
/// Validate every record before atomic replacement; preserve the previous encrypted file.
/// Caller must obtain explicit confirmation through the native dialog.
pub fn restore(&mut self, source: &Path, password: Zeroizing<Vec<u8>>) -> Result<usize> {
if !self.is_locked() {
return Err("CREDENTIALS_MUST_LOCK".into());
}
self.lock();
let _ownership = Self::acquire_ownership(&self.path)?;
let session = Self::load_snapshot(source, &password)?;
let keys = session
.store()?
.keys()
.map_err(|_| "CREDENTIAL_STORE_FAILED")?;
for key in &keys {
let id: CredentialId =
serde_json::from_slice(key).map_err(|_| "CREDENTIAL_STORE_CORRUPT")?;
if id.key()? != *key || session.read(key)?.is_none() {
return Err("CREDENTIAL_STORE_CORRUPT".into());
}
}
let parent = self.path.parent().ok_or("CREDENTIAL_PATH_INVALID")?;
if self.path.exists() {
let metadata = fs::symlink_metadata(&self.path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
if !metadata.is_file() || metadata.len() > MAX_FILE {
return Err("CREDENTIAL_STORE_CORRUPT".into());
}
let data = fs::read(&self.path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
if data.len() < 40 || &data[..8] != MAGIC {
return Err("SCHEMA_INCOMPATIBLE".into());
}
let mut salt = [0u8; 32];
salt.copy_from_slice(&data[8..40]);
let session = Unlocked::derive(&password, salt)?;
let mut temp = tempfile::NamedTempFile::new_in(
self.path.parent().ok_or("CREDENTIAL_PATH_INVALID")?,
)
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
temp.write_all(&data[40..])
let mut previous = tempfile::Builder::new()
.prefix("pre-restore-")
.suffix(".onxcred")
.tempfile_in(parent)
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
session
.stronghold
.load_client_from_snapshot(
CLIENT,
&session.provider()?,
&SnapshotPath::from_path(temp.path()),
)
.map_err(|_| "CREDENTIAL_UNLOCK_FAILED")?;
session
} else {
let mut salt = [0u8; 32];
rand::rngs::OsRng
.try_fill_bytes(&mut salt)
.map_err(|_| "CREDENTIAL_ENTROPY_FAILED")?;
let session = Unlocked::derive(&password, salt)?;
session
.stronghold
.create_client(CLIENT)
.map_err(|_| "CREDENTIAL_STORE_FAILED")?;
session.persist(&self.path)?;
session
};
self.unlocked = Some(session);
self.ownership = Some(ownership);
Ok(())
previous
.write_all(&fs::read(&self.path).map_err(|_| "CREDENTIAL_IO_FAILED")?)
.and_then(|_| previous.as_file().sync_all())
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
previous.keep().map_err(|_| "CREDENTIAL_IO_FAILED")?;
}
session.persist(&self.path)?;
// Restoration deliberately leaves the vault locked; no implicit permission grant.
Ok(keys.len())
}
pub fn list(&self) -> Result<Vec<CredentialId>> {
self.unlocked
.as_ref()
.ok_or("CREDENTIALS_LOCKED")?
self.session()?
.store()?
.keys()
.map_err(|_| "CREDENTIAL_STORE_FAILED")?
@@ -590,7 +674,7 @@ impl CredentialBroker {
.collect()
}
pub fn put(&mut self, id: &CredentialId, value: Zeroizing<Vec<u8>>) -> Result<()> {
let session = self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED")?;
let session = self.session()?;
let result = session
.write(id.key()?, &value)
.and_then(|_| session.persist(&self.path));
@@ -600,7 +684,7 @@ impl CredentialBroker {
result
}
pub fn delete(&mut self, id: &CredentialId) -> Result<()> {
let session = self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED")?;
let session = self.session()?;
session
.store()?
.delete(&id.key()?)
@@ -617,13 +701,10 @@ impl CredentialBroker {
if caller != &id.scope {
return Err("CREDENTIAL_SCOPE_DENIED".into());
}
self.unlocked
.as_ref()
.ok_or("CREDENTIALS_LOCKED")?
.read(&id.key()?)
self.session()?.read(&id.key()?)
}
pub fn change_password(&mut self, password: Zeroizing<Vec<u8>>) -> Result<()> {
let previous = self.unlocked.as_ref().ok_or("CREDENTIALS_LOCKED")?;
let previous = self.session()?;
let mut salt = [0u8; 32];
rand::rngs::OsRng
.try_fill_bytes(&mut salt)
@@ -653,6 +734,78 @@ mod tests {
Zeroizing::new(b"test-only-password-123".to_vec())
}
#[test]
fn encrypted_backup_restores_corrupt_store_without_overwrite_on_failure() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("credentials.v1");
let backup = temp.path().join("backup.onxcred");
let mut broker = CredentialBroker::new(path.clone());
broker.unlock(password()).unwrap();
let id = CredentialId::legacy("test-provider");
broker
.put(&id, Zeroizing::new(b"backup-test-secret".to_vec()))
.unwrap();
broker.backup(&backup).unwrap();
assert_eq!(
broker.backup(&backup).unwrap_err(),
"CREDENTIAL_BACKUP_EXISTS"
);
assert_eq!(
broker.restore(&backup, password()).unwrap_err(),
"CREDENTIALS_MUST_LOCK"
);
broker.lock();
fs::write(&path, b"corrupt-original").unwrap();
assert!(broker
.restore(&backup, Zeroizing::new(b"wrong-test-password".to_vec()))
.is_err());
assert_eq!(fs::read(&path).unwrap(), b"corrupt-original");
assert_eq!(broker.restore(&backup, password()).unwrap(), 1);
assert!(broker.is_locked());
let saved = fs::read_dir(temp.path())
.unwrap()
.filter_map(|e| e.ok())
.find(|e| e.file_name().to_string_lossy().starts_with("pre-restore-"))
.unwrap();
assert_eq!(fs::read(saved.path()).unwrap(), b"corrupt-original");
broker.unlock(password()).unwrap();
assert_eq!(
broker
.resolve(&Scope::Provider, &id)
.unwrap()
.unwrap()
.as_slice(),
b"backup-test-secret"
);
assert!(!fs::read(backup)
.unwrap()
.windows(18)
.any(|w| w == b"backup-test-secret"));
}
#[test]
fn session_revocation_denies_new_resolves_and_mutations() {
let temp = tempfile::tempdir().unwrap();
let mut broker = CredentialBroker::new(temp.path().join("credentials.v1"));
broker.unlock(password()).unwrap();
let id = CredentialId::legacy("test-provider");
broker
.put(&id, Zeroizing::new(b"test-value".to_vec()))
.unwrap();
broker.lock_signal().fetch_add(1, Ordering::SeqCst);
assert!(broker.is_locked());
assert_eq!(
broker.resolve(&Scope::Provider, &id).unwrap_err(),
"CREDENTIALS_LOCKED"
);
assert_eq!(
broker
.put(&id, Zeroizing::new(b"new-value".to_vec()))
.unwrap_err(),
"CREDENTIALS_LOCKED"
);
broker.unlock(password()).unwrap();
assert!(!broker.is_locked());
}
#[test]
fn python_fernet_migration_is_verified_idempotent_and_preserves_sources() {
let fixture: serde_json::Value =
serde_json::from_str(include_str!("../tests/fixtures/fernet-python.json")).unwrap();
frontend/src-tauri/src/lib.rs
+2
View File
@@ -6,4 +6,6 @@ pub mod recent;
#[cfg(feature = "desktop")]
pub mod request_lifecycle;
mod runtime_compat;
#[cfg(windows)]
pub mod session_lock;
pub mod workspace;
frontend/src-tauri/src/main.rs
+97
View File
@@ -22,6 +22,8 @@ struct Host {
recent: Mutex<Option<RecentVaultStore>>,
core: Arc<Mutex<Option<CoreSupervisor>>>,
credentials: Arc<Mutex<Option<CredentialBroker>>>,
#[cfg(windows)]
session_monitor: Mutex<Option<notesagent_host::session_lock::SessionMonitor>>,
streams: Arc<Mutex<HashMap<String, tauri::async_runtime::JoinHandle<()>>>>,
}
@@ -378,6 +380,15 @@ fn credentials_status(host: State<'_, Host>) -> Result<serde_json::Value, String
#[tauri::command]
async fn credentials_unlock(host: State<'_, Host>, password: String) -> Result<(), String> {
#[cfg(windows)]
if host
.session_monitor
.lock()
.map_err(|_| "HOST_BUSY")?
.is_none()
{
return Err("SESSION_MONITOR_UNAVAILABLE".into());
}
let broker = host.credentials.clone();
let password = Zeroizing::new(password.into_bytes());
tauri::async_runtime::spawn_blocking(move || {
@@ -455,6 +466,62 @@ async fn credentials_change_password(
.map_err(|_| "HOST_BUSY")?
}
#[tauri::command]
async fn credentials_backup(host: State<'_, Host>) -> Result<bool, String> {
let Some(path) = rfd::FileDialog::new()
.set_title("导出加密凭据备份（请选择新文件）")
.set_file_name("OpenNexus.onxcred")
.add_filter("OpenNexus credential backup", &["onxcred"])
.save_file()
else {
return Ok(false);
};
let broker = host.credentials.clone();
tauri::async_runtime::spawn_blocking(move || {
broker
.lock()
.map_err(|_| "HOST_BUSY")?
.as_ref()
.ok_or("HOST_NOT_READY")?
.backup(&path)?;
Ok(true)
})
.await
.map_err(|_| "HOST_BUSY")?
}
#[tauri::command]
async fn credentials_restore(
host: State<'_, Host>,
password: String,
) -> Result<Option<usize>, String> {
let password = Zeroizing::new(password.into_bytes());
let Some(path) = rfd::FileDialog::new()
.set_title("选择加密凭据备份")
.add_filter("OpenNexus credential backup", &["onxcred"])
.pick_file()
else {
return Ok(None);
};
if rfd::MessageDialog::new().set_title("恢复凭据备份")
.set_description("恢复将替换本机凭据库。当前加密文件会另存为恢复前备份；恢复后仍需解锁。笔记不会被替换。是否继续？")
.set_buttons(rfd::MessageButtons::YesNo).show() != rfd::MessageDialogResult::Yes {
return Ok(None);
}
let broker = host.credentials.clone();
tauri::async_runtime::spawn_blocking(move || {
broker
.lock()
.map_err(|_| "HOST_BUSY")?
.as_mut()
.ok_or("HOST_NOT_READY")?
.restore(&path, password)
.map(Some)
})
.await
.map_err(|_| "HOST_BUSY")?
}
#[tauri::command]
fn editor_capabilities(app: tauri::AppHandle, metadata_enabled: bool) -> Result<(), String> {
app.state::<tauri::menu::MenuItem<tauri::Wry>>()
@@ -602,6 +669,34 @@ fn main() {
.map_err(|_| std::io::Error::other("HOST_BUSY"))? = Some(CredentialBroker::new(
app.path().app_data_dir()?.join("credentials/stronghold.v1"),
));
#[cfg(windows)]
{
let signal = credential_state
.lock()
.map_err(|_| std::io::Error::other("HOST_BUSY"))?
.as_ref()
.ok_or_else(|| std::io::Error::other("HOST_NOT_READY"))?
.lock_signal();
*app.state::<Host>()
.session_monitor
.lock()
.map_err(|_| std::io::Error::other("HOST_BUSY"))? =
notesagent_host::session_lock::SessionMonitor::start(signal).ok();
}
let weak_credentials = Arc::downgrade(&credential_state);
std::thread::spawn(move || {
while let Some(state) = weak_credentials.upgrade() {
if let Ok(mut broker) = state.try_lock() {
if let Some(broker) = broker.as_mut() {
if broker.is_locked() {
broker.lock();
}
}
}
drop(state);
std::thread::sleep(Duration::from_millis(200));
}
});
let data_dir = app.path().app_data_dir()?.join("core-data");
// Debug builds use this worktree's interpreter; release builds only use bundled Core.
let core = if cfg!(debug_assertions) {
@@ -676,6 +771,8 @@ fn main() {
credentials_lock,
credentials_change_password,
credentials_import,
credentials_backup,
credentials_restore,
core_request,
core_request_prepare,
core_request_cancel,
frontend/src-tauri/src/session_lock.rs
+153
View File
@@ -0,0 +1,153 @@
//! Windows session notifications. Revocation is atomic and never waits for a KDF.
//! https://learn.microsoft.com/en-us/windows/win32/termserv/wm-wtssession-change
use std::cell::RefCell;
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use windows_sys::Win32::{
Foundation::*,
System::{LibraryLoader::GetModuleHandleW, RemoteDesktop::*},
UI::WindowsAndMessaging::*,
};
thread_local! { static SIGNAL: RefCell<Option<Arc<AtomicU64>>> = const { RefCell::new(None) }; }
unsafe extern "system" fn window_proc(
hwnd: HWND,
message: u32,
wparam: WPARAM,
lparam: LPARAM,
) -> LRESULT {
if message == WM_WTSSESSION_CHANGE
&& matches!(
wparam as u32,
WTS_SESSION_LOCK | WTS_SESSION_LOGOFF | WTS_CONSOLE_DISCONNECT | WTS_REMOTE_DISCONNECT
)
{
SIGNAL.with(|s| {
if let Some(signal) = s.borrow().as_ref() {
signal.fetch_add(1, Ordering::SeqCst);
}
});
}
if message == WM_DESTROY {
WTSUnRegisterSessionNotification(hwnd);
PostQuitMessage(0);
return 0;
}
DefWindowProcW(hwnd, message, wparam, lparam)
}
pub struct SessionMonitor {
window: usize,
thread: Option<std::thread::JoinHandle<()>>,
}
impl SessionMonitor {
pub fn start(signal: Arc<AtomicU64>) -> Result<Self, String> {
let (tx, rx) = std::sync::mpsc::sync_channel(1);
let thread = std::thread::spawn(move || unsafe {
SIGNAL.with(|s| *s.borrow_mut() = Some(signal));
let class: Vec<u16> = format!("OpenNexusSession-{}\0", uuid::Uuid::new_v4())
.encode_utf16()
.collect();
let module = GetModuleHandleW(std::ptr::null());
let descriptor = WNDCLASSW {
lpfnWndProc: Some(window_proc),
hInstance: module,
lpszClassName: class.as_ptr(),
..std::mem::zeroed()
};
if RegisterClassW(&descriptor) == 0 {
let _ = tx.send(Err("SESSION_MONITOR_UNAVAILABLE".to_string()));
return;
}
let window = CreateWindowExW(
0,
class.as_ptr(),
class.as_ptr(),
0,
0,
0,
0,
0,
std::ptr::null_mut(),
std::ptr::null_mut(),
module,
std::ptr::null(),
);
if window.is_null()
|| WTSRegisterSessionNotification(window, NOTIFY_FOR_THIS_SESSION) == 0
{
if !window.is_null() {
DestroyWindow(window);
}
UnregisterClassW(class.as_ptr(), module);
let _ = tx.send(Err("SESSION_MONITOR_UNAVAILABLE".to_string()));
return;
}
if tx.send(Ok(window as usize)).is_err() {
DestroyWindow(window);
} else {
let mut message: MSG = std::mem::zeroed();
while GetMessageW(&mut message, std::ptr::null_mut(), 0, 0) > 0 {
TranslateMessage(&message);
DispatchMessageW(&message);
}
}
UnregisterClassW(class.as_ptr(), module);
});
match rx
.recv()
.map_err(|_| "SESSION_MONITOR_UNAVAILABLE".to_string())?
{
Ok(window) => Ok(Self {
window,
thread: Some(thread),
}),
Err(error) => {
let _ = thread.join();
Err(error)
}
}
}
}
impl Drop for SessionMonitor {
fn drop(&mut self) {
unsafe {
PostMessageW(self.window as HWND, WM_CLOSE, 0, 0);
}
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn native_message_revokes_without_unlocking_on_session_return() {
let signal = Arc::new(AtomicU64::new(0));
let monitor = SessionMonitor::start(signal.clone()).unwrap();
// Inject only into our hidden test window; never lock the user's desktop.
unsafe {
SendMessageW(
monitor.window as HWND,
WM_WTSSESSION_CHANGE,
WTS_SESSION_LOCK as usize,
0,
);
}
assert_eq!(signal.load(Ordering::SeqCst), 1);
unsafe {
SendMessageW(
monitor.window as HWND,
WM_WTSSESSION_CHANGE,
WTS_SESSION_UNLOCK as usize,
0,
);
}
assert_eq!(signal.load(Ordering::SeqCst), 1);
}
}