rewrite txt_lookup and srv_lookup

rewrite

rewrite

lookup

l
This commit is contained in:
Luna Yao
2026-04-19 00:34:41 +02:00
parent 5a516a195c
commit 7ae7725fbf
3 changed files with 74 additions and 66 deletions
+9 -27
View File
@@ -1,22 +1,17 @@
use std::{net::SocketAddr, sync::Arc};
use super::{create_connector_by_url, http_connector::TunnelWithInfo};
use crate::utils::dns::{srv_lookup, txt_resolve};
use crate::{
common::{
error::Error,
global_ctx::ArcGlobalCtx,
log,
},
common::{error::Error, global_ctx::ArcGlobalCtx, log},
proto::common::TunnelInfo,
tunnel::{IpScheme, IpVersion, Tunnel, TunnelConnector, TunnelError, TunnelScheme},
};
use anyhow::Context;
use dashmap::DashSet;
use hickory_proto::rr::RData;
use hickory_resolver::proto::rr::rdata::SRV;
use rand::{seq::SliceRandom, Rng as _};
use rand::{Rng as _, seq::SliceRandom};
use strum::VariantArray;
use crate::utils::dns::{resolve_txt_record, RESOLVER};
fn weighted_choice<T>(options: &[(T, u64)]) -> Option<&T> {
let total_weight = options.iter().map(|(_, weight)| *weight).sum();
@@ -59,14 +54,13 @@ impl DnsTunnelConnector {
&self,
domain_name: &str,
) -> Result<Box<dyn TunnelConnector>, Error> {
let txt_data = resolve_txt_record(domain_name)
let txt_data = txt_resolve(domain_name)
.await
.with_context(|| format!("resolve txt record failed, domain_name: {}", domain_name))?;
let candidate_urls = txt_data
.split(" ")
.map(|s| s.to_string())
.filter_map(|s| url::Url::parse(s.as_str()).ok())
.iter()
.filter_map(|s| url::Url::parse(s).ok())
.collect::<Vec<_>>();
// shuffle candidate_urls and get the first one
@@ -74,7 +68,7 @@ impl DnsTunnelConnector {
.choose(&mut rand::thread_rng())
.with_context(|| {
format!(
"no valid url found, txt_data: {}, expecting an url list splitted by space",
"no valid url found, txt_data: {:?}, expecting an url list split by space",
txt_data
)
})?;
@@ -84,7 +78,7 @@ impl DnsTunnelConnector {
Ok(connector)
}
fn handle_one_srv_record(record: &SRV, protocol: IpScheme) -> Result<(url::Url, u64), Error> {
fn handle_one_srv_record(record: SRV, protocol: IpScheme) -> Result<(url::Url, u64), Error> {
// port must be non-zero
if record.port == 0 {
return Err(anyhow::anyhow!("port must be non-zero").into());
@@ -120,21 +114,9 @@ impl DnsTunnelConnector {
let srv_lookup_tasks = srv_domains
.iter()
.map(|(protocol, srv_domain)| {
let resolver = RESOLVER.clone();
let responses = responses.clone();
async move {
let response = resolver.srv_lookup(srv_domain).await.with_context(|| {
format!("srv_lookup failed, srv_domain: {}", srv_domain)
})?;
tracing::info!(?response, ?srv_domain, "srv_lookup response");
for record in response
.answers()
.iter()
.filter_map(|record| match record.data {
RData::SRV(ref srv) => Some(srv),
_ => None,
})
{
for record in srv_lookup(srv_domain).await? {
let parsed_record = Self::handle_one_srv_record(record, **protocol);
tracing::info!(?parsed_record, ?srv_domain, "parsed_record");
if let Err(e) = &parsed_record {