From 5d92fde875d785594cab2a0b60db74425bbf4fb7 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sat, 8 Mar 2025 00:24:50 +0800 Subject: [PATCH] feat(shadowsocks-service): local-fake-dns switched db engine to rocksdb - ref #1895 BREAKING CHANGE: Database file must be recreated because engine have been switched from sled to rocksdb. --- Cargo.lock | 243 +++++++++------ crates/shadowsocks-service/Cargo.toml | 4 +- .../src/local/fake_dns/manager.rs | 279 ++++++++++-------- 3 files changed, 315 insertions(+), 211 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e29e7e75..bfdbc9d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -283,6 +283,26 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags 2.7.0", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.94", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -431,6 +451,16 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +[[package]] +name = "bzip2-sys" +version = "0.1.13+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "c2rust-bitfields" version = "0.19.0" @@ -484,6 +514,15 @@ dependencies = [ "subtle", ] +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "1.0.0" @@ -543,6 +582,17 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + [[package]] name = "clap" version = "4.5.31" @@ -875,6 +925,12 @@ dependencies = [ "signature", ] +[[package]] +name = "either" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7914353092ddf589ad78f25c5c1c21b7f80b0ff8621e7c814c3485b5306da9d" + [[package]] name = "elliptic-curve" version = "0.13.8" @@ -1051,16 +1107,6 @@ dependencies = [ "percent-encoding", ] -[[package]] -name = "fs2" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" -dependencies = [ - "libc", - "winapi", -] - [[package]] name = "fsevent-sys" version = "4.1.0" @@ -1175,15 +1221,6 @@ dependencies = [ "slab", ] -[[package]] -name = "fxhash" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" -dependencies = [ - "byteorder", -] - [[package]] name = "generator" version = "0.8.4" @@ -1249,6 +1286,12 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "glob" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" + [[package]] name = "group" version = "0.13.0" @@ -1392,7 +1435,7 @@ dependencies = [ "ipconfig", "moka", "once_cell", - "parking_lot 0.12.3", + "parking_lot", "quinn", "rand 0.9.0", "resolv-conf", @@ -1763,15 +1806,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "instant" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" -dependencies = [ - "cfg-if", -] - [[package]] name = "ipconfig" version = "0.3.2" @@ -1805,6 +1839,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.14" @@ -1887,6 +1930,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + [[package]] name = "libc" version = "0.2.170" @@ -1921,7 +1970,33 @@ checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d" dependencies = [ "bitflags 2.7.0", "libc", - "redox_syscall 0.5.8", + "redox_syscall", +] + +[[package]] +name = "librocksdb-sys" +version = "0.17.1+9.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b7869a512ae9982f4d46ba482c2a304f1efd80c6412a3d4bf57bb79a619679f" +dependencies = [ + "bindgen", + "bzip2-sys", + "cc", + "libc", + "libz-sys", + "lz4-sys", + "zstd-sys", +] + +[[package]] +name = "libz-sys" +version = "1.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9b68e50e6e0b26f672573834882eb57759f6db9b3be2ea3c35c91188bb4eaa" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", ] [[package]] @@ -1977,7 +2052,7 @@ dependencies = [ "log", "log-mdc", "once_cell", - "parking_lot 0.12.3", + "parking_lot", "rand 0.8.5", "serde", "serde-value", @@ -2008,6 +2083,16 @@ version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9106e1d747ffd48e6be5bb2d97fa706ed25b144fbee4d5c02eae110cd8d6badd" +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "managed" version = "0.8.0" @@ -2060,6 +2145,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.5" @@ -2091,7 +2182,7 @@ dependencies = [ "crossbeam-epoch", "crossbeam-utils", "loom", - "parking_lot 0.12.3", + "parking_lot", "portable-atomic", "rustc_version", "smallvec", @@ -2129,6 +2220,16 @@ dependencies = [ "libc", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "notify" version = "8.0.0" @@ -2311,17 +2412,6 @@ version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" -[[package]] -name = "parking_lot" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" -dependencies = [ - "instant", - "lock_api", - "parking_lot_core 0.8.6", -] - [[package]] name = "parking_lot" version = "0.12.3" @@ -2329,21 +2419,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" dependencies = [ "lock_api", - "parking_lot_core 0.9.10", -] - -[[package]] -name = "parking_lot_core" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60a2cfe6f0ad2bfc16aefa463b497d5c7a5ecd44a23efa72aa342d90177356dc" -dependencies = [ - "cfg-if", - "instant", - "libc", - "redox_syscall 0.2.16", - "smallvec", - "winapi", + "parking_lot_core", ] [[package]] @@ -2354,7 +2430,7 @@ checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" dependencies = [ "cfg-if", "libc", - "redox_syscall 0.5.8", + "redox_syscall", "smallvec", "windows-targets 0.52.6", ] @@ -2576,7 +2652,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.0", "rustls", "socket2", "thiserror 2.0.12", @@ -2594,7 +2670,7 @@ dependencies = [ "getrandom 0.2.15", "rand 0.8.5", "ring", - "rustc-hash", + "rustc-hash 2.1.0", "rustls", "rustls-pki-types", "slab", @@ -2694,15 +2770,6 @@ dependencies = [ "zerocopy 0.8.17", ] -[[package]] -name = "redox_syscall" -version = "0.2.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" -dependencies = [ - "bitflags 1.3.2", -] - [[package]] name = "redox_syscall" version = "0.5.8" @@ -2862,6 +2929,16 @@ dependencies = [ "signature", ] +[[package]] +name = "rocksdb" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26ec73b20525cb235bad420f911473b69f9fe27cc856c5461bccd7e4af037f43" +dependencies = [ + "libc", + "librocksdb-sys", +] + [[package]] name = "rpassword" version = "7.3.1" @@ -2909,6 +2986,12 @@ version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.0" @@ -3341,10 +3424,10 @@ dependencies = [ "pin-project", "rand 0.9.0", "regex", + "rocksdb", "rustls-native-certs", "serde", "shadowsocks", - "sled", "smoltcp", "socket2", "spin", @@ -3407,22 +3490,6 @@ dependencies = [ "autocfg", ] -[[package]] -name = "sled" -version = "0.34.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f96b4737c2ce5987354855aed3797279def4ebf734436c6aa4552cf8e169935" -dependencies = [ - "crc32fast", - "crossbeam-epoch", - "crossbeam-utils", - "fs2", - "fxhash", - "libc", - "log", - "parking_lot 0.11.2", -] - [[package]] name = "sm4" version = "0.5.1" @@ -3765,7 +3832,7 @@ dependencies = [ "bytes", "libc", "mio", - "parking_lot 0.12.3", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", diff --git a/crates/shadowsocks-service/Cargo.toml b/crates/shadowsocks-service/Cargo.toml index 26bf610e..dfd3fe28 100644 --- a/crates/shadowsocks-service/Cargo.toml +++ b/crates/shadowsocks-service/Cargo.toml @@ -83,7 +83,7 @@ local-socks4 = ["local"] # Enable Tun interface protocol for sslocal local-tun = ["local", "etherparse", "tun", "smoltcp"] # Enable Fake DNS -local-fake-dns = ["local", "trust-dns", "sled", "bson"] +local-fake-dns = ["local", "trust-dns", "rocksdb", "bson"] # sslocal support online URL (SIP008 Online Configuration Delivery) # https://shadowsocks.org/doc/sip008.html local-online-config = [ @@ -135,7 +135,7 @@ bytes = "1.7" byte_string = "1.0" byteorder = "1.5" rand = { version = "0.9", features = ["small_rng"] } -sled = { version = "0.34.7", optional = true } +rocksdb = { version = "0.23", optional = true } futures = "0.3" tokio = { version = "1.38", features = [ diff --git a/crates/shadowsocks-service/src/local/fake_dns/manager.rs b/crates/shadowsocks-service/src/local/fake_dns/manager.rs index dc2c2781..90d56c18 100644 --- a/crates/shadowsocks-service/src/local/fake_dns/manager.rs +++ b/crates/shadowsocks-service/src/local/fake_dns/manager.rs @@ -10,17 +10,40 @@ use std::{ use hickory_resolver::proto::rr::Name; use ipnet::{Ipv4AddrRange, Ipv4Net, Ipv6AddrRange, Ipv6Net}; -use log::{trace, warn}; -use sled::{Config as SledConfig, Db as SledDatabase}; +use log::{error, trace, warn}; +use rocksdb::DB as RocksDB; use tokio::sync::Mutex; use super::proto; -const FAKE_DNS_MANAGER_STORAGE_VERSION: u32 = 2; +const FAKE_DNS_MANAGER_STORAGE_VERSION: u32 = 3; + +/// Error type of FakeDns manager +#[derive(thiserror::Error, Debug)] +pub enum FakeDnsError { + /// std::io::Error wrapper + #[error("{0}")] + IoError(#[from] io::Error), + /// rocksdb::Error + #[error("{0}")] + RocksDBError(#[from] rocksdb::Error), +} + +impl From for io::Error { + fn from(value: FakeDnsError) -> Self { + match value { + FakeDnsError::IoError(e) => e, + FakeDnsError::RocksDBError(e) => io::Error::new(io::ErrorKind::Other, e), + } + } +} + +/// FakeDns Api Result type +pub type FakeDnsResult = Result; /// Fake DNS manager pub struct FakeDnsManager { - db: SledDatabase, + db: Mutex, ipv4_network: Mutex>, ipv6_network: Mutex>, expire_duration: Duration, @@ -28,15 +51,14 @@ pub struct FakeDnsManager { macro_rules! map_domain_ip { ($self:ident, $domain:ident, $addr_ty:ty, $addr_field:ident, $network_field:ident) => {{ + let db = $self.db.lock().await; let name2ip_key = FakeDnsManager::get_name2ip_key($domain); loop { - let name2ip_value = $self.db.get(&name2ip_key)?; - let mut domain_name_mapping = proto::DomainNameMapping::default(); - if let Some(ref v) = name2ip_value { - domain_name_mapping = proto::DomainNameMapping::decode(v)?; + if let Some(v) = db.get(&name2ip_key)? { + domain_name_mapping = proto::DomainNameMapping::decode(&v)?; if !domain_name_mapping.$addr_field.is_empty() { match domain_name_mapping.$addr_field.parse::<$addr_ty>() { @@ -45,43 +67,12 @@ macro_rules! map_domain_ip { let expire_secs = FakeDnsManager::get_current_timestamp() + $self.expire_duration.as_secs() as i64; - let mut reallocate = true; if domain_name_mapping.expire_time >= now { // Not expired yet. - reallocate = false; - } else { - // Expired. Try to reuse. - - let ip2name_key = FakeDnsManager::get_ip2name_key(i.into()); - let ip2name_value = $self.db.get(&ip2name_key)?; - if let Some(ref v) = ip2name_value { - let mut ip_mapping = proto::IpAddrMapping::decode(v)?; - if ip_mapping.domain_name == $domain.to_string() { - // Try to extend its expire time - ip_mapping.expire_time = expire_secs; - let nv = ip_mapping.encode_to_vec()?; - if let Ok(..) = - $self - .db - .compare_and_swap(&ip2name_key, ip2name_value.as_ref(), Some(nv)) - { - reallocate = false; - } else { - // CAS Failed, retry - continue; - } - } - } - } - - if !reallocate { domain_name_mapping.expire_time = expire_secs; let nv = domain_name_mapping.encode_to_vec()?; - // Ignore update error. It is ok if expire_time is updated by another thread. - let _ = $self - .db - .compare_and_swap(&name2ip_key, name2ip_value.as_ref(), Some(nv)); + db.put(&name2ip_key, nv)?; trace!( "fakedns mapping {} -> {}, expires {}", $domain, @@ -89,6 +80,27 @@ macro_rules! map_domain_ip { domain_name_mapping.expire_time ); return Ok((i, $self.expire_duration)); + } else { + // Expired. Try to reuse. + + let ip2name_key = FakeDnsManager::get_ip2name_key(i.into()); + if let Some(v) = db.get(&ip2name_key)? { + let mut ip_mapping = proto::IpAddrMapping::decode(&v)?; + if ip_mapping.domain_name == $domain.to_string() { + // Try to extend its expire time + ip_mapping.expire_time = expire_secs; + let nv = ip_mapping.encode_to_vec()?; + + db.put(&ip2name_key, nv)?; + trace!( + "fakedns mapping {} -> {}, expires {}", + $domain, + i, + domain_name_mapping.expire_time + ); + return Ok((i, $self.expire_duration)); + } + } } } Err(..) => { @@ -99,54 +111,42 @@ macro_rules! map_domain_ip { } // Allocate a new IPv4 address for this domain - 'alloc_network: while let Some(ip) = $self.$network_field.lock().await.next() { + while let Some(ip) = $self.$network_field.lock().await.next() { let ip2name_key = FakeDnsManager::get_ip2name_key(ip.into()); - loop { - let ip2name_value = $self.db.get(&ip2name_key)?; - if let Some(ref v) = ip2name_value { - let ip_mapping = proto::IpAddrMapping::decode(v)?; + if let Some(v) = db.get(&ip2name_key)? { + let ip_mapping = proto::IpAddrMapping::decode(&v)?; - let now = FakeDnsManager::get_current_timestamp(); - if ip_mapping.expire_time > now { - break; - } - } - - let mut ip_mapping = proto::IpAddrMapping::default(); - - let expire_secs = FakeDnsManager::get_current_timestamp() + $self.expire_duration.as_secs() as i64; - ip_mapping.expire_time = expire_secs; - ip_mapping.domain_name = $domain.to_string(); - - let nv = ip_mapping.encode_to_vec()?; - - if let Ok(..) = $self.db.compare_and_swap(&ip2name_key, ip2name_value, Some(nv.clone())) { - // Replace name2ip - - domain_name_mapping.$addr_field = ip.to_string(); - domain_name_mapping.expire_time = ip_mapping.expire_time; - let nv = domain_name_mapping.encode_to_vec()?; - - if let Ok(..) = $self - .db - .compare_and_swap(&name2ip_key, name2ip_value.as_ref(), Some(nv)) - { - trace!( - "fakedns mapping {} -> {}, expires {} created", - $domain, - ip, - domain_name_mapping.expire_time - ); - - return Ok((ip, $self.expire_duration)); - } else { - // name2ip CAS failed. Some other thread already allocated an address for this name. - let _ = $self.db.remove(&ip2name_key); - break 'alloc_network; - } + let now = FakeDnsManager::get_current_timestamp(); + if ip_mapping.expire_time > now { + continue; } } + + let mut ip_mapping = proto::IpAddrMapping::default(); + + let expire_secs = FakeDnsManager::get_current_timestamp() + $self.expire_duration.as_secs() as i64; + ip_mapping.expire_time = expire_secs; + ip_mapping.domain_name = $domain.to_string(); + + let nv = ip_mapping.encode_to_vec()?; + + db.put(&ip2name_key, nv)?; + // Replace name2ip + + domain_name_mapping.$addr_field = ip.to_string(); + domain_name_mapping.expire_time = ip_mapping.expire_time; + let nv = domain_name_mapping.encode_to_vec()?; + + db.put(&name2ip_key, nv)?; + trace!( + "fakedns mapping {} -> {}, expires {} created", + $domain, + ip, + domain_name_mapping.expire_time + ); + + return Ok((ip, $self.expire_duration)); } } }}; @@ -158,13 +158,25 @@ impl FakeDnsManager { ipv4_network: Ipv4Net, ipv6_network: Ipv6Net, expire_duration: Duration, - ) -> io::Result { - let db = SledConfig::new() - .cache_capacity(10 * 1024 * 1024) - .mode(sled::Mode::HighThroughput) - .flush_every_ms(Some(1_000)) - .path(db_path) - .open()?; + ) -> FakeDnsResult { + let db_path = db_path.as_ref(); + + // https://github.com/facebook/rocksdb/wiki/Setup-Options-and-Basic-Tuning + let mut db_options = rocksdb::Options::default(); + db_options.create_if_missing(true); + db_options.set_compression_type(rocksdb::DBCompressionType::Zstd); + db_options.set_bottommost_compression_type(rocksdb::DBCompressionType::Zstd); + db_options.set_bottommost_zstd_max_train_bytes(0, true); + db_options.set_max_background_jobs(6); + db_options.set_bytes_per_sync(1048576); + db_options.set_compaction_pri(rocksdb::CompactionPri::MinOverlappingRatio); + let mut db = match RocksDB::open(&db_options, db_path) { + Ok(db) => db, + Err(err) => { + error!("failed to open rocksdb, path: {}, error: {}", db_path.display(), err); + return Err(err.into()); + } + }; let ipv4_network_str = ipv4_network.to_string(); let ipv6_network_str = ipv6_network.to_string(); @@ -172,27 +184,50 @@ impl FakeDnsManager { let mut recreate_database = true; let key = "shadowsocks_fakedns_meta"; - if let Some(v) = db.get(key)? { - if let Ok(c) = proto::StorageMeta::decode(&v) { - if c.version == FAKE_DNS_MANAGER_STORAGE_VERSION { - if ipv4_network_str != c.ipv4_network || ipv6_network_str != c.ipv6_network { - warn!( - "IPv4 network {} (storage {}), IPv6 network {} (storage {}) not match", - ipv4_network_str, c.ipv4_network, ipv6_network_str, c.ipv6_network - ); + match db.get(key) { + Ok(Some(v)) => { + if let Ok(c) = proto::StorageMeta::decode(&v) { + if c.version == FAKE_DNS_MANAGER_STORAGE_VERSION { + if ipv4_network_str != c.ipv4_network || ipv6_network_str != c.ipv6_network { + warn!( + "IPv4 network {} (storage {}), IPv6 network {} (storage {}) not match", + ipv4_network_str, c.ipv4_network, ipv6_network_str, c.ipv6_network + ); + } else { + recreate_database = false; + } } else { - recreate_database = false; + warn!("storage version {} not match, recreating database", c.version); } } else { - warn!("storage version {} not match, recreating database", c.version); + warn!("storage meta parse failed. recreating database"); } - } else { - warn!("storage meta parse failed. recreating database"); + } + Ok(None) => { + // New DB without an META + } + Err(err) => { + error!("failed to get {}, error: {}", key, err); + return Err(err.into()); } } if recreate_database { - let _ = db.clear(); + drop(db); + let _ = RocksDB::destroy(&db_options, db_path); + + // Re-create by Open + db = match RocksDB::open(&db_options, db_path) { + Ok(db) => db, + Err(err) => { + error!( + "failed to recreate rocksdb, path: {}, error: {}", + db_path.display(), + err + ); + return Err(err.into()); + } + }; let c = proto::StorageMeta { ipv4_network: ipv4_network_str, @@ -201,13 +236,16 @@ impl FakeDnsManager { }; let v = c.encode_to_vec()?; - db.insert(key, v)?; + if let Err(err) = db.put(key, v) { + error!("failed to init storage, key: {}, error: {}", key, err); + return Err(err.into()); + } trace!("FakeDNS database created. {:?}", c); } Ok(FakeDnsManager { - db, + db: Mutex::new(db), ipv4_network: Mutex::new(ipv4_network.hosts().cycle()), ipv6_network: Mutex::new(ipv6_network.hosts().cycle()), expire_duration, @@ -233,33 +271,33 @@ impl FakeDnsManager { } /// Get or create an IPv4 mapping for `domain` - pub async fn map_domain_ipv4(&self, domain: &Name) -> io::Result<(Ipv4Addr, Duration)> { + pub async fn map_domain_ipv4(&self, domain: &Name) -> FakeDnsResult<(Ipv4Addr, Duration)> { map_domain_ip!(self, domain, Ipv4Addr, ipv4_addr, ipv4_network) } /// Get or create an IPv6 mapping for `domain` - pub async fn map_domain_ipv6(&self, domain: &Name) -> io::Result<(Ipv6Addr, Duration)> { + pub async fn map_domain_ipv6(&self, domain: &Name) -> FakeDnsResult<(Ipv6Addr, Duration)> { map_domain_ip!(self, domain, Ipv6Addr, ipv6_addr, ipv6_network) } /// Get IP mapped domain name - pub async fn map_ip_domain(&self, ip: IpAddr) -> io::Result> { + pub async fn map_ip_domain(&self, ip: IpAddr) -> FakeDnsResult> { + let db = self.db.lock().await; + + // ip -> domain_name let ip2name_key = FakeDnsManager::get_ip2name_key(ip); - - let ip2name_value = self.db.get(&ip2name_key)?; - match ip2name_value { + match db.get(&ip2name_key)? { None => Ok(None), - Some(ref v) => { - let mut ip_mapping = proto::IpAddrMapping::decode(v)?; + Some(v) => { + // Got ip -> domain_name + let mut ip_mapping = proto::IpAddrMapping::decode(&v)?; let now = FakeDnsManager::get_current_timestamp(); if ip_mapping.expire_time >= now { // Ok. It is not expired yet. Try to extend its expire time. ip_mapping.expire_time = now + self.expire_duration.as_secs() as i64; let nv = ip_mapping.encode_to_vec()?; - let _ = self - .db - .compare_and_swap(&ip2name_key, ip2name_value.as_ref(), Some(nv))?; + let _ = db.put(&ip2name_key, nv)?; // Update name2ip's expire time @@ -270,13 +308,12 @@ impl FakeDnsManager { { let name2ip_key = FakeDnsManager::get_name2ip_key(&name); - let name2ip_value = self.db.get(&name2ip_key)?; - match name2ip_value { - Some(ref v) => { - let mut domain_name_mapping = proto::DomainNameMapping::decode(v)?; + match db.get(&name2ip_key)? { + Some(v) => { + let mut domain_name_mapping = proto::DomainNameMapping::decode(&v)?; domain_name_mapping.expire_time = ip_mapping.expire_time; let nv = domain_name_mapping.encode_to_vec()?; - let _ = self.db.compare_and_swap(&name2ip_key, name2ip_value.as_ref(), Some(nv)); + let _ = db.put(&name2ip_key, nv)?; } None => { // Interesting. No name2ip.