dns: Multiple fixes after testing (5ea0493b) · Commits · ankitgusai19 / onionmasq

crates/onion-tunnel/src/dns.rs

+30 −11

Original line number	Diff line number	Diff line
		use std::{
		collections::HashMap,
		net::Ipv4Addr,
		net::{Ipv4Addr, IpAddr},
		sync::{Arc, Mutex},
		};

		@@ -17,7 +17,10 @@ use crate::socket::UdpSocket;
		pub type LockedDnsCookies = Arc<Mutex<DnsCookies>>;

		pub struct DnsCookies {
		cookies: HashMap<IpAddress, String>,
		cookies: HashMap<IpAddr, String>,
		// Reverse cookie. This is so when we get a new hostname, we can quickly learn if we have it in
		// the cookie hashmap or not. XXX: This is not optimal.
		rev_cookies: HashMap<String, IpAddr>,
		next_v4_cookie: Ipv4Addr,
		}

		@@ -25,12 +28,28 @@ impl DnsCookies {
		fn new() -> Self {
		Self {
		cookies: HashMap::new(),
		next_v4_cookie: Ipv4Addr::new(1, 0, 0, 0),
		rev_cookies: HashMap::new(),
		next_v4_cookie: Ipv4Addr::new(1, 0, 0, 1),
		}
		}

		fn next_v4(&mut self, cookie: String) -> Ipv4Addr {
		self.cookies.insert(self.next_v4_cookie.into(), cookie);
		fn next_v4(&mut self, h: String) -> Ipv4Addr {
		let mut hostname = h.clone();
		if hostname.ends_with('.') {
		hostname.pop();
		}

		if let Some(ip) = self.rev_cookies.get(&hostname) {
		match ip {
		IpAddr::V4(v4) => return v4.clone(),
		_ => (),
		}
		}

		info!("Added hostname cookie: {} -> {}", self.next_v4_cookie, hostname);
		self.cookies.insert(self.next_v4_cookie.into(), hostname.clone());
		self.rev_cookies.insert(hostname, self.next_v4_cookie.into());

		let current = self.next_v4_cookie;

		// XXX: This is obviously problematic, it needs to be bounded.
		@@ -41,7 +60,7 @@ impl DnsCookies {
		}

		pub fn get(&self, addr: &IpAddress) -> Option<&String> {
		self.cookies.get(addr)
		self.cookies.get(&addr.clone().into())
		}
		}

		@@ -95,10 +114,10 @@ impl DnsManager {
		let flags = dns_message_parser::Flags {
		qr: true,
		opcode: dns_message_parser::Opcode::Query,
		aa: true,
		aa: false,
		tc: false,
		rd: true,
		ra: false,
		ra: true,
		ad: false,
		cd: false,
		rcode: dns_message_parser::RCode::NoError,
		@@ -106,10 +125,10 @@ impl DnsManager {
		dns_message_parser::Dns {
		id: query.id,
		flags,
		questions: Vec::new(),
		questions: query.questions.clone(),
		answers: vec![rr],
		authorities: Vec::new(),
		additionals: Vec::new(),
		authorities: query.authorities.clone(),
		additionals: query.additionals.clone(),
		}
		}

crates/onion-tunnel/src/proxy.rs

+3 −3

Original line number	Diff line number	Diff line
		@@ -41,9 +41,9 @@ impl ArtiProxy {
		tor_addr = (ip_addr, dest.1).into_tor_addr_dangerously();
		}

		info!("Connecting to: {:?}", dest);
		let mut arti_stream = self.arti.connect(tor_addr.unwrap()).await.unwrap();
		info!("Connected to: {:?}", dest);
		info!("Connecting to: {:?}", tor_addr.as_ref().unwrap());
		let mut arti_stream = self.arti.connect(tor_addr.as_ref().unwrap()).await.unwrap();
		info!("Connected to: {:?}", tor_addr.as_ref().unwrap());

		loop {
		let mut arti_buf = Vec::new();