fix: 串行化凭据所有权并响应桌面请求取消
This commit is contained in:
@@ -201,6 +201,9 @@ impl Unlocked {
|
||||
pub struct CredentialBroker {
|
||||
path: PathBuf,
|
||||
unlocked: Option<Unlocked>,
|
||||
// Separate stable inode: snapshots are atomically replaced, so locking the
|
||||
// snapshot itself would not protect the next writer after replacement.
|
||||
ownership: Option<fs::File>,
|
||||
}
|
||||
|
||||
impl CredentialBroker {
|
||||
@@ -479,6 +482,7 @@ impl CredentialBroker {
|
||||
Self {
|
||||
path,
|
||||
unlocked: None,
|
||||
ownership: None,
|
||||
}
|
||||
}
|
||||
pub fn is_locked(&self) -> bool {
|
||||
@@ -486,9 +490,50 @@ impl CredentialBroker {
|
||||
}
|
||||
pub fn lock(&mut self) {
|
||||
self.unlocked.take();
|
||||
self.ownership.take();
|
||||
}
|
||||
pub fn unlock(&mut self, password: Zeroizing<Vec<u8>>) -> Result<()> {
|
||||
self.lock();
|
||||
use fs2::FileExt;
|
||||
let parent = self.path.parent().ok_or("CREDENTIAL_PATH_INVALID")?;
|
||||
fs::create_dir_all(parent).map_err(|_| "CREDENTIAL_IO_FAILED")?;
|
||||
let mut lock_name = self
|
||||
.path
|
||||
.file_name()
|
||||
.ok_or("CREDENTIAL_PATH_INVALID")?
|
||||
.to_os_string();
|
||||
lock_name.push(".lock");
|
||||
let lock_path = parent.join(lock_name);
|
||||
if let Ok(metadata) = fs::symlink_metadata(&lock_path) {
|
||||
if !metadata.is_file() || metadata.file_type().is_symlink() {
|
||||
return Err("CREDENTIAL_PATH_INVALID".into());
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
use std::os::windows::fs::MetadataExt;
|
||||
if metadata.file_attributes() & 0x400 != 0 {
|
||||
return Err("CREDENTIAL_PATH_INVALID".into());
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut options = fs::OpenOptions::new();
|
||||
options.read(true).write(true).create(true).truncate(false);
|
||||
#[cfg(windows)]
|
||||
{
|
||||
use std::os::windows::fs::OpenOptionsExt;
|
||||
options.share_mode(0x1 | 0x2); // Do not allow replacing the held lock file.
|
||||
}
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
let ownership = options
|
||||
.open(&lock_path)
|
||||
.map_err(|_| "CREDENTIAL_IO_FAILED")?;
|
||||
ownership
|
||||
.try_lock_exclusive()
|
||||
.map_err(|_| "CREDENTIALS_BUSY")?;
|
||||
let session = if self.path.exists() {
|
||||
let metadata = fs::symlink_metadata(&self.path).map_err(|_| "CREDENTIAL_IO_FAILED")?;
|
||||
if !metadata.is_file() || metadata.len() > MAX_FILE {
|
||||
@@ -530,6 +575,7 @@ impl CredentialBroker {
|
||||
session
|
||||
};
|
||||
self.unlocked = Some(session);
|
||||
self.ownership = Some(ownership);
|
||||
Ok(())
|
||||
}
|
||||
pub fn list(&self) -> Result<Vec<CredentialId>> {
|
||||
|
||||
@@ -3,5 +3,7 @@
|
||||
pub mod core;
|
||||
pub mod credentials;
|
||||
pub mod recent;
|
||||
#[cfg(feature = "desktop")]
|
||||
pub mod request_lifecycle;
|
||||
mod runtime_compat;
|
||||
pub mod workspace;
|
||||
|
||||
+152
-100
@@ -6,6 +6,7 @@ use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
|
||||
use notesagent_host::core::CoreSupervisor;
|
||||
use notesagent_host::credentials::CredentialBroker;
|
||||
use notesagent_host::recent::{RecentVault, RecentVaultStore};
|
||||
use notesagent_host::request_lifecycle::Requests;
|
||||
use notesagent_host::workspace::{portable_path_string, Document, Entry, Workspace};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
@@ -16,6 +17,7 @@ use zeroize::Zeroizing;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Host {
|
||||
requests: Requests,
|
||||
workspace: Mutex<Option<Workspace>>,
|
||||
recent: Mutex<Option<RecentVaultStore>>,
|
||||
core: Arc<Mutex<Option<CoreSupervisor>>>,
|
||||
@@ -87,6 +89,14 @@ fn core_url(path: &str) -> Result<String, String> {
|
||||
mod core_proxy_tests {
|
||||
use super::{core_url, is_json_content_type};
|
||||
|
||||
#[test]
|
||||
fn request_dto_accepts_camel_case_and_rejects_unowned_headers() {
|
||||
let mut payload = serde_json::json!({"requestId":"fixture-reservation","method":"POST","path":"/api/tasks","body":{"title":"fixture"},"contentType":"application/json"});
|
||||
assert!(serde_json::from_value::<super::CoreRequest>(payload.clone()).is_ok());
|
||||
payload["authorization"] = serde_json::json!("must-not-be-forwarded");
|
||||
assert!(serde_json::from_value::<super::CoreRequest>(payload).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_allows_expected_loopback_paths() {
|
||||
assert_eq!(core_url("/health").unwrap(), "http://127.0.0.1:8000/health");
|
||||
@@ -118,115 +128,155 @@ mod core_proxy_tests {
|
||||
|
||||
/// Authenticated process-local transport; session headers are owned by Rust.
|
||||
#[tauri::command]
|
||||
async fn core_request(
|
||||
fn core_request_prepare(host: State<'_, Host>, timeout_ms: u64) -> Result<String, String> {
|
||||
host.requests.prepare(timeout_ms)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn core_request_cancel(host: State<'_, Host>, request_id: String) -> Result<(), String> {
|
||||
host.requests.cancel(&request_id)
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct CoreRequest {
|
||||
request_id: String,
|
||||
method: String,
|
||||
path: String,
|
||||
body: Option<serde_json::Value>,
|
||||
body_base64: Option<String>,
|
||||
content_type: Option<String>,
|
||||
idempotency_key: Option<String>,
|
||||
host: State<'_, Host>,
|
||||
) -> Result<CoreResponse, String> {
|
||||
let core = host.core.clone();
|
||||
let core_path = path.clone();
|
||||
let session = tauri::async_runtime::spawn_blocking(move || {
|
||||
core.lock()
|
||||
.map_err(|_| "HOST_BUSY")?
|
||||
.as_mut()
|
||||
.ok_or("CORE_UNAVAILABLE")?
|
||||
.request_session(&core_path)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "CORE_UNAVAILABLE")??;
|
||||
let method =
|
||||
reqwest::Method::from_bytes(method.as_bytes()).map_err(|_| "CORE_METHOD_DENIED")?;
|
||||
if !matches!(
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn core_request(request: CoreRequest, host: State<'_, Host>) -> Result<CoreResponse, String> {
|
||||
let CoreRequest {
|
||||
request_id,
|
||||
method,
|
||||
reqwest::Method::GET
|
||||
| reqwest::Method::POST
|
||||
| reqwest::Method::PUT
|
||||
| reqwest::Method::PATCH
|
||||
| reqwest::Method::DELETE
|
||||
) {
|
||||
return Err("CORE_METHOD_DENIED".into());
|
||||
}
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.no_proxy()
|
||||
.build()
|
||||
.map_err(|_| "CORE_CLIENT_ERROR")?;
|
||||
let mut request = client
|
||||
.request(method, &session.url)
|
||||
.header(
|
||||
reqwest::header::AUTHORIZATION,
|
||||
session.authorization.as_str(),
|
||||
)
|
||||
.header("X-Core-Generation", &session.generation);
|
||||
if let Some(value) = body {
|
||||
request = request.json(&value);
|
||||
}
|
||||
if let Some(encoded) = body_base64 {
|
||||
if encoded.len() > MAX_CORE_RESPONSE_BYTES * 4 / 3 + 4 {
|
||||
return Err("CORE_REQUEST_TOO_LARGE".into());
|
||||
}
|
||||
let bytes = BASE64_STANDARD
|
||||
.decode(encoded)
|
||||
.map_err(|_| "CORE_BODY_INVALID")?;
|
||||
if bytes.len() > MAX_CORE_RESPONSE_BYTES {
|
||||
return Err("CORE_REQUEST_TOO_LARGE".into());
|
||||
}
|
||||
let content_type = content_type
|
||||
.as_deref()
|
||||
.unwrap_or("application/octet-stream");
|
||||
if !matches!(content_type, "application/octet-stream" | "application/zip") {
|
||||
return Err("CORE_CONTENT_TYPE_DENIED".into());
|
||||
}
|
||||
request = request
|
||||
.header(reqwest::header::CONTENT_TYPE, content_type)
|
||||
.body(bytes);
|
||||
}
|
||||
if let Some(key) = idempotency_key {
|
||||
if key.len() > 128 || !key.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-') {
|
||||
return Err("CORE_HEADER_INVALID".into());
|
||||
}
|
||||
request = request.header("Idempotency-Key", key);
|
||||
}
|
||||
let mut response = request.send().await.map_err(|_| "CORE_UNAVAILABLE")?;
|
||||
let status = response.status().as_u16();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|length| length > MAX_CORE_RESPONSE_BYTES as u64)
|
||||
{
|
||||
return Err("CORE_RESPONSE_TOO_LARGE".into());
|
||||
}
|
||||
let mut bytes = Vec::new();
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| "CORE_RESPONSE_ERROR")? {
|
||||
if bytes.len().saturating_add(chunk.len()) > MAX_CORE_RESPONSE_BYTES {
|
||||
return Err("CORE_RESPONSE_TOO_LARGE".into());
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
let (body, body_base64) = if is_json_content_type(&content_type) {
|
||||
(
|
||||
String::from_utf8(bytes.to_vec()).map_err(|_| "CORE_RESPONSE_ERROR")?,
|
||||
None,
|
||||
)
|
||||
} else {
|
||||
(String::new(), Some(BASE64_STANDARD.encode(&bytes)))
|
||||
};
|
||||
Ok(CoreResponse {
|
||||
status,
|
||||
content_type,
|
||||
path,
|
||||
body,
|
||||
body_base64,
|
||||
})
|
||||
content_type,
|
||||
idempotency_key,
|
||||
} = request;
|
||||
let mut lease = host.requests.claim(&request_id)?;
|
||||
let checkpoint = lease.checkpoint();
|
||||
lease
|
||||
.run(async {
|
||||
if body.is_some() && body_base64.is_some() {
|
||||
return Err("CORE_BODY_INVALID".into());
|
||||
}
|
||||
if body
|
||||
.as_ref()
|
||||
.is_some_and(|value| value.to_string().len() > MAX_CORE_RESPONSE_BYTES)
|
||||
{
|
||||
return Err("CORE_REQUEST_TOO_LARGE".into());
|
||||
}
|
||||
let core = host.core.clone();
|
||||
let core_path = path.clone();
|
||||
let session = tauri::async_runtime::spawn_blocking(move || {
|
||||
core.lock()
|
||||
.map_err(|_| "HOST_BUSY")?
|
||||
.as_mut()
|
||||
.ok_or("CORE_UNAVAILABLE")?
|
||||
.request_session(&core_path)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "CORE_UNAVAILABLE")??;
|
||||
let method =
|
||||
reqwest::Method::from_bytes(method.as_bytes()).map_err(|_| "CORE_METHOD_DENIED")?;
|
||||
if !matches!(
|
||||
method,
|
||||
reqwest::Method::GET
|
||||
| reqwest::Method::POST
|
||||
| reqwest::Method::PUT
|
||||
| reqwest::Method::PATCH
|
||||
| reqwest::Method::DELETE
|
||||
) {
|
||||
return Err("CORE_METHOD_DENIED".into());
|
||||
}
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(600))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.no_proxy()
|
||||
.build()
|
||||
.map_err(|_| "CORE_CLIENT_ERROR")?;
|
||||
let mut request = client
|
||||
.request(method, &session.url)
|
||||
.header(
|
||||
reqwest::header::AUTHORIZATION,
|
||||
session.authorization.as_str(),
|
||||
)
|
||||
.header("X-Core-Generation", &session.generation)
|
||||
.header("X-Request-Id", &request_id);
|
||||
if let Some(value) = body {
|
||||
request = request.json(&value);
|
||||
}
|
||||
if let Some(encoded) = body_base64 {
|
||||
if encoded.len() > MAX_CORE_RESPONSE_BYTES * 4 / 3 + 4 {
|
||||
return Err("CORE_REQUEST_TOO_LARGE".into());
|
||||
}
|
||||
let bytes = BASE64_STANDARD
|
||||
.decode(encoded)
|
||||
.map_err(|_| "CORE_BODY_INVALID")?;
|
||||
if bytes.len() > MAX_CORE_RESPONSE_BYTES {
|
||||
return Err("CORE_REQUEST_TOO_LARGE".into());
|
||||
}
|
||||
let content_type = content_type
|
||||
.as_deref()
|
||||
.unwrap_or("application/octet-stream");
|
||||
if !matches!(content_type, "application/octet-stream" | "application/zip") {
|
||||
return Err("CORE_CONTENT_TYPE_DENIED".into());
|
||||
}
|
||||
request = request
|
||||
.header(reqwest::header::CONTENT_TYPE, content_type)
|
||||
.body(bytes);
|
||||
}
|
||||
if let Some(key) = idempotency_key {
|
||||
if key.len() > 128 || !key.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-') {
|
||||
return Err("CORE_HEADER_INVALID".into());
|
||||
}
|
||||
request = request.header("Idempotency-Key", key);
|
||||
}
|
||||
checkpoint()?;
|
||||
let mut response = request.send().await.map_err(|_| "CORE_UNAVAILABLE")?;
|
||||
let status = response.status().as_u16();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|length| length > MAX_CORE_RESPONSE_BYTES as u64)
|
||||
{
|
||||
return Err("CORE_RESPONSE_TOO_LARGE".into());
|
||||
}
|
||||
let mut bytes = Vec::new();
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| "CORE_RESPONSE_ERROR")? {
|
||||
if bytes.len().saturating_add(chunk.len()) > MAX_CORE_RESPONSE_BYTES {
|
||||
return Err("CORE_RESPONSE_TOO_LARGE".into());
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
let (body, body_base64) = if is_json_content_type(&content_type) {
|
||||
(
|
||||
String::from_utf8(bytes.to_vec()).map_err(|_| "CORE_RESPONSE_ERROR")?,
|
||||
None,
|
||||
)
|
||||
} else {
|
||||
(String::new(), Some(BASE64_STANDARD.encode(&bytes)))
|
||||
};
|
||||
Ok(CoreResponse {
|
||||
status,
|
||||
content_type,
|
||||
body,
|
||||
body_base64,
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
@@ -627,6 +677,8 @@ fn main() {
|
||||
credentials_change_password,
|
||||
credentials_import,
|
||||
core_request,
|
||||
core_request_prepare,
|
||||
core_request_cancel,
|
||||
core_stream,
|
||||
core_stream_cancel,
|
||||
editor_capabilities,
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
//! Reserve before dispatch so cancellation cannot race a delayed IPC invocation.
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
future::Future,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
use tokio::sync::watch;
|
||||
|
||||
struct Entry {
|
||||
deadline: Instant,
|
||||
claimed: bool,
|
||||
cancel: watch::Sender<bool>,
|
||||
}
|
||||
#[derive(Default, Clone)]
|
||||
pub struct Requests {
|
||||
entries: Arc<Mutex<HashMap<String, Entry>>>,
|
||||
}
|
||||
pub struct Lease {
|
||||
id: String,
|
||||
owner: Requests,
|
||||
deadline: Instant,
|
||||
cancel: watch::Receiver<bool>,
|
||||
}
|
||||
|
||||
impl Requests {
|
||||
pub fn prepare(&self, timeout_ms: u64) -> Result<String, String> {
|
||||
if !(1..=600_000).contains(&timeout_ms) {
|
||||
return Err("CORE_TIMEOUT_INVALID".into());
|
||||
}
|
||||
let mut entries = self.entries.lock().map_err(|_| "HOST_BUSY")?;
|
||||
entries.retain(|_, entry| entry.claimed || entry.deadline > Instant::now());
|
||||
if entries.len() >= 64 {
|
||||
return Err("CORE_REQUEST_LIMIT".into());
|
||||
}
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let (cancel, _) = watch::channel(false);
|
||||
entries.insert(
|
||||
id.clone(),
|
||||
Entry {
|
||||
deadline: Instant::now() + Duration::from_millis(timeout_ms),
|
||||
claimed: false,
|
||||
cancel,
|
||||
},
|
||||
);
|
||||
Ok(id)
|
||||
}
|
||||
pub fn claim(&self, id: &str) -> Result<Lease, String> {
|
||||
let mut entries = self.entries.lock().map_err(|_| "HOST_BUSY")?;
|
||||
let entry = entries.get_mut(id).ok_or("CORE_REQUEST_NOT_PREPARED")?;
|
||||
if entry.claimed {
|
||||
return Err("CORE_REQUEST_ALREADY_STARTED".into());
|
||||
}
|
||||
if entry.deadline <= Instant::now() {
|
||||
entries.remove(id);
|
||||
return Err("REQUEST_TIMEOUT".into());
|
||||
}
|
||||
entry.claimed = true;
|
||||
Ok(Lease {
|
||||
id: id.into(),
|
||||
owner: self.clone(),
|
||||
deadline: entry.deadline,
|
||||
cancel: entry.cancel.subscribe(),
|
||||
})
|
||||
}
|
||||
pub fn cancel(&self, id: &str) -> Result<(), String> {
|
||||
if let Some(entry) = self.entries.lock().map_err(|_| "HOST_BUSY")?.remove(id) {
|
||||
entry.cancel.send_replace(true);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl Lease {
|
||||
/// Recheck after synchronous encoding/validation, immediately before network IO.
|
||||
pub fn checkpoint(&self) -> impl Fn() -> Result<(), String> + Send + 'static {
|
||||
let cancel = self.cancel.clone();
|
||||
let deadline = self.deadline;
|
||||
move || {
|
||||
if *cancel.borrow() {
|
||||
return Err("REQUEST_CANCELLED".into());
|
||||
}
|
||||
if deadline <= Instant::now() {
|
||||
return Err("REQUEST_TIMEOUT".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
pub async fn run<T>(
|
||||
&mut self,
|
||||
operation: impl Future<Output = Result<T, String>>,
|
||||
) -> Result<T, String> {
|
||||
// Check current state before polling an operation with possible side effects.
|
||||
if *self.cancel.borrow() {
|
||||
return Err("REQUEST_CANCELLED".into());
|
||||
}
|
||||
if self.deadline <= Instant::now() {
|
||||
return Err("REQUEST_TIMEOUT".into());
|
||||
}
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = self.cancel.changed() => Err("REQUEST_CANCELLED".into()),
|
||||
_ = tokio::time::sleep_until(self.deadline.into()) => Err("REQUEST_TIMEOUT".into()),
|
||||
value = operation => value,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl Drop for Lease {
|
||||
fn drop(&mut self) {
|
||||
if let Ok(mut entries) = self.owner.entries.lock() {
|
||||
entries.remove(&self.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[tokio::test]
|
||||
async fn cancellation_before_dispatch_and_replay_never_run_work() {
|
||||
let requests = Requests::default();
|
||||
let id = requests.prepare(1000).unwrap();
|
||||
requests.cancel(&id).unwrap();
|
||||
assert!(requests.claim(&id).is_err());
|
||||
let id = requests.prepare(1000).unwrap();
|
||||
let mut lease = requests.claim(&id).unwrap();
|
||||
assert!(requests.claim(&id).is_err());
|
||||
requests.cancel(&id).unwrap();
|
||||
let result: Result<(), String> = lease
|
||||
.run(async { panic!("cancelled work was polled") })
|
||||
.await;
|
||||
assert_eq!(result.unwrap_err(), "REQUEST_CANCELLED");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn timeout_and_cancel_drop_inflight_work_and_release_capacity() {
|
||||
let requests = Requests::default();
|
||||
let id = requests.prepare(10).unwrap();
|
||||
let mut lease = requests.claim(&id).unwrap();
|
||||
assert_eq!(
|
||||
lease
|
||||
.run(std::future::pending::<Result<(), String>>())
|
||||
.await
|
||||
.unwrap_err(),
|
||||
"REQUEST_TIMEOUT"
|
||||
);
|
||||
drop(lease);
|
||||
assert!(requests.entries.lock().unwrap().is_empty());
|
||||
let id = requests.prepare(1000).unwrap();
|
||||
let mut lease = requests.claim(&id).unwrap();
|
||||
let (result, ()) = tokio::join!(
|
||||
lease.run(std::future::pending::<Result<(), String>>()),
|
||||
async {
|
||||
tokio::task::yield_now().await;
|
||||
requests.cancel(&id).unwrap();
|
||||
}
|
||||
);
|
||||
assert_eq!(result.unwrap_err(), "REQUEST_CANCELLED");
|
||||
drop(lease);
|
||||
assert!(requests.entries.lock().unwrap().is_empty());
|
||||
}
|
||||
#[test]
|
||||
fn capacity_and_expired_reservations_are_bounded() {
|
||||
let requests = Requests::default();
|
||||
for _ in 0..64 {
|
||||
requests.prepare(1000).unwrap();
|
||||
}
|
||||
assert_eq!(requests.prepare(1000).unwrap_err(), "CORE_REQUEST_LIMIT");
|
||||
for entry in requests.entries.lock().unwrap().values_mut() {
|
||||
entry.deadline = Instant::now();
|
||||
}
|
||||
assert!(requests.prepare(1000).is_ok());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn cancelling_a_real_response_closes_its_socket() {
|
||||
use std::io::{Read, Write};
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let endpoint = listener.local_addr().unwrap();
|
||||
let (started, received) = tokio::sync::oneshot::channel();
|
||||
let server = std::thread::spawn(move || {
|
||||
let (mut socket, _) = listener.accept().unwrap();
|
||||
socket
|
||||
.set_read_timeout(Some(Duration::from_secs(5)))
|
||||
.unwrap();
|
||||
let mut request = Vec::new();
|
||||
while !request.ends_with(b"\r\n\r\n") {
|
||||
let mut byte = [0];
|
||||
socket.read_exact(&mut byte).unwrap();
|
||||
request.push(byte[0]);
|
||||
}
|
||||
socket
|
||||
.write_all(
|
||||
b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\nx",
|
||||
)
|
||||
.unwrap();
|
||||
started.send(()).unwrap();
|
||||
match socket.read(&mut [0]) {
|
||||
Ok(0) => true,
|
||||
Err(error) => matches!(
|
||||
error.kind(),
|
||||
std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
|
||||
),
|
||||
_ => false,
|
||||
}
|
||||
});
|
||||
let requests = Requests::default();
|
||||
let id = requests.prepare(10000).unwrap();
|
||||
let mut lease = requests.claim(&id).unwrap();
|
||||
let operation = async {
|
||||
let mut response = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.unwrap()
|
||||
.get(format!("http://{endpoint}/"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| "HTTP_FAILURE".to_string())?;
|
||||
while response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|_| "BODY_FAILURE".to_string())?
|
||||
.is_some()
|
||||
{}
|
||||
Ok(())
|
||||
};
|
||||
let (result, ()) = tokio::join!(lease.run(operation), async {
|
||||
received.await.unwrap();
|
||||
requests.cancel(&id).unwrap();
|
||||
});
|
||||
assert_eq!(result.unwrap_err(), "REQUEST_CANCELLED");
|
||||
assert!(
|
||||
tokio::task::spawn_blocking(move || server.join().unwrap())
|
||||
.await
|
||||
.unwrap(),
|
||||
"HTTP socket was not closed on cancellation"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user