idoh 0.1.11

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
103
104
105
106
107
108
109
110
111

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,
              });
            }
          }
        }
        Ok(Some(r))
      })
      .await
  }
}

#[cfg(feature = "cache")]
pub mod cache {
  use std::ops::Deref;

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

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

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

  pub struct Cache;

  /// Wrapper around RefVal that derefs to [Mx] instead of Vec<Mx>
  pub struct MxRef<'a> {
    inner: expire_cache::map::RefVal<'a, String, Vec<Mx>>,
  }

  impl<'a> Deref for MxRef<'a> {
    type Target = [Mx];
    fn deref(&self) -> &Self::Target {
      &self.inner
    }
  }

  impl super::MxLookup for Cache {
    type VecMx<'a> = MxRef<'a>;
    async fn mx<'a>(
      &'a self,
      domain: impl AsRef<str> + Send + 'a,
    ) -> Result<Option<Self::VecMx<'a>>> {
      let domain = domain.as_ref();
      let domain_key = domain.to_string();
      if let Some(result) = CACHE.get(&domain_key) {
        return Ok(Some(MxRef { inner: result }));
      }
      let result = match Resolve.mx(domain).await {
        Ok(r) => r,
        Err(err) => {
          log::warn!("{domain} MX: {err}");
          Default::default()
        }
      };
      if let Some(vec) = result {
        CACHE.insert(domain_key.clone(), vec);
      }
      Ok(CACHE.get(&domain_key).map(|inner| MxRef { inner }))
    }
  }
}