Skip to main content

common/
dns.rs

1use std::error::Error;
2
3use rand::distributions::Alphanumeric;
4use rand::{thread_rng, Rng};
5use std::convert::TryInto;
6
7pub trait Encoder: Send {
8    fn calculate_max_decoded_size(&self, max_encoded_size: usize) -> usize;
9    fn encode(&self, data: &[u8]) -> Result<String, Box<dyn Error>>;
10}
11
12pub trait Decoder: Send {
13    fn decode(&self, data: &str) -> Result<Vec<u8>, Box<dyn Error>>;
14}
15
16pub struct HexEncoder {}
17
18impl Encoder for HexEncoder {
19    fn calculate_max_decoded_size(&self, max_encoded_size: usize) -> usize {
20        max_encoded_size / 2
21    }
22    fn encode(&self, data: &[u8]) -> Result<String, Box<dyn Error>> {
23        Ok(hex::encode(data))
24    }
25}
26
27pub struct HexDecoder {}
28
29impl Decoder for HexDecoder {
30    fn decode(&self, data: &str) -> Result<Vec<u8>, Box<dyn Error>> {
31        match hex::decode(data) {
32            Ok(v) => Ok(v),
33            Err(e) => Err(e.into()),
34        }
35    }
36}
37
38pub const CLIENT_ID_SIZE_IN_BYTES: usize = 4;
39pub type ClientId = [u8; CLIENT_ID_SIZE_IN_BYTES];
40
41pub fn new_client_id() -> ClientId {
42    let values: Vec<u8> = thread_rng()
43        .sample_iter(&Alphanumeric)
44        .take(CLIENT_ID_SIZE_IN_BYTES)
45        .collect();
46    values.try_into().unwrap()
47}
48
49pub struct ClientIdSuffixEncoder<E: Encoder> {
50    encoder: E,
51}
52
53impl<E: Encoder> ClientIdSuffixEncoder<E> {
54    pub fn new(encoder: E) -> Self {
55        Self { encoder }
56    }
57}
58
59impl<E: Encoder> Encoder for ClientIdSuffixEncoder<E> {
60    fn calculate_max_decoded_size(&self, max_encoded_size: usize) -> usize {
61        self.encoder.calculate_max_decoded_size(max_encoded_size) - CLIENT_ID_SIZE_IN_BYTES
62    }
63
64    fn encode(&self, data: &[u8]) -> Result<String, Box<dyn Error>> {
65        if data.len() < CLIENT_ID_SIZE_IN_BYTES {
66            return Err(String::from("not enough data to encode").into());
67        }
68
69        let (data, client_id) = data.split_at(data.len() - CLIENT_ID_SIZE_IN_BYTES);
70        let mut res = self.encoder.encode(data)?;
71        let encoded_client_id = std::str::from_utf8(client_id).unwrap();
72        res.push_str(encoded_client_id);
73        Ok(res)
74    }
75}
76
77pub struct ClientIdSuffixDecoder<D: Decoder> {
78    decoder: D,
79}
80
81impl<D: Decoder> ClientIdSuffixDecoder<D> {
82    pub fn new(decoder: D) -> Self {
83        Self { decoder }
84    }
85}
86
87impl<D: Decoder> Decoder for ClientIdSuffixDecoder<D> {
88    fn decode(&self, data: &str) -> Result<Vec<u8>, Box<dyn Error>> {
89        if data.len() < CLIENT_ID_SIZE_IN_BYTES {
90            return Err(String::from("not enough data to decode").into());
91        }
92
93        let (data, client_id) = data.split_at(data.len() - CLIENT_ID_SIZE_IN_BYTES);
94        let mut res = self.decoder.decode(data)?;
95        let decoded_client_id = client_id.as_bytes();
96        res.extend_from_slice(decoded_client_id);
97        Ok(res)
98    }
99}
100
101pub struct AppendSuffixEncoder<E: Encoder> {
102    encoder: E,
103    suffix: String,
104}
105
106impl<E: Encoder> AppendSuffixEncoder<E> {
107    pub fn new(encoder: E, suffix: String) -> Self {
108        Self { encoder, suffix }
109    }
110}
111
112impl<E: Encoder> Encoder for AppendSuffixEncoder<E> {
113    fn calculate_max_decoded_size(&self, max_encoded_size: usize) -> usize {
114        self.encoder.calculate_max_decoded_size(max_encoded_size) - self.suffix.len()
115    }
116
117    fn encode(&self, data: &[u8]) -> Result<String, Box<dyn Error>> {
118        let mut res = self.encoder.encode(data)?;
119        res.push_str(self.suffix.as_str());
120        Ok(res)
121    }
122}
123
124pub struct AppendSuffixDecoder<D: Decoder> {
125    decoder: D,
126    suffix: String,
127}
128
129impl<D: Decoder> AppendSuffixDecoder<D> {
130    pub fn new(decoder: D, suffix: String) -> Self {
131        Self { decoder, suffix }
132    }
133}
134
135impl<D: Decoder> Decoder for AppendSuffixDecoder<D> {
136    fn decode(&self, data: &str) -> Result<Vec<u8>, Box<dyn Error>> {
137        let data_without_suffix = data
138            .strip_suffix(self.suffix.as_str())
139            .ok_or_else(|| String::from("missing suffix"))?;
140        self.decoder.decode(data_without_suffix)
141    }
142}
143
144#[cfg(test)]
145mod tests {
146
147    use mockall::mock;
148
149    use super::*;
150
151    mock! {
152        Encoder{}
153        impl Encoder for Encoder {
154            fn calculate_max_decoded_size(&self, max_encoded_size: usize) -> usize;
155            fn encode(&self, data: &[u8]) -> Result<String, Box<dyn Error>>;
156        }
157    }
158
159    mock! {
160        Decoder{}
161        impl Decoder for Decoder {
162            fn decode(&self, data: &str) -> Result<Vec<u8>, Box<dyn Error>>;
163        }
164    }
165
166    #[test]
167    fn client_id_suffix_encoder_not_enough_data() -> Result<(), Box<dyn Error>> {
168        let encoder_mock = MockEncoder::new();
169        let encoder = ClientIdSuffixEncoder::new(encoder_mock);
170        let res = encoder.encode(b"123");
171        assert!(res.is_err());
172        Ok(())
173    }
174
175    #[test]
176    fn client_id_suffix_encoder_internal_encoder_failed() -> Result<(), Box<dyn Error>> {
177        let mut encoder_mock = MockEncoder::new();
178        encoder_mock
179            .expect_encode()
180            .returning(|_| Err(String::from("bla").into()));
181        let encoder = ClientIdSuffixEncoder::new(encoder_mock);
182        let res = encoder.encode(b"bla1234");
183        assert!(res.is_err());
184        Ok(())
185    }
186
187    #[test]
188    fn client_id_suffix_encoder_success() -> Result<(), Box<dyn Error>> {
189        let mut encoder_mock = MockEncoder::new();
190        encoder_mock
191            .expect_encode()
192            .returning(|_| Ok(String::from("encoded")));
193        let encoder = ClientIdSuffixEncoder::new(encoder_mock);
194        let res = encoder.encode(b"bla1234")?;
195        assert_eq!("encoded1234", res);
196        Ok(())
197    }
198
199    #[test]
200    fn client_id_suffix_decoder_not_enough_data() -> Result<(), Box<dyn Error>> {
201        let decoder_mock = MockDecoder::new();
202        let decoder = ClientIdSuffixDecoder::new(decoder_mock);
203        let res = decoder.decode("123");
204        assert!(res.is_err());
205        Ok(())
206    }
207
208    #[test]
209    fn client_id_suffix_decoder_internal_decoder_failed() -> Result<(), Box<dyn Error>> {
210        let mut decoder_mock = MockDecoder::new();
211        decoder_mock
212            .expect_decode()
213            .returning(|_| Err(String::from("bla").into()));
214        let decoder = ClientIdSuffixDecoder::new(decoder_mock);
215        let res = decoder.decode("bla1234");
216        assert!(res.is_err());
217        Ok(())
218    }
219
220    #[test]
221    fn client_id_suffix_decoder_success() -> Result<(), Box<dyn Error>> {
222        let mut decoder_mock = MockDecoder::new();
223        decoder_mock
224            .expect_decode()
225            .returning(|_| Ok(String::from("decoded").into_bytes()));
226        let decoder = ClientIdSuffixDecoder::new(decoder_mock);
227        let res = decoder.decode("bla1234")?;
228        assert_eq!(b"decoded1234", res.as_slice());
229        Ok(())
230    }
231
232    #[test]
233    fn append_suffix_encoder_internal_encoder_failed() -> Result<(), Box<dyn Error>> {
234        let suffix = String::from("xyz");
235        let mut encoder_mock = MockEncoder::new();
236        encoder_mock
237            .expect_encode()
238            .returning(|_| Err(String::from("bla").into()));
239        let encoder = AppendSuffixEncoder::new(encoder_mock, suffix);
240        let res = encoder.encode(b"bla");
241        assert!(res.is_err());
242        Ok(())
243    }
244
245    #[test]
246    fn append_suffix_encoder_success() -> Result<(), Box<dyn Error>> {
247        let suffix = String::from("xyz");
248        let mut encoder_mock = MockEncoder::new();
249        encoder_mock
250            .expect_encode()
251            .returning(|_| Ok(String::from("encoded")));
252        let encoder = AppendSuffixEncoder::new(encoder_mock, suffix);
253        let res = encoder.encode(b"bla")?;
254        assert_eq!("encodedxyz", res);
255        Ok(())
256    }
257
258    #[test]
259    fn append_suffix_decoder_missing_suffix() -> Result<(), Box<dyn Error>> {
260        let suffix = String::from("xyz");
261        let decoder_mock = MockDecoder::new();
262        let decoder = AppendSuffixDecoder::new(decoder_mock, suffix);
263        let res = decoder.decode("123");
264        assert!(res.is_err());
265        Ok(())
266    }
267
268    #[test]
269    fn append_suffix_decoder_internal_decoder_failed() -> Result<(), Box<dyn Error>> {
270        let suffix = String::from("xyz");
271        let mut decoder_mock = MockDecoder::new();
272        decoder_mock
273            .expect_decode()
274            .returning(|_| Err(String::from("bla").into()));
275        let decoder = AppendSuffixDecoder::new(decoder_mock, suffix);
276        let res = decoder.decode("blaxyz");
277        assert!(res.is_err());
278        Ok(())
279    }
280
281    #[test]
282    fn append_suffix_decoder_success() -> Result<(), Box<dyn Error>> {
283        let suffix = String::from("xyz");
284        let mut decoder_mock = MockDecoder::new();
285        decoder_mock
286            .expect_decode()
287            .returning(|_| Ok(String::from("decoded").into_bytes()));
288        let decoder = AppendSuffixDecoder::new(decoder_mock, suffix);
289        let res = decoder.decode("blaxyz")?;
290        assert_eq!(b"decoded", res.as_slice());
291        Ok(())
292    }
293}