refactor(network): split client.rs into pool, stream, headers and tests, and factor execute_stream

client.rs held the connection pool, the streaming response type, header
helpers, tests and a 274-line execute_stream in one 1220-line file (audit
RS-15). Move the self-contained parts into client/ and break
execute_stream into phase helpers so the redirect loop reads top-down.

Pure moves (verbatim apart from pub(super) on items the parent uses):
- client/pool.rs: Connector and Pool.
- client/stream.rs: ResponseBody and HttpStream (re-exported from client,
  so crate::HttpStream is unchanged).
- client/headers.rs: update_explicit_cookies, forward_sensitive_headers
  and canonical_header. error_chain stays in client.rs; it is not a
  header helper.
- client/tests.rs: the existing unit tests.

execute_stream helpers (base line ranges refer to main:client.rs). Each
body is the statements it replaces; the loop, await points, error
mapping, redirect limit and header mutations stay in the same order:
- prepare_request: 720-769, request validation, method, body, headers,
  user agent, default JSON content type and explicit referer, returned
  as PreparedRequest. Only the mut markers moved to the call site.
- outgoing_headers: 773-802, per-attempt cookie header and gzip
  negotiation, returning (outgoing, gzip). header(&url) became
  header(url) because url is now a reference.
- send: 803-845, the request, its 45 s header timeout and the native
  HTTPS-to-HTTP fallback retry. The permit is passed by value and
  returned, so drop(permit) still runs before the retry acquires a new
  one. matches!(method, ..) became matches!(*method, ..).
- redirect_target: 862-869, Location resolution, credential rejection
  and redirect policy validation, returning the target.
- release_redirect_body: 906-928, the 20 ms drain of a small HTTP/1.1
  redirect body.
- response_reader: 932-959, discard eligibility and gzip decoder setup,
  returning (headers, small_http1, reader).
- canonical_headers: 960-970, canonical response header map.

The redirect request rewrite (871-905) stays inline: it updates the
loop-carried method, body and headers and reads six other values, so a
helper would need nine parameters or a state struct that rewrites every
statement in the loop.

hyper is not a direct dependency, so the response type is named through
the pooled client's tower Service implementation (HyperResponse).
No public path or signature changes.
This commit is contained in:
zarzet committed 2026-10-19 00:00:04 +07:00
1 parent 12d422afa3
commit 42c90dbf62
5 files changed
+767 -658

No files matched your search

