From 1504539b33ae8b265a3c76e39d4f98d7c39188f1 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 18 Feb 2025 19:33:07 +0800 Subject: [PATCH] feat(client): Optionally enable asynchronous DNS resolver --- Cargo.toml | 18 +++++++++++++---- examples/client.py | 1 + src/client.rs | 10 +++++++++- src/dns.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++ src/error.rs | 1 + src/lib.rs | 1 + src/param/client.rs | 5 +++++ 7 files changed, 79 insertions(+), 5 deletions(-) create mode 100644 src/dns.rs diff --git a/Cargo.toml b/Cargo.toml index 47e1a330..a9e24c9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,14 @@ doctest = false [dependencies] tokio = { version = "1.0", features = ["sync"] } -pyo3 = { version = "0.23.0", features = ["anyhow", "indexmap", "multiple-pymethods", "abi3-py39", "generate-import-lib", "experimental-inspect"] } +pyo3 = { version = "0.23.0", features = [ + "anyhow", + "indexmap", + "multiple-pymethods", + "abi3-py39", + "generate-import-lib", + "experimental-inspect", +] } pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } pyo3-stub-gen = "0.7.0" anyhow = "1.0" @@ -27,15 +34,18 @@ indexmap = { version = "2.7.0", features = ["serde"] } cookie = "0.18.0" arc-swap = "1.7.1" url = "2.5" -rquest = { version = "2.1.5", features = ["full", "websocket"] } +rquest = { version = "2.1.5", features = ["full", "websocket", "hickory-dns"] } futures-util = { version = "0.3.0", default-features = false } [target.'cfg(not(target_env = "msvc"))'.dependencies] -jemallocator = { package = "tikv-jemallocator", version = "0.6", features = ["disable_initial_exec_tls", "unprefixed_malloc_on_supported_platforms"] } +jemallocator = { package = "tikv-jemallocator", version = "0.6", features = [ + "disable_initial_exec_tls", + "unprefixed_malloc_on_supported_platforms", +] } [profile.release] lto = true opt-level = 3 codegen-units = 1 strip = true -panic = "abort" \ No newline at end of file +panic = "abort" diff --git a/examples/client.py b/examples/client.py index 7c234a2c..1e9b7b0b 100644 --- a/examples/client.py +++ b/examples/client.py @@ -6,6 +6,7 @@ async def main(): client = Client( impersonate=Impersonate.Firefox133, user_agent="rnet", + async_dns=True, ) resp = await client.get("https://httpbin.org/stream/20") print("Status Code: ", resp.status_code) diff --git a/src/client.rs b/src/client.rs index 990f61c1..f069ac6b 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,4 +1,5 @@ use crate::{ + dns, error::{ wrap_invali_header_name_error, wrap_invali_header_value_error, wrap_rquest_error, wrap_url_parse_error, @@ -12,6 +13,7 @@ use arc_swap::{ArcSwap, Guard}; use pyo3::prelude::*; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; use rquest::{ + dns::LookupIpStrategy, header::{HeaderMap, HeaderName, HeaderValue}, redirect::Policy, Url, @@ -466,7 +468,7 @@ impl Client { #[pyo3(signature = (**kwds))] fn new(mut kwds: Option) -> PyResult { let params = kwds.get_or_insert_default(); - let mut builder = rquest::Client::builder(); + let mut builder = rquest::Client::builder().no_hickory_dns(); // Impersonation options. if let Some(impersonate) = params.impersonate.take() { @@ -534,6 +536,12 @@ impl Client { // Cookie store options. apply_option!(apply_if_some, builder, params.cookie_store, cookie_store); + // Async resolver options. + if params.async_dns.unwrap_or(false) { + let hickory_dns_resolver = dns::get_or_try_init(LookupIpStrategy::Ipv4AndIpv6)?; + builder = builder.dns_resolver(hickory_dns_resolver); + } + // Timeout options. apply_option!( apply_transformed_option, diff --git a/src/dns.rs b/src/dns.rs new file mode 100644 index 00000000..5fd35f46 --- /dev/null +++ b/src/dns.rs @@ -0,0 +1,48 @@ +use crate::error::DNSResolverError; +use rquest::dns::{HickoryDnsResolver, LookupIpStrategy}; +use std::sync::{Arc, OnceLock}; + +/// Initializes and returns a DNS resolver with the specified strategy. +/// +/// This function initializes a global DNS resolver using the provided lookup IP strategy. +/// If the DNS resolver has already been initialized, it returns the existing instance. +/// +/// # Arguments +/// +/// * `strategy` - An optional `LookupIpStrategy` to use for the DNS resolver. +/// +/// # Returns +/// +/// A `Result` containing an `Arc` to the `HickoryDnsResolver` instance, or an error if initialization fails. +/// +/// # Errors +/// +/// This function returns an error if the DNS resolver fails to initialize. +/// +/// # Examples +/// +/// ```rust +/// use rnet::dns::get_or_try_init; +/// use rquest::dns::LookupIpStrategy; +/// +/// let resolver = get_or_try_init(LookupIpStrategy::default()).unwrap(); +/// ``` +pub fn get_or_try_init(strategy: S) -> crate::Result> +where + S: Into>, +{ + static DNS_RESOLVER: OnceLock, &'static str>> = OnceLock::new(); + + DNS_RESOLVER + .get_or_init(move || { + HickoryDnsResolver::new(strategy.into()) + .map(Arc::new) + .map_err(|err| { + eprintln!("failed to initialize the DNS resolver: {}", err); + "failed to initialize the DNS resolver" + }) + }) + .as_ref() + .map(Arc::clone) + .map_err(DNSResolverError::new_err) +} diff --git a/src/error.rs b/src/error.rs index 37534dc9..fbec9058 100644 --- a/src/error.rs +++ b/src/error.rs @@ -19,6 +19,7 @@ Potential solutions: "#; create_exception!(exceptions, BorrowingError, PyRuntimeError); +create_exception!(exceptions, DNSResolverError, PyRuntimeError); create_exception!(exceptions, BaseError, PyException); create_exception!(exceptions, BodyError, BaseError); diff --git a/src/lib.rs b/src/lib.rs index f7d10314..c0bd196b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ mod client; +mod dns; mod error; mod param; mod response; diff --git a/src/param/client.rs b/src/param/client.rs index f9aa75f5..249cbe5b 100644 --- a/src/param/client.rs +++ b/src/param/client.rs @@ -76,6 +76,10 @@ pub struct ClientParams { #[pyo3(get)] pub cookie_store: Option, + /// Whether to use async DNS resolver. + #[pyo3(get)] + pub async_dns: Option, + // ========= Timeout options ========= /// The timeout to use for the request. (in seconds) #[pyo3(get)] @@ -234,6 +238,7 @@ impl<'py> FromPyObject<'py> for ClientParams { extract_option!(ob, params, referer); extract_option!(ob, params, allow_redirects); extract_option!(ob, params, cookie_store); + extract_option!(ob, params, async_dns); extract_option!(ob, params, timeout); extract_option!(ob, params, connect_timeout);