any_dns/
custom_handler.rs

1#![allow(unused)]
2
3use async_trait::async_trait;
4use dyn_clone::DynClone;
5use std::fmt::Debug;
6
7use crate::dns_socket::DnsSocket;
8
9#[derive(thiserror::Error, Debug)]
10pub enum CustomHandlerError {
11    #[error(transparent)]
12    IO(#[from] crate::dns_socket::RequestError),
13
14    #[error("Query is not processed by handler. Fallback to ICANN.")]
15    Unhandled,
16}
17
18/**
19 * Trait to implement to make AnyDns use a custom handler.
20 * Important: Handler must be clonable so it can be used by multiple threads.
21 */
22#[async_trait]
23pub trait CustomHandler: DynClone + Send + Sync {
24    async fn lookup(
25        &mut self,
26        query: &Vec<u8>,
27        socket: DnsSocket,
28    ) -> Result<Vec<u8>, CustomHandlerError>;
29}
30
31/**
32 * Clonable handler holder
33 */
34pub struct HandlerHolder {
35    pub func: Box<dyn CustomHandler>,
36}
37
38impl Clone for HandlerHolder {
39    fn clone(&self) -> Self {
40        Self {
41            func: dyn_clone::clone_box(&*self.func),
42        }
43    }
44}
45
46impl Debug for HandlerHolder {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("HandlerHolder")
49            .field("func", &"HandlerHolder")
50            .finish()
51    }
52}
53
54impl HandlerHolder {
55    /**
56     * Bootstrap a holder from a struct that implements the CustomHandler.
57     */
58    pub fn new(f: impl CustomHandler + 'static) -> Self {
59        HandlerHolder { func: Box::new(f) }
60    }
61
62    pub async fn call(
63        &mut self,
64        query: &Vec<u8>,
65        socket: DnsSocket,
66    ) -> Result<Vec<u8>, CustomHandlerError> {
67        self.func.lookup(query, socket).await
68    }
69}
70
71#[derive(Clone)]
72pub struct EmptyHandler {}
73
74impl EmptyHandler {
75    pub fn new() -> Self {
76        EmptyHandler {}
77    }
78}
79
80#[async_trait]
81impl CustomHandler for EmptyHandler {
82    async fn lookup(
83        &mut self,
84        _query: &Vec<u8>,
85        _socket: DnsSocket,
86    ) -> Result<Vec<u8>, CustomHandlerError> {
87        Err(CustomHandlerError::Unhandled)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use crate::dns_socket::DnsSocket;
94    use async_trait::async_trait;
95    use std::net::SocketAddr;
96
97    use super::{CustomHandler, CustomHandlerError, HandlerHolder};
98
99    struct ClonableStruct {
100        value: String,
101    }
102
103    impl Clone for ClonableStruct {
104        fn clone(&self) -> Self {
105            Self {
106                value: format!("{} cloned", self.value.clone()),
107            }
108        }
109    }
110
111    #[derive(Clone)]
112    pub struct TestHandler {
113        value: ClonableStruct,
114    }
115
116    impl TestHandler {
117        pub fn new(value: &str) -> Self {
118            TestHandler {
119                value: ClonableStruct {
120                    value: value.to_string(),
121                },
122            }
123        }
124    }
125    #[async_trait]
126    impl CustomHandler for TestHandler {
127        async fn lookup(
128            &mut self,
129            _query: &Vec<u8>,
130            _socket: DnsSocket,
131        ) -> Result<Vec<u8>, CustomHandlerError> {
132            println!("value {}", self.value.value);
133            Err(CustomHandlerError::Unhandled)
134        }
135    }
136
137    #[tokio::test]
138    async fn run_processor() {
139        let test1 = TestHandler::new("test1");
140        let holder1 = HandlerHolder::new(test1);
141        let mut cloned = holder1.clone();
142        let icann_fallback: SocketAddr = "8.8.8.8:53".parse().unwrap();
143
144        let socket = DnsSocket::new(
145            "0.0.0.0:18293".parse().unwrap(),
146            icann_fallback,
147            holder1.clone()
148        )
149        .await
150        .unwrap();
151        let result = cloned.call(&vec![], socket).await;
152        assert!(result.is_err());
153    }
154}