From 79fee04aefa3da0884a3bfa4933f772ebecd0c8c Mon Sep 17 00:00:00 2001 From: Xnoe Date: Fri, 5 Dec 2025 22:03:11 +0000 Subject: [PATCH] Squash commits. --- .gitignore | 3 + Cargo.lock | 459 ++++++++++++++++++++++++++++++++++++++ Cargo.toml | 10 + src/cursor.rs | 60 +++++ src/dns_parser.rs | 131 +++++++++++ src/ip_pool.rs | 47 ++++ src/main.rs | 551 ++++++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 1261 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/cursor.rs create mode 100644 src/dns_parser.rs create mode 100644 src/ip_pool.rs create mode 100644 src/main.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d21e953 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/target +test.ini +.vscode \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..8d1a9ba --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,459 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli", +] + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "anyhow" +version = "1.0.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" + +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", + "windows-link", +] + +[[package]] +name = "bitflags" +version = "2.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2261d10cca569e4643e526d8dc2e62e433cc8aba21ab764233731f8d369bf394" + +[[package]] +name = "bytes" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" + +[[package]] +name = "cfg-if" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "dlv-list" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "442039f5147480ba31067cb00ada1adae6892028e40e45fc5de7b7df6dcc1b5f" +dependencies = [ + "const-random", +] + +[[package]] +name = "getrandom" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "io-uring" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "046fa2d4d00aea763528b4950358d0ead425372445dc8ff86312b3c69ff7727b" +dependencies = [ + "bitflags", + "cfg-if", + "libc", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.176" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58f929b4d672ea937a23a1ab494143d968337a5f47e56d0815df1e0890ddf174" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "memchr" +version = "2.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", +] + +[[package]] +name = "mio" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78bed444cc8a2160f01cbcf811ef18cac863ad68ae8ca62092e8db51d51c761c" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" + +[[package]] +name = "ordered-multimap" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49203cdcae0030493bad186b28da2fa25645fa276a51b6fec8010d281e02ef79" +dependencies = [ + "dlv-list", + "hashbrown", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" + +[[package]] +name = "proc-macro2" +version = "1.0.101" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rust-dnsproxy-thing" +version = "0.1.0" +dependencies = [ + "anyhow", + "lazy_static", + "rust-ini", + "tokio", +] + +[[package]] +name = "rust-ini" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "796e8d2b6696392a43bea58116b667fb4c29727dc5abd27d6acf338bb4f688c7" +dependencies = [ + "cfg-if", + "ordered-multimap", +] + +[[package]] +name = "rustc-demangle" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f7d92ca342cea22a06f2121d944b4fd82af56988c270852495420f961d4ace" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "signal-hook-registry" +version = "1.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a4719bff48cee6b39d12c020eeb490953ad2443b7055bd0b21fca26bd8c28b" +dependencies = [ + "libc", +] + +[[package]] +name = "slab" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a2ae44ef20feb57a68b23d846850f861394c2e02dc425a50098ae8c90267589" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "socket2" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "233504af464074f9d066d7b5416c5f9b894a5862a6506e306f7b816cdd6f1807" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "tokio" +version = "1.47.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89e49afdadebb872d3145a5638b59eb0691ea23e46ca484037cfab3b76b95038" +dependencies = [ + "backtrace", + "bytes", + "io-uring", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "slab", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "windows-link" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45e46c0661abb7180e7b9c281db115305d49ca1709ab8242adf09666d2173c65" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..e19d946 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "rust-dnsproxy-thing" +version = "0.1.0" +edition = "2024" + +[dependencies] +tokio = { version = "1.47.1", features = ["full"] } +anyhow = "1.0.100" +lazy_static = "1.5.0" +rust-ini = "0.21.3" \ No newline at end of file diff --git a/src/cursor.rs b/src/cursor.rs new file mode 100644 index 0000000..1174e64 --- /dev/null +++ b/src/cursor.rs @@ -0,0 +1,60 @@ +#[derive(Clone)] +pub struct Cursor<'a, T> { + buf: &'a [T], + index: usize, +} + +impl<'a, T> Cursor<'a, T> { + pub fn from(buf: &'a [T]) -> Self { + Self { buf, index: 0 } + } + + pub fn index(&mut self) -> usize { + self.index + } + + pub fn next(&mut self) -> Option<&T> { + let next_index = self.index + 1; + if next_index >= self.buf.len() { + None + } else { + let v = &self.buf[self.index]; + self.index = next_index; + Some(v) + } + } + + pub fn seek(&mut self, location: usize) -> Result<(), ()> { + if location >= self.buf.len() { + Err(()) + } else { + self.index = location; + Ok(()) + } + } + + pub fn next_slice(&mut self, amount: usize) -> Option<&'a [T]> { + let next_index = self.index + amount; + if next_index >= self.buf.len() { + None + } else { + let slice = &self.buf[self.index..next_index]; + self.index = next_index; + Some(slice) + } + } + + pub fn next_array(&mut self) -> Option<[T; N]> where [T; N]: TryFrom<&'a [T]> { + Some(self.next_slice(N)?.try_into().ok()?) + } + + pub fn forward(&mut self, amount: usize) -> Result<(), ()> { + let next_index = self.index + amount; + if next_index >= self.buf.len() { + Err(()) + } else { + self.index = next_index; + Ok(()) + } + } +} \ No newline at end of file diff --git a/src/dns_parser.rs b/src/dns_parser.rs new file mode 100644 index 0000000..9ba8272 --- /dev/null +++ b/src/dns_parser.rs @@ -0,0 +1,131 @@ +use crate::cursor::Cursor; + +pub fn dns_name_to_parts(cursor: &mut Cursor) -> Option>> { + let mut parts: Vec> = Vec::new(); + let mut ptr_depth = 0; + let mut restore_point = None; + while let Some(&(mut byte)) = cursor.next() { + if byte == 0 { + break; + } + if byte >= 192 { + if let None = restore_point { + restore_point = Some(cursor.index() + 1); + } + if ptr_depth >= 16 { + return None; + } + let ptr_lsb = *cursor.next()?; + cursor + .seek(u16::from_be_bytes([byte & 0b0011_1111, ptr_lsb]) as usize) + .ok()?; + byte = *cursor.next()?; + ptr_depth += 1; + } + parts.push(cursor.next_slice(byte as usize)?.to_vec()); + } + if let Some(position) = restore_point { + cursor.seek(position).ok()?; + } + Some(parts) +} + +pub fn dns_name_len(buffer: &mut Cursor) -> Option { + let mut length = 0; + while let Some(&byte) = buffer.next() { + if byte == 0 { + return Some(length + 1); + } + if byte >= 192 { + return Some(length + 2); + } + buffer.forward(byte as usize).ok()?; + length += byte as usize + 1; + } + Some(length) +} + +pub fn parts_to_dns_name(parts: &Vec>) -> Vec { + let mut result = Vec::new(); + for part in parts { + result.push(part.len() as u8); + result.extend_from_slice(&part); + } + result.push(0u8); + result +} + +pub fn dns_parts_to_string(parts: &Vec>) -> String { + let mut result = String::new(); + for part in parts { + result.push_str(&String::from_utf8_lossy(&part)); + result.push('.'); + } + return result; +} + +pub fn string_to_dns_name(string: String) -> Vec { + let mut result = Vec::new(); + + for part in string.split('.') { + result.push(part.len() as u8); + result.extend_from_slice(part.as_bytes()); + } + result.push(0u8); + return result; +} + +pub struct AnswerIterator<'a> { + cursor: Cursor<'a, u8>, + ancount: u16, +} + +impl<'a> AnswerIterator<'a> { + pub fn from(buf: &'a [u8]) -> Option { + let mut cursor = Cursor::from(buf); + cursor.seek(4).ok()?; + let qdcount = u16::from_be_bytes(cursor.next_array::<2>()?); + let ancount = u16::from_be_bytes(cursor.next_array::<2>()?); + cursor.forward(4).ok()?; + + // Skip past the question section + for _ in 0..qdcount { + dns_name_len(&mut cursor)?; + cursor.forward(4).ok()?; + } + Some(Self { cursor, ancount }) + } +} + +#[derive(PartialEq, Debug)] +pub struct Answer { + pub name: Vec>, + pub rrtype: u16, + pub rrclass: u16, + pub ttl: u32, + pub rdata: Vec, +} + +impl<'a> Iterator for AnswerIterator<'a> { + type Item = Answer; + + fn next(&mut self) -> Option { + if self.ancount == 0 { + return None; + } + self.ancount -= 1; + let name = dns_name_to_parts(&mut self.cursor)?; + let rrtype = u16::from_be_bytes(self.cursor.next_array::<2>()?); + let rrclass = u16::from_be_bytes(self.cursor.next_array::<2>()?); + let ttl = u32::from_be_bytes(self.cursor.next_array::<4>()?); + let rdlength = u16::from_be_bytes(self.cursor.next_array::<2>()?) as usize; + let rdata = self.cursor.next_slice(rdlength)?.to_vec(); + Some(Answer { + name, + rrtype, + rrclass, + ttl, + rdata, + }) + } +} diff --git a/src/ip_pool.rs b/src/ip_pool.rs new file mode 100644 index 0000000..8a0a572 --- /dev/null +++ b/src/ip_pool.rs @@ -0,0 +1,47 @@ +use std::collections::{HashSet, VecDeque}; +use std::net::Ipv4Addr; + +pub struct IpPool { + free_ips: VecDeque, + allocated_ips: HashSet, +} + +impl IpPool { + pub fn new(base_addr: Ipv4Addr, subnet_prefix_size: u8) -> Option { + if subnet_prefix_size == 32 { + return None; + } + let subnet_prefix_mask: u32 = u32::MAX << (32 - subnet_prefix_size); + let base_addr_int = >::into(base_addr); + if base_addr_int & !subnet_prefix_mask != 0 { + return None; + } + let last_address_number = u32::MAX & !subnet_prefix_mask; + let mut pool = VecDeque::with_capacity(last_address_number as usize); + for number in 0..=(last_address_number as usize) { + pool.push_back(Ipv4Addr::from(base_addr_int + (number as u32))); + } + + Some(Self { + free_ips: pool, + allocated_ips: HashSet::new(), + }) + } + + pub fn acquire(&mut self) -> Option { + match self.free_ips.pop_front() { + ip @ Some(addr) => { + self.allocated_ips.insert(addr); + ip + } + None => None, + } + } + + pub fn release(&mut self, addr: Ipv4Addr) { + if self.allocated_ips.contains(&addr) { + self.allocated_ips.remove(&addr); + self.free_ips.push_back(addr); + } + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..a38f77c --- /dev/null +++ b/src/main.rs @@ -0,0 +1,551 @@ +#[macro_use] +extern crate lazy_static; + +mod cursor; +mod dns_parser; +mod ip_pool; + +use std::collections::HashMap; +use std::net::Ipv4Addr; +use std::process::{Command, Output}; +use std::sync::OnceLock; +use std::time::{Duration, Instant}; + +use ini::Ini; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream, UdpSocket}, + sync::Mutex, +}; + +use crate::cursor::Cursor; +use crate::dns_parser::{ + Answer, AnswerIterator, dns_name_to_parts, dns_parts_to_string, + parts_to_dns_name, string_to_dns_name, +}; +use crate::ip_pool::IpPool; + +fn run_command_string(command: String) -> Result { + let command_vec = { + if COMMAND_PREFIX.get().unwrap().eq("") { + command.split(' ').collect::>() + } else { + COMMAND_PREFIX + .get() + .unwrap() + .split(' ') + .chain(command.split(' ')) + .collect::>() + } + }; + + Command::new(command_vec[0]) + .args(&command_vec[1..]) + .output() +} + +async fn setup_forwarding( + forwarding_map: &mut HashMap, Vec>, + dns_name: &Vec, + ttl: u32, + original_ip: Ipv4Addr, + forged_ip: Ipv4Addr, +) { + let forwarding = Forwarding { + expires: Instant::now() + Duration::from_secs(ttl as u64), + forged_ip: forged_ip, + original_ip: original_ip, + ttl: ttl, + }; + + match forwarding_map.get_mut(dns_name) { + Some(forwarding_list) => { + forwarding_list.push(forwarding); + } + None => { + forwarding_map.insert(dns_name.clone(), vec![forwarding]); + } + } + + let mark_command = format!( + "iptables -t mangle -A PREROUTING -d {} -j MARK --set-mark {}", + forged_ip, + FwmarkConfigMap.lock().await.get(dns_name).unwrap() + ); + let _mark_output = run_command_string(mark_command).unwrap(); + + let dnat_command = format!( + "iptables -t nat -A PREROUTING -d {} -j DNAT --to {}", + forged_ip, original_ip + ); + let _dnat_output = run_command_string(dnat_command).unwrap(); +} + +async fn teardown_forwarding( + forwarding_map: &mut HashMap, Vec>, + dns_name: &Vec, + original_ip: Ipv4Addr, +) { + let forwarding_list = match forwarding_map.get_mut(dns_name) { + Some(f) => f, + None => return, + }; + + let forwarding = match forwarding_list + .iter() + .find(|f| f.original_ip == original_ip) + { + Some(f) => f, + None => return, + } + .clone(); + + forwarding_list.retain(|f| f.original_ip != original_ip); + + if forwarding_list.len() == 0 { + forwarding_map.remove(dns_name); + } + + IpAllocator.lock().await.release(forwarding.forged_ip); + let mark_command = format!( + "iptables -t mangle -D PREROUTING -d {} -j MARK --set-mark {}", + forwarding.forged_ip, + FwmarkConfigMap.lock().await.get(dns_name).unwrap() + ); + let _mark_output = run_command_string(mark_command).unwrap(); + + let dnat_command = format!( + "iptables -t nat -D PREROUTING -d {} -j DNAT --to {}", + forwarding.forged_ip, forwarding.original_ip + ); + let _dnat_output = run_command_string(dnat_command).unwrap(); +} + +async fn query_upstream_resolvers(original_message: &[u8], tcp: bool) -> anyhow::Result> { + let mut upstream_reply = Vec::new(); + + if tcp { + let mut upstream = TcpStream::connect(UPSTREAM_DNS.get().unwrap()).await?; + //upstream.write(&(original_message.len() as u16).to_be_bytes()).await?; + upstream.write(&original_message).await?; + let mut size_buffer = [0u8; 2]; + upstream.read(&mut size_buffer).await?; + let size: u16 = u16::from_be_bytes(size_buffer); + upstream_reply.resize(size as usize + 2, 0u8); + upstream.read(&mut upstream_reply[2..]).await?; + upstream_reply[0] = size_buffer[0]; + upstream_reply[1] = size_buffer[1]; + } else { + let upstream = UdpSocket::bind("0.0.0.0:0").await?; + upstream.connect(UPSTREAM_DNS.get().unwrap()).await?; + upstream.send(original_message).await?; + upstream_reply.resize(512, 0); + upstream.recv_from(&mut upstream_reply[..512]).await?; + } + + return Ok(upstream_reply); +} + +fn forge_replies(replies: &Vec, dns_name_string: String, qname_parts: Vec>, original_message: &[u8], reply_buf: &mut Vec, tcp: bool) { + let reply: [u8; 12] = [ + 0u8, + 0, // ID + 0b1000_0000, + 0b0000_0000, // Flags + 0, + 1, // Qdcount + 0, + 1, // Ancount + 0, + 0, // Nscount + 0, + 0, // Arcount + ]; + + let offset = if tcp { 2 } else { 0 }; + + let now = Instant::now(); + + println!( + "Forging reply for {}, original IPs: {:?}, forged IPs: {:?}", + dns_name_string, + replies + .iter() + .map(|e| e.original_ip) + .collect::>(), + replies + .iter() + .map(|e| e.forged_ip) + .collect::>() + ); + + let mut new_reply = reply.clone().to_vec(); + new_reply[0..2].copy_from_slice(&original_message[offset..][0..=1]); + new_reply[6..8].copy_from_slice(&(replies.len() as u16).to_be_bytes()); + new_reply.extend_from_slice(&original_message[offset..][12..=12 + qname_parts.iter().map(|p| p.len()+1).sum::() + 1 + 3]); + for reply in replies { + new_reply.extend_from_slice(&parts_to_dns_name(&qname_parts)); + new_reply.extend_from_slice(&[0, 1]); + new_reply.extend_from_slice(&[0, 1]); + new_reply.extend_from_slice(&((reply.expires - now).as_secs() as u32).to_be_bytes()); + new_reply.extend_from_slice(&[0, 4]); + new_reply.extend_from_slice(&reply.forged_ip.octets()); + } + if tcp { + reply_buf.extend_from_slice(&(new_reply.len() as u16).to_be_bytes()); + } + reply_buf.extend_from_slice(&new_reply); +} + +async fn handle_dns_response(buf: &[u8], reply_buf: &mut Vec, tcp: bool) -> anyhow::Result<()> { + + let offset: usize = if tcp { 2 } else { 0 }; + let mut cursor = Cursor::from(&buf[offset..]); + + // Identify some metadata from the query + if cursor.seek(4).is_err() { + return Err(anyhow::Error::msg("Failed to seek to QDCount")); + } + let qdcount = u16::from_be_bytes(cursor.next_array::<2>().ok_or(anyhow::Error::msg("Failed to read QDCount"))?); + if qdcount != 1 { + eprintln!("Got qdcount: {}", qdcount); + return Err(anyhow::Error::msg( + "Missing question from query. Got qdcount {}", + )); + } + if cursor.forward(6).is_err() { + return Err(anyhow::Error::msg("Failed to seek to question section")); + } + let qname_parts = + dns_name_to_parts(&mut cursor).ok_or(anyhow::Error::msg("Failed to decode QName."))?; + let qtype = u16::from_be_bytes(cursor.next_array::<2>().ok_or(anyhow::Error::msg("Failed to read QType."))?); + let qclass = u16::from_be_bytes(cursor.next_array::<2>().ok_or(anyhow::Error::msg("Failed to read QClass."))?); + + // If the query is for anything other than qclass IN or qtype A, just forward the query upstream + if qclass != 1 || qtype != 1 { + let upstream_reply = query_upstream_resolvers(buf, tcp).await?; + reply_buf.extend_from_slice(&upstream_reply); + return Ok(()) + } + + let now = Instant::now(); + + let dns_name = parts_to_dns_name(&qname_parts); + let dns_name_string = dns_parts_to_string(&qname_parts); + + let entries = match ForwardingMap.lock().await.get_mut(&dns_name) { + Some(forwardings) => forwardings.clone(), + None => Vec::new(), + }; + + // Let's first lookup the qname in the Forwardings to see if we have non-expired answers + if entries.len() > 0 && entries.iter().all(|e| e.expires > now) { + forge_replies(&entries, dns_name_string, qname_parts, buf, reply_buf, tcp); + return Ok(()); + } else { + let upstream_reply = query_upstream_resolvers(buf, tcp).await?; + + // Try an answer from the upstream response that has type A. + let a_answers = match AnswerIterator::from(&upstream_reply[offset..]) { + Some(answers) => answers + .filter(|Answer { rrtype, .. }| *rrtype == 1) + .collect::>(), + None => { + return Err(anyhow::Error::msg( + "Failed to extract answers from upstream reply!", + )); + } + }; + + let mut forge = true; + if qtype != 1 { + // Only forge for queries with an A qtype + eprintln!("Not forging due non-A type question."); + forge = false; + } + + if a_answers.len() == 0 { + // If no A type answer, don't forge + eprintln!("Not forging due to no returned A type answers."); + forge = false; + } + + if a_answers.iter().any(|a| a.rdata.len() != 4) { + eprintln!("Not forging due to malformed A type answer.",); + forge = false; + } + + if !FwmarkConfigMap + .lock() + .await + .contains_key(&parts_to_dns_name(&qname_parts)) + { + eprintln!("Not forging due to non-matching qname."); + forge = false; + } + + if forge { + // Normalise a_answers so we're working with Ipv4Addr + let normalised_answers: Vec<(Vec>, u32, Ipv4Addr)> = a_answers + .iter() + .map(|answer| { + ( + answer.name.clone(), + answer.ttl, + Ipv4Addr::from( + as TryInto<[u8; 4]>>::try_into(answer.rdata.clone()).unwrap(), + ), + ) + }) + .collect::<_>(); + + let mut replies: Vec = Vec::new(); + + // Determine if we need to create or renew our entries + if entries.len() > 0 && entries.iter().all(|e| e.expires > now) { + println!("Found not expired forwardings for {}", dns_name_string); + replies.extend(entries.clone() as Vec); + } else { + // We want to identify which of our A answers already exist in the ForwardingMap + let existing_entries = entries + .iter() + .filter(|e| { + normalised_answers + .iter() + .any(|(_, _, addr)| e.original_ip == *addr) + }) + .collect::>(); + + // Let's also identify which entries don't match any of the replies + let nonexisting_entries = entries + .iter() + .filter(|e| { + !normalised_answers + .iter() + .any(|(_, _, addr)| e.original_ip == *addr) + }) + .collect::>(); + + // And now let's find the replies that don't match any of the current entries + let new_answers = normalised_answers + .iter() + .filter(|answer| entries.iter().all(|e| e.original_ip != answer.2)) + .collect::>(); + + // Acquire the forwarding map + let mut forwarding_map = ForwardingMap.lock().await; + + // Remove all non-existing entries + for entry in nonexisting_entries { + println!( + "Removed forwarding for {} with real IP {} / forged IP {} as the upstream no longer returns this value", + dns_name_string, entry.original_ip, entry.forged_ip + ); + teardown_forwarding(&mut forwarding_map, &dns_name, entry.original_ip).await; + } + + // Add new answers + for (_, ttl, original_ip) in new_answers { + let forged_ip = IpAllocator.lock().await.acquire().unwrap(); + setup_forwarding( + &mut forwarding_map, + &dns_name, + *ttl, + *original_ip, + forged_ip, + ) + .await; + println!( + "Added forwarding for {} with real IP {} / forged IP {}", + dns_name_string, original_ip, forged_ip + ); + } + + // For all the existing entries, update their TTL if they're expired + for entry in existing_entries.iter() { + println!( + "Updating TTL of existing entry {}/{}", + dns_name_string, entry.forged_ip + ); + if entry.expires < now { + teardown_forwarding(&mut forwarding_map, &dns_name, entry.original_ip).await; + setup_forwarding( + &mut forwarding_map, + &dns_name, + entry.ttl, + entry.original_ip, + entry.forged_ip, + ) + .await; + } + } + + let entries = forwarding_map.get(&dns_name).unwrap().clone(); + replies.extend(entries); + } + + forge_replies(&replies, dns_name_string, qname_parts, buf, reply_buf, tcp); + } else { + reply_buf.extend_from_slice(&upstream_reply); + } + + Ok(()) + } +} + +#[derive(Clone, Debug)] +struct Forwarding { + expires: Instant, + forged_ip: Ipv4Addr, + original_ip: Ipv4Addr, + ttl: u32, +} + +lazy_static! { + static ref FwmarkConfigMap: Mutex, u32>> = Mutex::new(HashMap::new()); + static ref ForwardingMap: Mutex, Vec>> = Mutex::new(HashMap::new()); + static ref IpAllocator: Mutex = + Mutex::new(IpPool::new(Ipv4Addr::new(100, 64, 0, 0), 24).unwrap()); +} + +static COMMAND_PREFIX: OnceLock = OnceLock::new(); +static UPSTREAM_DNS: OnceLock = OnceLock::new(); +static LISTEN_ADDR: OnceLock = OnceLock::new(); + +async fn tcp_handler(listener: TcpListener) { + loop { + let accepted = listener.accept().await; + if let Err(e) = &accepted { + eprintln!("[TCP] Failed to accept TCP socket with error: {:?}", e); + } + + let (mut socket, addr) = accepted.unwrap(); + let mut size_buf = [0u8; 2]; + let Ok(2) = socket.read(&mut size_buf).await else { + eprintln!("[TCP] Failed to read message size from {}", addr); + continue; + }; + let message_size = u16::from_be_bytes(size_buf) as usize; + let mut message_buffer = Vec::new(); + message_buffer.resize(message_size + 2, 0); + match socket.read(&mut message_buffer[2..]).await { + Ok(n) => { + if n != message_size { + eprintln!( + "[TCP] Received too few bytes than expected from {}. Message size indicated {}, read {}", + addr, message_size, n + ); + continue; + } + } + Err(e) => { + eprintln!( + "[TCP] Failed to read message from {} with error: {:?}", + addr, e + ); + continue; + } + } + message_buffer[0] = size_buf[0]; + message_buffer[1] = size_buf[1]; + let mut reply = Vec::new(); + if let Err(e) = handle_dns_response(&message_buffer, &mut reply, true).await { + eprintln!( + "[TCP] Received error when handling response for {}: {:?}", + addr, e + ); + continue; + } + if let Err(e) = socket.write(&reply).await { + eprintln!( + "[TCP] Received error when sending response to {}: {:?}", + addr, e + ); + continue; + } + } +} + +async fn udp_handler(socket: UdpSocket) { + loop { + let mut message_buffer = [0u8; 512]; + match socket.recv_from(&mut message_buffer).await { + Ok((_, addr)) => { + let mut reply = Vec::new(); + if let Err(e) = handle_dns_response(&message_buffer, &mut reply, false).await { + eprintln!( + "[UDP] Received error when handling response for {}: {:?}", + addr, e + ); + continue; + } + if let Err(e) = socket.send_to(&reply, addr).await { + eprintln!( + "[UDP] Received error when sending response to {}: {:?}", + addr, e + ); + continue; + } + } + Err(e) => { + eprintln!("[UDP] Failed to read message with error: {:?}", e); + continue; + } + } + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let mut args = std::env::args(); + let _ = args.next(); + let config_file_path = match args.next() { + Some(p) => p, + None => String::from("/etc/rust-dns-selective-routing/config.ini"), + }; + let i = Ini::load_from_file(config_file_path).unwrap(); + + let map: HashMap> = i + .into_iter() + .filter_map(|(k, it)| match k { + Some(s) => Some((s, it.into_iter().collect())), + None => None, + }) + .collect(); + + if let Some(v) = map.get("config") { + if let Some(cmd) = v.get("command_prefix") { + COMMAND_PREFIX.set(cmd.clone()).unwrap(); + } else { + COMMAND_PREFIX.set(String::from("")).unwrap(); + } + if let Some(upstream) = v.get("upstream_dns") { + UPSTREAM_DNS.set(upstream.clone()).unwrap(); + } else { + UPSTREAM_DNS.set(String::from("1.1.1.1:53")).unwrap(); + } + if let Some(listen_addr) = v.get("listen_addr") { + LISTEN_ADDR.set(listen_addr.clone()).unwrap(); + } else { + LISTEN_ADDR.set(String::from("127.0.0.1:53")).unwrap(); + } + } + + if let Some(v) = map.get("domains") { + let mut fwmark_config = FwmarkConfigMap.lock().await; + for (domain, fwmark) in v.iter() { + if let Ok(fwmark) = fwmark.parse::() { + let _ = fwmark_config.insert(string_to_dns_name(domain.to_string()), fwmark); + } + } + } + + let recv_socket_tcp = TcpListener::bind(LISTEN_ADDR.get().unwrap()).await?; + let recv_socket_udp = UdpSocket::bind(LISTEN_ADDR.get().unwrap()).await?; + + tokio::try_join!( + tokio::spawn(tcp_handler(recv_socket_tcp)), + tokio::spawn(udp_handler(recv_socket_udp)), + )?; + Ok(()) +}