diff --git a/src/error.rs b/src/error.rs index 1a57783..b0c916d 100644 --- a/src/error.rs +++ b/src/error.rs @@ -33,6 +33,12 @@ pub enum Error { #[error("std::str::Utf8Error {0:?}")] Utf8(#[from] std::str::Utf8Error), + #[error("TryFromSliceError {0:?}")] + TryFromSlice(#[from] std::array::TryFromSliceError), + + #[error("ProtoError {0:?}")] + ProtoError(#[from] trust_dns_proto::error::ProtoError), + #[cfg(target_os = "android")] #[error("jni::errors::Error {0:?}")] Jni(#[from] jni::errors::Error), diff --git a/src/lib.rs b/src/lib.rs index 84c3d11..9588a4f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,14 @@ -use crate::{error::Error, http::HttpManager, socks::SocksProxyManager, tun2proxy::TunToProxy}; +use crate::{ + error::Error, + http::HttpManager, + socks::SocksProxyManager, + tun2proxy::{ConnectionManager, TunToProxy}, +}; use socks5_impl::protocol::{UserKey, Version}; use std::{ net::{SocketAddr, ToSocketAddrs}, rc::Rc, }; -use tun2proxy::ConnectionManager; mod android; mod dns; diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 29675e2..8e99aff 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -508,7 +508,8 @@ impl<'a> TunToProxy<'a> { let port = connection_info.dst.port(); if let (Some(virtual_dns), true) = (&mut self.options.virtual_dns, port == 53) { let payload = &frame[payload_offset..payload_offset + payload_size]; - if let Some(response) = virtual_dns.receive_query(payload) { + let response = virtual_dns.receive_query(payload)?; + { let rx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]); let tx_buffer = diff --git a/src/virtdns.rs b/src/virtdns.rs index 0d13aa6..6f8949a 100644 --- a/src/virtdns.rs +++ b/src/virtdns.rs @@ -1,3 +1,4 @@ +use crate::error::Result; use hashlink::{linked_hash_map::RawEntryMut, LruCache}; use smoltcp::wire::Ipv4Cidr; use std::{ @@ -10,19 +11,6 @@ use std::{ const MAPPING_TIMEOUT: u64 = 60; // Mapping timeout in seconds -#[derive(Eq, PartialEq, Debug)] -#[allow(dead_code, clippy::upper_case_acronyms)] -enum DnsRecordType { - A = 1, - AAAA = 28, -} - -#[derive(Eq, PartialEq, Debug)] -#[allow(dead_code)] -enum DnsClass { - IN = 1, -} - struct NameCacheEntry { name: String, expiry: Instant, @@ -57,25 +45,32 @@ impl VirtualDns { } // /* - pub fn receive_query(&mut self, data: &[u8]) -> Option> { + pub fn receive_query(&mut self, data: &[u8]) -> Result> { use crate::dns; - let mut dns_block = || { - let message = dns::parse_data_to_dns_message(data, false)?; - let qname = dns::extract_domain_from_dns_message(&message)?; - if let Some(ip) = self.allocate_ip(qname.clone()) { - let message = dns::build_dns_response(message, &qname, ip, 5)?; - message.to_vec() - } else { - Err("Virtual IP space for DNS exhausted".into()) - } - }; - dns_block().ok() + let message = dns::parse_data_to_dns_message(data, false)?; + let qname = dns::extract_domain_from_dns_message(&message)?; + let ip = self.allocate_ip(qname.clone())?; + let message = dns::build_dns_response(message, &qname, ip, 5)?; + Ok(message.to_vec()?) } // */ /* - pub fn receive_query(&mut self, data: &[u8]) -> Option> { + pub fn receive_query(&mut self, data: &[u8]) -> Result> { + #[derive(Eq, PartialEq, Debug)] + #[allow(dead_code, clippy::upper_case_acronyms)] + enum DnsRecordType { + A = 1, + AAAA = 28, + } + + #[derive(Eq, PartialEq, Debug)] + #[allow(dead_code)] + enum DnsClass { + IN = 1, + } + if data.len() < 17 { - return None; + return Err("Invalid DNS query".into()); } // bit 1: Message is a query (0) // bits 2 - 5: Standard query opcode (0) @@ -83,22 +78,22 @@ impl VirtualDns { // bit 7: Message is not truncated (0) // bit 8: Recursion desired (1) let is_supported_query = (data[2] & 0b11111011) == 0b00000001; - let num_queries = u16::from_be_bytes(data[4..6].try_into().ok()?); + let num_queries = u16::from_be_bytes(data[4..6].try_into()?); if !is_supported_query || num_queries != 1 { - return None; + return Err("Invalid DNS query".into()); } - let (qname, offset) = VirtualDns::parse_qname(data, 12)?; + let (qname, offset) = VirtualDns::parse_qname(data, 12).ok_or("parse_qname")?; if offset + 3 >= data.len() { - return None; + return Err("Invalid DNS query".into()); } - let qtype = u16::from_be_bytes(data[offset..offset + 2].try_into().ok()?); - let qclass = u16::from_be_bytes(data[offset + 2..offset + 4].try_into().ok()?); + let qtype = u16::from_be_bytes(data[offset..offset + 2].try_into()?); + let qclass = u16::from_be_bytes(data[offset + 2..offset + 4].try_into()?); if qtype != DnsRecordType::A as u16 && qtype != DnsRecordType::AAAA as u16 || qclass != DnsClass::IN as u16 { - return None; + return Err("Invalid DNS query".into()); } if qtype == DnsRecordType::A as u16 { @@ -131,7 +126,7 @@ impl VirtualDns { response[10] = 0; response[11] = 0; if qtype == DnsRecordType::A as u16 { - if let Some(ip) = self.allocate_ip(qname) { + if let Ok(ip) = self.allocate_ip(qname) { response.extend(&[ 0xc0, 0x0c, // Question name pointer 0, 1, // Record type: A @@ -154,10 +149,10 @@ impl VirtualDns { } else { response[7] = 0; // No answers } - Some(response) + Ok(response) } // */ - fn increment_ip(addr: IpAddr) -> Option { + fn increment_ip(addr: IpAddr) -> Result { let mut ip_bytes = match addr as IpAddr { IpAddr::V4(ip) => Vec::::from(ip.octets()), IpAddr::V6(ip) => Vec::::from(ip.octets()), @@ -176,36 +171,29 @@ impl VirtualDns { } } let addr = if addr.is_ipv4() { - let bytes: [u8; 4] = ip_bytes.as_slice().try_into().ok()?; + let bytes: [u8; 4] = ip_bytes.as_slice().try_into()?; IpAddr::V4(Ipv4Addr::from(bytes)) } else { - let bytes: [u8; 16] = ip_bytes.as_slice().try_into().ok()?; + let bytes: [u8; 16] = ip_bytes.as_slice().try_into()?; IpAddr::V6(Ipv6Addr::from(bytes)) }; - Some(addr) + Ok(addr) } // This is to be called whenever we receive or send a packet on the socket // which connects the tun interface to the client, so existing IP address to name // mappings to not expire as long as the connection is active. - pub fn touch_ip(&mut self, addr: &IpAddr) -> bool { - match self.lru_cache.get_mut(addr) { - None => false, - Some(entry) => { - entry.expiry = Instant::now() + Duration::from_secs(MAPPING_TIMEOUT); - true - } - } + pub fn touch_ip(&mut self, addr: &IpAddr) { + _ = self.lru_cache.get_mut(addr).map(|entry| { + entry.expiry = Instant::now() + Duration::from_secs(MAPPING_TIMEOUT); + }); } pub fn resolve_ip(&mut self, addr: &IpAddr) -> Option<&String> { - match self.lru_cache.get(addr) { - None => None, - Some(entry) => Some(&entry.name), - } + self.lru_cache.get(addr).map(|entry| &entry.name) } - fn allocate_ip(&mut self, name: String) -> Option { + fn allocate_ip(&mut self, name: String) -> Result { let now = Instant::now(); loop { @@ -223,9 +211,9 @@ impl VirtualDns { } if let Some(ip) = self.name_to_ip.get(&name) { - let result = Some(*ip); - self.touch_ip(&ip.clone()); - return result; + let ip = *ip; + self.touch_ip(&ip); + return Ok(ip); } let started_at = self.next_addr; @@ -235,16 +223,10 @@ impl VirtualDns { self.lru_cache.raw_entry_mut().from_key(&self.next_addr) { let expiry = Instant::now() + Duration::from_secs(MAPPING_TIMEOUT); - vacant.insert( - self.next_addr, - NameCacheEntry { - name: name.clone(), - expiry, - }, - ); - // e.insert(name.clone()); - self.name_to_ip.insert(name, self.next_addr); - return Some(self.next_addr); + let name0 = name.clone(); + vacant.insert(self.next_addr, NameCacheEntry { name, expiry }); + self.name_to_ip.insert(name0, self.next_addr); + return Ok(self.next_addr); } self.next_addr = Self::increment_ip(self.next_addr)?; if self.next_addr == self.broadcast_addr { @@ -252,7 +234,7 @@ impl VirtualDns { self.next_addr = self.network_addr; } if self.next_addr == started_at { - return None; + return Err("Virtual IP space for DNS exhausted".into()); } } }