https_dns/
upstream.rs

1use crate::bootstrap::BootstrapClient;
2use crate::cache::Cache;
3use crate::error::UpstreamError::{self, Build, Resolve};
4use reqwest::{
5    header::{HeaderMap, HeaderValue, CONTENT_TYPE},
6    Client,
7};
8use std::{net::IpAddr, time::Duration};
9use tracing::info;
10use trust_dns_proto::op::message::Message;
11
12#[derive(Clone, Debug)]
13pub struct HttpsClient {
14    host: String,
15    port: u16,
16    https_client: Client,
17    cache: Cache,
18}
19
20impl HttpsClient {
21    pub async fn new(host: String, port: u16) -> Result<Self, UpstreamError> {
22        let mut headers = HeaderMap::new();
23        headers.insert(
24            CONTENT_TYPE,
25            HeaderValue::from_str("application/dns-message").unwrap(),
26        );
27
28        let mut client_builder = Client::builder()
29            .default_headers(headers)
30            .https_only(true)
31            .gzip(true)
32            .brotli(true)
33            .timeout(Duration::from_secs(10));
34
35        if host.parse::<IpAddr>().is_err() {
36            let bootstrap_client = match BootstrapClient::new() {
37                Ok(bootstrap_client) => bootstrap_client,
38                Err(error) => return Err(error),
39            };
40            let ip_addr = match bootstrap_client.bootstrap(&host).await {
41                Ok(ip_addr) => ip_addr,
42                Err(error) => return Err(error),
43            };
44            client_builder = client_builder.resolve(&host, ip_addr);
45        }
46
47        let https_client = match client_builder.build() {
48            Ok(https_client) => https_client,
49            Err(_) => return Err(Build),
50        };
51        info!("connected to https://{}:{}", host, port);
52
53        Ok(HttpsClient {
54            host,
55            port,
56            https_client,
57            cache: Cache::new(),
58        })
59    }
60
61    pub async fn process(&mut self, request_message: Message) -> Result<Message, UpstreamError> {
62        if let Some(response_message) = self.cache.get(&request_message) {
63            return Ok(response_message);
64        }
65
66        let raw_request_message = match request_message.to_vec() {
67            Ok(raw_request_message) => raw_request_message,
68            Err(_) => return Err(Resolve),
69        };
70
71        let url = format!("https://{}:{}/dns-query", self.host, self.port);
72        let request = self.https_client.post(url).body(raw_request_message);
73        let response = match request.send().await {
74            Ok(response) => response,
75            Err(_) => return Err(Resolve),
76        };
77
78        let raw_response_message = match response.bytes().await {
79            Ok(response_bytes) => response_bytes,
80            Err(_) => return Err(Resolve),
81        };
82
83        let message = match Message::from_vec(&raw_response_message) {
84            Ok(message) => message,
85            Err(_) => return Err(Resolve),
86        };
87
88        self.cache.put(message.clone());
89        Ok(message)
90    }
91}