dns: Refactor a bit to be cleaner (ad05dc95) · Commits · ankitgusai19 / onionmasq

crates/onion-tunnel/src/dns.rs

+47 −32

Original line number	Diff line number	Diff line
		@@ -17,17 +17,26 @@ pub type LockedDnsCookies = Arc<Mutex<DnsCookies>>;

		pub struct DnsCookies {
		cookies: HashMap<IpAddress, String>,
		next_v4_cookie: Ipv4Addr,
		}

		impl DnsCookies {
		fn new() -> Self {
		Self {
		cookies: HashMap::new(),
		next_v4_cookie: Ipv4Addr::new(1, 0, 0, 0),
		}
		}

		fn add(&mut self, addr: IpAddress, cookie: String) {
		self.cookies.insert(addr, cookie);
		fn next_v4(&mut self, cookie: String) -> Ipv4Addr {
		self.cookies.insert(self.next_v4_cookie.into(), cookie);
		let current = self.next_v4_cookie;

		// XXX: This is obviously problematic, it needs to be bounded.
		let u_addr: u32 = self.next_v4_cookie.into();
		self.next_v4_cookie = (u_addr + 1).into();

		current
		}

		pub fn get(&self, addr: &IpAddress) -> Option<&String> {
		@@ -37,7 +46,6 @@ impl DnsCookies {

		pub struct DnsManager {
		listener_udp_v4: UdpSocket,

		cookies: LockedDnsCookies,
		}

		@@ -61,19 +69,13 @@ impl DnsManager {
		self.cookies.clone()
		}

		pub async fn start(&mut self) {
		info!("Starting DNS manager");

		loop {
		tokio::select! {
		r = self.listener_udp_v4.next() => match r {
		Some((payload, endpoint)) => {
		let ret = dns_message_parser::Dns::decode(payload);
		match ret {
		Ok(p) => {
		let cookie = Ipv4Addr::new(10, 42, 42, 254);
		let q = p.questions.get(0).unwrap();
		info!("DNS Question: {} from {}", q, endpoint);
		fn build_answer(&mut self, query: &dns_message_parser::Dns) -> dns_message_parser::Dns {
		let q = query.questions.get(0).unwrap();
		let cookie = self
		.cookies
		.lock()
		.unwrap()
		.next_v4(q.domain_name.to_string().clone());
		let a = dns_message_parser::rr::A {
		domain_name: q.domain_name.clone(),
		ttl: 10,
		@@ -91,16 +93,29 @@ impl DnsManager {
		cd: false,
		rcode: dns_message_parser::RCode::NoError,
		};
		let reply = dns_message_parser::Dns {
		id: p.id,
		dns_message_parser::Dns {
		id: query.id,
		flags,
		questions: Vec::new(),
		answers: vec![rr],
		authorities: Vec::new(),
		additionals: Vec::new(),
		};
		}
		}

		pub async fn start(&mut self) {
		info!("Starting DNS manager");

		loop {
		tokio::select! {
		r = self.listener_udp_v4.next() => match r {
		Some((payload, endpoint)) => {
		let ret = dns_message_parser::Dns::decode(payload);
		match ret {
		Ok(p) => {
		info!("DNS Question: {} from {}", p, endpoint);
		let reply = self.build_answer(&p);
		let _ = self.listener_udp_v4.send(reply.encode().unwrap().as_ref(), endpoint);
		self.cookies.lock().unwrap().add(cookie.into(), q.domain_name.to_string().clone());
		}
		Err(e) => error!("Unable to parse DNS request: {}", e),
		}