idoh 0.1.13

Fast and secure DNS over HTTPS resolution / 极速安全的 DoH 解析库
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

use aok::Result;

use crate::{Resolver, record_type};

/// MX record structure
#[derive(Debug, Clone)]
pub struct Mx {
  /// Priority of the mail server (lower values have higher priority)
  pub priority: u16,
  /// Mail server hostname
  pub server: String,
  /// Time to live in seconds
  pub ttl: u64,
}

pub trait MxLookup {
  type VecMx<'a>: std::ops::Deref<Target = [Mx]> + 'a
  where
    Self: 'a;
  fn mx<'a>(
    &'a self,
    domain: impl AsRef<str> + Send + 'a,
  ) -> impl std::future::Future<Output = Result<Option<Self::VecMx<'a>>>> + Send + 'a;
}

impl<T: Resolver + Sync> MxLookup for T {
  type VecMx<'a>
    = Vec<Mx>
  where
    Self: 'a;
  async fn mx<'a>(
    &'a self,
    domain: impl AsRef<str> + Send + 'a,
  ) -> Result<Option<Self::VecMx<'a>>> {
    self
      .resolve(domain, "MX", |li| {
        let mut r = Vec::with_capacity(li.len());
        for i in li {
          if i.r#type == record_type::MX {
            // MX data format: "priority server"
            // Example: "10 mail.example.com."
            if let Some((priority_str, server)) = i.data.split_once(' ')
              && let Ok(priority) = priority_str.parse::<u16>()
            {
              r.push(Mx {
                priority,
                server: server.trim_end_matches('.').to_string(),
                ttl: i.ttl,
              });
            }
          }
        }
        r.sort_unstable_by_key(|i| i.priority);
        Ok(Some(r))
      })
      .await
  }
}

#[cfg(feature = "cache")]
pub mod cache {

  use aok::Result;
  use dashmap::DashMap;
  use expire_cache::Expire;

  use super::Mx;
  use crate::Resolve;

  pub struct MxRef<'a>(expire_cache::map::RefVal<'a, String, Option<Vec<Mx>>>);

  impl<'a> std::ops::Deref for MxRef<'a> {
    type Target = [Mx];
    fn deref(&self) -> &Self::Target {
      match &*self.0 {
        None => &[],
        Some(li) => li,
      }
    }
  }

  #[static_init::dynamic]
  pub static CACHE: Expire<DashMap<String, Option<Vec<Mx>>>> = Expire::new(600);

  pub struct Cache;

  impl super::MxLookup for Cache {
    type VecMx<'a> = MxRef<'a>;
    async fn mx<'a>(&'a self, domain: impl AsRef<str> + Send + 'a) -> Result<Option<MxRef<'a>>> {
      let domain = domain.as_ref();
      let r = CACHE
        .get_or_init_async(domain.to_owned(), move |domain| {
          Box::pin(Resolve.mx(domain.clone()))
        })
        .await?;
      if r.is_some() {
        return Ok(Some(MxRef(r)));
      }
      Ok(None)
    }
  }
}