dns: Wrap IP cookies a bit better (44307d4c) · Commits · ankitgusai19 / onionmasq

crates/onion-tunnel/src/dns.rs

+26 −4

Original line number	Diff line number	Diff line
		@@ -13,10 +13,32 @@ use smoltcp::{

		use crate::socket::UdpSocket;

		pub type LockedDnsCookies = Arc<Mutex<DnsCookies>>;

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

		impl DnsCookies {
		fn new() -> Self {
		Self {
		cookies: HashMap::new(),
		}
		}

		fn add(&mut self, addr: IpAddress, cookie: String) {
		self.cookies.insert(addr, cookie);
		}

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

		pub struct DnsManager {
		listener_udp_v4: UdpSocket,

		cookies: Arc<Mutex<HashMap<IpAddress, String>>>,
		cookies: LockedDnsCookies,
		}

		impl DnsManager {
		@@ -31,11 +53,11 @@ impl DnsManager {

		Self {
		listener_udp_v4,
		cookies: Arc::new(Mutex::new(HashMap::new())),
		cookies: Arc::new(Mutex::new(DnsCookies::new())),
		}
		}

		pub fn cookies(&self) -> Arc<Mutex<HashMap<IpAddress, String>>> {
		pub fn cookies(&self) -> Arc<Mutex<DnsCookies>> {
		self.cookies.clone()
		}

		@@ -78,7 +100,7 @@ impl DnsManager {
		additionals: Vec::new(),
		};
		let _ = self.listener_udp_v4.send(reply.encode().unwrap().as_ref(), endpoint);
		self.cookies.lock().unwrap().insert(cookie.into(), q.domain_name.to_string().clone());
		self.cookies.lock().unwrap().add(cookie.into(), q.domain_name.to_string().clone());
		}
		Err(e) => error!("Unable to parse DNS request: {}", e),
		}

crates/onion-tunnel/src/lib.rs

+19 −8

Original line number	Diff line number	Diff line
		@@ -6,7 +6,7 @@ mod socket;

		use arti_client::{TorClient, TorClientConfig};
		use device::VirtualDevice;
		use dns::DnsManager;
		use dns::{DnsManager, LockedDnsCookies};
		use log::{debug, info};
		use proxy::ArtiProxy;
		use smoltcp::{
		@@ -15,7 +15,7 @@ use smoltcp::{
		wire::{IpAddress, IpCidr, Ipv4Address},
		};
		use std::{
		collections::{BTreeMap, HashMap},
		collections::BTreeMap,
		sync::{Arc, Mutex},
		};
		use tor_rtcompat::tokio::TokioNativeTlsRuntime;
		@@ -27,6 +27,7 @@ pub struct OnionTunnel {
		iface_name: String,
		iface: Arc<Mutex<Interface<'static, VirtualDevice<TunTapInterface>>>>,
		arti: TorClient<TokioNativeTlsRuntime>,
		dns_cookies: Option<LockedDnsCookies>,
		}

		impl OnionTunnel {
		@@ -59,6 +60,7 @@ impl OnionTunnel {
		iface_fd,
		iface_name: iface_name.to_string(),
		iface,
		dns_cookies: None,
		}
		}

		@@ -71,18 +73,27 @@ impl OnionTunnel {
		}
		}

		fn proxy(&mut self, socket: TcpSocket, cookies: Arc<Mutex<HashMap<IpAddress, String>>>) {
		let mut proxy = ArtiProxy::new(socket, self.arti.clone(), cookies);
		fn proxy(&mut self, socket: TcpSocket) {
		let mut proxy = ArtiProxy::new(
		socket,
		self.arti.clone(),
		self.dns_cookies.as_ref().unwrap().clone(),
		);
		tokio::spawn(async move { proxy.start().await });
		}

		fn dns(&mut self) {
		// Spawn the DNS manager service.
		let mut dns_manager = DnsManager::new(self.iface.clone());
		self.dns_cookies = Some(dns_manager.cookies());
		tokio::spawn(async move { dns_manager.start().await });
		}

		pub async fn start(&mut self) {
		info!("Starting onion tunnel on TUN iface {}", self.iface_name);

		// Spawn the DNS manager service.
		let mut dns_manager = DnsManager::new(self.iface.clone());
		let cookies = dns_manager.cookies();
		tokio::spawn(async move { dns_manager.start().await });
		self.dns();

		loop {
		let timestamp = smoltcp::time::Instant::now();
		@@ -96,7 +107,7 @@ impl OnionTunnel {
		while let Some(packet) = packets.pop_front() {
		if let Some(tcp_socket) = Parser::parse(packet.clone()).take() {
		let socket = TcpSocket::new(self.iface.clone(), tcp_socket);
		self.proxy(socket, cookies.clone());
		self.proxy(socket);
		}
		self.iface.lock().unwrap().device_mut().put_packet(packet);
		}

crates/onion-tunnel/src/proxy.rs

+4 −8

Original line number	Diff line number	Diff line
		use std::collections::HashMap;
		use std::net::IpAddr;
		use std::sync::Arc;
		use std::sync::Mutex;

		use arti_client::DangerouslyIntoTorAddr;
		use arti_client::TorAddr;
		use arti_client::TorClient;
		use log::info;
		use smoltcp::wire::IpAddress;
		use std::net::IpAddr;
		use tokio::io::{AsyncReadExt, AsyncWriteExt};
		use tor_rtcompat::tokio::TokioNativeTlsRuntime;

		use crate::dns::LockedDnsCookies;
		use crate::socket::TcpSocket;

		pub struct ArtiProxy {
		socket: TcpSocket,
		arti: TorClient<TokioNativeTlsRuntime>,
		cookies: Arc<Mutex<HashMap<IpAddress, String>>>,
		cookies: LockedDnsCookies,
		}

		impl ArtiProxy {
		pub fn new(
		socket: TcpSocket,
		arti: TorClient<TokioNativeTlsRuntime>,
		cookies: Arc<Mutex<HashMap<IpAddress, String>>>,
		cookies: LockedDnsCookies,
		) -> Self {
		Self {
		socket,

crates/onion-tunnel/src/socket.rs

+0 −1

Original line number	Diff line number	Diff line
		use std::{
		net::IpAddr,
		pin::Pin,
		sync::{Arc, Mutex},
		task::{Context, Poll},