File diff suppressed because it is too large. Load diff
@@ -0,0 +1,57 @@
use super::*;
pub(super) fn update_explicit_cookies(headers: &mut HeaderMap, response: &HeaderMap) {
let Some(original) = headers
.get(header::COOKIE)
.and_then(|value| value.to_str().ok())
else {
return;
};
let names: Vec<_> = response
.get_all(header::SET_COOKIE)
.iter()
.filter_map(|value| value.to_str().ok())
.filter_map(|value| cookie::Cookie::parse(value).ok())
.map(|cookie| cookie.name().to_owned())
.collect();
let mut changed = false;
let mut remaining: Vec<_> = original
.split(';')
.map(str::trim)
.filter(|pair| {
let remove = pair
.split_once('=')
.is_some_and(|(name, _)| names.iter().any(|replaced| replaced == name));
changed |= remove;
!remove
})
.collect();
if changed {
remaining.sort_unstable();
let value = remaining
.join("; ")
.parse()
.expect("existing HTTP cookie header");
headers.insert(header::COOKIE, value);
}
}
pub(super) fn forward_sensitive_headers(original: &str, target: &str) -> bool {
let original = original.to_lowercase();
let target = target.to_lowercase();
original == target
|| (!target.contains([':', '%']) && target.ends_with(&format!(".{original}")))
}
pub(super) fn canonical_header(name: &str) -> String {
name.split('-')
.map(|word| {
let mut bytes = word.as_bytes().to_vec();
if let Some(first) = bytes.first_mut() {
first.make_ascii_uppercase();
}
String::from_utf8(bytes).expect("ASCII HTTP header name")
})
.collect::<Vec<_>>()
.join("-")
}
@@ -0,0 +1,101 @@
use super::*;
#[derive(Clone)]
pub(super) struct Connector(Transport);
impl Service<Uri> for Connector {
type Response = <Transport as Service<Uri>>::Response;
type Error = BoxError;
type Future = BoxFuture<'static, Result<Self::Response, BoxError>>;
fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.0.poll_ready(context)
}
fn call(&mut self, uri: Uri) -> Self::Future {
let future = self.0.call(uri);
Box::pin(async move {
tokio::time::timeout(Duration::from_secs(10), future)
.await
.map_err(|_| {
BoxError::from(io::Error::new(
io::ErrorKind::TimedOut,
"connection timeout",
))
})?
})
}
}
pub(super) struct Pool {
pub(super) client: Client<Connector, Full<Bytes>>,
pub(super) allow_private: bool,
pub(super) allow_http_fallback: bool,
permits: Mutex<HashMap<String, Weak<Semaphore>>>,
}
impl Pool {
pub(super) fn new(
tls: rustls::ClientConfig,
dns: Arc<Dns>,
allow_private: bool,
allow_http_fallback: bool,
) -> Self {
let mut http = HttpConnector::new_with_resolver(Resolver { dns, allow_private });
http.enforce_http(false);
http.set_connect_timeout(Some(Duration::from_secs(10)));
http.set_happy_eyeballs_timeout(Some(Duration::from_millis(300)));
http.set_keepalive(Some(Duration::from_secs(30)));
// Go's dialer default: small request/control writes are not held back
// by Nagle while waiting for a delayed ACK.
http.set_nodelay(true);
let connector = HttpsConnectorBuilder::new()
.with_tls_config(tls)
.https_or_http()
.enable_http1()
.enable_http2()
.wrap_connector(http);
let client = Client::builder(TokioExecutor::new())
.pool_timer(TokioTimer::new())
.pool_idle_timeout(Duration::from_secs(60))
.pool_max_idle_per_host(10)
.http1_max_buf_size(10 << 20)
// Go's 4 MiB per-stream window doubles one stream's ceiling per
// RTT. The 5 MiB connection window, and so peak buffering per
// connection, stays hyper's default.
.http2_initial_stream_window_size(4 << 20)
.build(Connector(connector));
Self {
client,
allow_private,
allow_http_fallback,
permits: Mutex::default(),
}
}
pub(super) async fn acquire(&self, uri: &Uri) -> OwnedSemaphorePermit {
let key = format!(
"{}://{}",
uri.scheme_str().unwrap_or(""),
uri.authority().map_or("", |authority| authority.as_str())
);
let semaphore = {
let mut permits = self.permits.lock().expect("HTTP connection limit lock");
if permits.len() > 256 {
permits.retain(|_, semaphore| semaphore.strong_count() > 0);
}
match permits.get(&key).and_then(Weak::upgrade) {
Some(semaphore) => semaphore,
None => {
let semaphore = Arc::new(Semaphore::new(20));
permits.insert(key, Arc::downgrade(&semaphore));
semaphore
}
}
};
semaphore
.acquire_owned()
.await
.expect("HTTP semaphore is never closed")
}
}
@@ -0,0 +1,125 @@
use super::*;
pub(super) struct ResponseBody {
pub(super) reader: Box<dyn AsyncRead + Unpin + Send>,
pub(super) small_http1: bool,
pub(super) _permit: OwnedSemaphorePermit,
}
/// A pull-based media response. Read and progress callbacks run on the caller's
/// worker; JavaScript never executes inside the Tokio executor. Dropping this
/// response immediately releases its body and per-origin permit.
pub struct HttpStream {
/// Response metadata. The body is delivered exclusively through `read`.
pub response: HttpResponse,
pub(super) body: Option<ResponseBody>,
pub(super) service: Arc<NetworkService>,
pub(super) _pool: Arc<Pool>,
pub(super) generation: u64,
pub(super) started: Instant,
pub(super) last_progress: Instant,
pub(super) timeout: Duration,
pub(super) stall_timeout: Duration,
pub(super) failure: Option<String>,
}
impl HttpStream {
/// Read at most 64 KiB without collecting the response in memory. Errors
/// close the stream so callers cannot accidentally resume a failed reader.
pub fn read(
&mut self,
buffer: &mut [u8],
check: impl Fn() -> Result<(), String>,
) -> Result<usize, String> {
if let Some(error) = &self.failure {
return Err(error.clone());
}
let result = (|| {
check()?;
if self.service.generation.load(Ordering::Acquire) != self.generation {
return Err("network policy changed".into());
}
let Some(body) = self.body.as_mut() else {
return Ok(0);
};
if buffer.is_empty() {
return Ok(0);
}
let total = self.timeout.saturating_sub(self.started.elapsed());
let idle = self
.stall_timeout
.saturating_sub(self.last_progress.elapsed());
if total.is_zero() {
return Err("HTTP request timeout exceeded".into());
}
if idle.is_zero() {
return Err(self.stall_message());
}
let size = buffer.len().min(64 << 10);
let count = self
.service
.run(self.generation, total.min(idle), &check, async {
body.reader
.read(&mut buffer[..size])
.await
.map_err(|error| error.to_string())
})
.map_err(|error| {
if error == "HTTP request timeout exceeded" && idle <= total {
self.stall_message()
} else {
error
}
})?;
if count > 0 {
self.last_progress = Instant::now();
}
Ok(count)
})();
match &result {
Err(error) => {
self.body = None;
self.failure = Some(error.clone());
}
Ok(0) if !buffer.is_empty() => self.body = None,
_ => {}
}
result
}
/// Release an unwanted body (a range probe or retryable error). A small
/// HTTP/1.1 body is finished within 20 ms so its connection can return to
/// the pool; HTTP/2, compressed, large or slow bodies are dropped at once.
pub fn discard(&mut self, check: impl Fn() -> Result<(), String>) {
let Some(mut body) = self.body.take() else {
return;
};
if !body.small_http1 || self.failure.is_some() {
return;
}
let _ = self
.service
.run(self.generation, Duration::from_millis(20), &check, async {
let mut buffer = [0; 2048];
let mut received = 0;
loop {
let count = body
.reader
.read(&mut buffer)
.await
.map_err(|error| error.to_string())?;
received += count;
if count == 0 || received > 2048 {
return Ok(());
}
}
});
}
fn stall_message(&self) -> String {
format!(
"download stalled: no data received for {}s (network timeout)",
self.stall_timeout.as_secs()
)
}
}
@@ -0,0 +1,163 @@
use super::*;
use std::future::poll_fn;
use std::sync::{atomic::AtomicBool, mpsc};
#[test]
fn direct_media_keeps_manifest_permissions_and_https_requirement() {
let service = NetworkService::new().unwrap();
let original = service.session(
NetworkPermissions {
domains: vec!["cdn.example.test".into()],
allow_http: false,
},
Duration::from_secs(30),
);
let media = original.direct_media();
assert!(media.validate_url("https://cdn.example.test/audio").is_ok());
assert!(
media
.validate_url("https://unlisted.example.test/audio")
.is_err()
);
assert!(media.validate_url("http://cdn.example.test/audio").is_err());
assert!(!media.follow_redirects);
assert!(media.native_media);
}
#[test]
fn direct_media_does_not_send_metadata_cookies_or_follow_redirects() {
use std::io::{Read, Write};
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let server = std::thread::spawn(move || {
for response in [
"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok",
"HTTP/1.1 302 Found\r\nLocation: /unexpected\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
] {
let (mut stream, _) = listener.accept().unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(3)))
.unwrap();
let mut request = Vec::new();
while !request.ends_with(b"\r\n\r\n") {
let mut byte = [0];
stream.read_exact(&mut byte).unwrap();
request.push(byte[0]);
assert!(request.len() < 8192);
}
assert!(
!String::from_utf8(request)
.unwrap()
.to_lowercase()
.contains("\r\ncookie:")
);
stream.write_all(response.as_bytes()).unwrap();
}
});
let service = NetworkService::new().unwrap();
service.set_allow_private_network(true);
let original = service.session(
NetworkPermissions {
domains: vec!["127.0.0.1".into()],
allow_http: true,
},
Duration::from_secs(3),
);
let mut cookies = http::HeaderMap::new();
cookies.insert(
http::header::SET_COOKIE,
"metadata=private; Path=/".parse().unwrap(),
);
original
.cookies
.lock()
.unwrap()
.store(&UrlParts::parse(&base).unwrap(), &cookies);
let media = original.direct_media();
let request =
|| serde_json::from_value(serde_json::json!({"url":format!("{base}/media")})).unwrap();
assert_eq!(media.request(request(), || Ok(())).unwrap().body, b"ok");
assert!(
media
.request(request(), || Ok(()))
.unwrap_err()
.contains("media redirect blocked")
);
server.join().unwrap();
}
#[test]
fn cancellation_after_operation_is_pending_drops_it_before_return() {
struct Dropped<'a>(&'a AtomicBool);
impl Drop for Dropped<'_> {
fn drop(&mut self) {
self.0.store(true, Ordering::Release);
}
}
let service = NetworkService::new().unwrap();
// Sample two phases; timings are diagnostics, not shared-host limits.
for (waiting_for, wake) in [
(Duration::ZERO, false),
(Duration::from_millis(15), false),
(Duration::from_millis(15), true),
] {
let cancelled = AtomicBool::new(false);
let dropped = AtomicBool::new(false);
let (waiting_tx, waiting_rx) = mpsc::sync_channel(1);
let (finished_tx, finished_rx) = mpsc::sync_channel(1);
std::thread::scope(|scope| {
let worker = scope.spawn(|| {
let mut waiting = Some(waiting_tx);
let result = service.run(
0,
Duration::from_secs(2),
|| {
if cancelled.load(Ordering::Acquire) {
Err("cancelled".into())
} else {
Ok(())
}
},
async {
let _dropped = Dropped(&dropped);
if !waiting_for.is_zero() {
tokio::time::sleep(waiting_for).await;
}
poll_fn(|_| {
if let Some(waiting) = waiting.take() {
waiting.send(()).unwrap();
}
Poll::<Result<(), String>>::Pending
})
.await
},
);
finished_tx
.send((result, Instant::now(), dropped.load(Ordering::Acquire)))
.unwrap();
});
waiting_rx.recv_timeout(Duration::from_secs(1)).unwrap();
let began = Instant::now();
cancelled.store(true, Ordering::Release);
if wake {
service.cancellation_waker().wake();
}
let (result, returned, dropped_before_return) =
finished_rx.recv_timeout(Duration::from_secs(1)).unwrap();
worker.join().unwrap();
assert_eq!(result, Err("cancelled".into()));
assert!(dropped_before_return);
println!(
"pending operation cancellation (wait={}ms, wake={wake}): return={}us join={}us",
waiting_for.as_millis(),
returned.duration_since(began).as_micros(),
began.elapsed().as_micros(),
);
});
}
assert_eq!(
service.run(0, Duration::from_secs(1), || Ok(()), async { Ok(7) }),
Ok(7)
);
}