forked from cloudflare/pingora
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dns.rs
144 lines (136 loc) · 4.53 KB
/
dns.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
use crate::discovery::ServiceDiscovery;
use crate::Backend;
use async_trait::async_trait;
use hickory_client::client::{Client, SyncClient};
use hickory_client::op::DnsResponse;
use hickory_client::rr::{DNSClass, Name, RData, Record, RecordType};
use hickory_client::udp::UdpClientConnection;
use pingora_error::Error;
use pingora_error::ErrorType::Custom;
use rand::Rng;
use std::collections::{BTreeSet, HashMap};
use std::env;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str::FromStr;
pub struct DNS {
domain: Box<str>,
resolver: Vec<SocketAddr>,
port: u32,
}
static DNS_ADDR_ENV: &str = "DNS_ADDR";
static GOOGLE_DNS: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53);
impl DNS {
///
/// # Example
/// Dns::new("www.example.com:8088")
/// if this domain port haven't set ,default 80
pub fn new(service: &str) -> pingora_error::Result<Self, Error> {
let mut resolver = vec![];
if service.is_empty() {
return Err(*Error::new(Custom("the service is empty is illegal")));
}
// from ENV get dns addr
if let Ok(result) = env::var(DNS_ADDR_ENV) {
result
.split(",")
.into_iter()
.for_each(|item| resolver.push(item.parse().unwrap()));
} else {
// default
resolver.push(GOOGLE_DNS);
}
let addr: Vec<&str> = service.split(":").map(|field| field.trim()).collect();
if addr.len() == 1 {
return Ok(DNS {
domain: service.into(),
resolver,
port: 80,
});
}
let port = addr.get(1).unwrap().parse().unwrap();
return Ok(DNS {
domain: service.into(),
resolver,
port,
});
}
///
/// new DNS
/// # Example
/// let _ = DNS::new("www.example.com",vec!["8.8.8.8:53"])
///
pub fn new_with_resolve(
service: &str,
resolver: Vec<String>,
) -> pingora_error::Result<Self, Error> {
if service.is_empty() {
return Err(*Error::new(Custom("the service is empty is illegal")));
}
if resolver.is_empty() {
return Err(*Error::new(Custom("the resolver is empty is illegal")));
}
let resolver = resolver.iter().map(|addr| addr.parse().unwrap()).collect();
let addr: Vec<&str> = service.split(":").map(|field| field.trim()).collect();
if addr.len() == 1 {
return Ok(DNS {
domain: service.into(),
resolver,
port: 80,
});
}
let port = addr.get(1).unwrap().parse().unwrap();
return Ok(DNS {
domain: service.into(),
resolver,
port,
});
}
fn insert_item(
&self,
remote_addr: &str,
health: &mut HashMap<u64, bool>,
tree: &mut BTreeSet<Backend>,
) {
let addr = format!("{}:{}", remote_addr, self.port);
let value = Backend::new(addr.as_str()).unwrap();
health.insert(value.hash_key(), true);
tree.insert(value);
}
}
///
/// implemented DNS discovery
#[async_trait]
impl ServiceDiscovery for DNS {
async fn discover(&self) -> pingora_error::Result<(BTreeSet<Backend>, HashMap<u64, bool>)> {
// no readiness
let mut health: HashMap<u64, bool> = HashMap::new();
let mut tree: BTreeSet<Backend> = BTreeSet::new();
// random dns address
let mut rng = rand::thread_rng();
let index;
if self.resolver.len() > 1 {
index = rng.gen_range(0..self.resolver.len());
} else {
index = 0
}
// random request dns addr
let address = self.resolver[index];
let conn = UdpClientConnection::new(address).unwrap();
let client = SyncClient::new(conn);
let name = Name::from_str(self.domain.as_ref()).unwrap();
let response: DnsResponse = client.query(&name, DNSClass::IN, RecordType::A).unwrap();
let answers: &[Record] = response.answers();
for x in answers {
match x.data() {
Some(RData::A(ref remote_addr)) => {
self.insert_item(remote_addr.to_string().as_str(), &mut health, &mut tree)
}
Some(RData::AAAA(ref remote_addr)) => {
self.insert_item(remote_addr.to_string().as_str(), &mut health, &mut tree)
}
_ => {}
}
}
Ok((tree, health))
}
}