diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 25212a7..70537cd 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -2315,6 +2315,7 @@ dependencies = [ "tauri-plugin-opener", "tempfile", "tokio", + "tokio-util", "tracing", "tracing-subscriber", "uuid", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index d0f3984..86ce60a 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -9,6 +9,7 @@ chrono = { version = "0.4", features = ["serde"] } uuid = { version = "1.8", features = ["v4", "serde"] } thiserror = "1.0" tokio = { version = "1.37", features = ["full"] } +tokio-util = "0.7" tempfile = "3.10" dirs = "5" parking_lot = "0.12" @@ -36,6 +37,7 @@ serde_json = { workspace = true } chrono = { workspace = true } uuid = { workspace = true } tokio = { workspace = true } +tokio-util = { workspace = true } reqwest = { version = "0.13", features = ["json"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/src-tauri/crates/mhost-core/src/models.rs b/src-tauri/crates/mhost-core/src/models.rs index 9b6ca61..f0ed5d9 100644 --- a/src-tauri/crates/mhost-core/src/models.rs +++ b/src-tauri/crates/mhost-core/src/models.rs @@ -190,6 +190,87 @@ pub enum RuleSource { AdBlock(ExternalSource), } +// --------------------------------------------------------------------------- +// AdBlock (issue #130) +// --------------------------------------------------------------------------- +// +// 广告屏蔽状态与配置。**仅在 DNS 模式下生效** —— hosts 模式不再承担 +// 广告屏蔽职责(早期尝试因 `/etc/hosts` 膨胀失败)。 +// +// 数据布局(mhost-storage): +// {root}/adblock.json # AdBlockState 整体序列化 +// {root}/adblock-cache/{id}.txt # 每源原始 hosts 内容 + +/// Per-source response when a domain hits an ad block rule. +/// +/// `ZeroAddress` returns a 0.0.0.0 A record (clients typically retry-then-fail +/// after a timeout). `NxDomain` returns NXDOMAIN immediately (more aggressive +/// but some clients surface it as an error). +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum AdBlockResponse { + #[default] + ZeroAddress, + NxDomain, +} + +/// A remote ad block subscription source. +/// +/// One source = one URL of hosts-format blocklist. Persisted as part of +/// `AdBlockState`. The cached fetched content lives at +/// `{root}/adblock-cache/{source_id}.txt` so that DNS mode can keep serving +/// blocklist hits even when the remote URL is temporarily unreachable. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AdBlockSource { + pub source_id: SourceId, + pub name: String, + pub url: String, + pub enabled: bool, + pub response: AdBlockResponse, + /// RFC 3339 timestamp of the last successful fetch. `None` if never fetched. + #[serde(skip_serializing_if = "Option::is_none")] + pub last_fetched_at: Option>, + /// Last fetch error message (transport or non-2xx). Cleared on next success. + #[serde(skip_serializing_if = "Option::is_none")] + pub last_error: Option, + /// Number of rules parsed from the last successful fetch. + pub rule_count: usize, + /// HTTP ETag from the last successful fetch (reserved for future + /// conditional GETs — unused in v1 but persisted so we don't need a + /// migration when conditional fetch lands). + #[serde(skip_serializing_if = "Option::is_none")] + pub etag: Option, +} + +/// Persistent state for the DNS-mode ad block subsystem. +/// +/// Stored as a single JSON document (`{root}/adblock.json`). All fields are +/// mutable from the IPC surface; the Rust side owns the on-disk write path +/// (`mhost-storage/src/adblock.rs`). +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AdBlockState { + pub enabled: bool, + pub sources: Vec, + pub whitelist: Vec, + pub auto_refresh_enabled: bool, + /// Hours between background refreshes. `0` disables background refresh + /// regardless of `auto_refresh_enabled`. Clamped to a sane range at the + /// IPC boundary. + pub refresh_interval_hours: u32, +} + +impl Default for AdBlockState { + fn default() -> Self { + Self { + enabled: false, + sources: Vec::new(), + whitelist: Vec::new(), + auto_refresh_enabled: true, + refresh_interval_hours: 24, + } + } +} + // --------------------------------------------------------------------------- // ExportFormat // --------------------------------------------------------------------------- @@ -579,6 +660,110 @@ mod tests { assert_eq!(source, restored); } + // ----------------------------------------------------------------------- + // AdBlock tests (issue #130) + // ----------------------------------------------------------------------- + + #[test] + fn test_ad_block_response_default_is_zero_address() { + assert_eq!(AdBlockResponse::default(), AdBlockResponse::ZeroAddress); + } + + #[test] + fn test_ad_block_response_serde_roundtrip() { + for variant in [AdBlockResponse::ZeroAddress, AdBlockResponse::NxDomain] { + let json = serde_json::to_string(&variant).unwrap(); + let restored: AdBlockResponse = serde_json::from_str(&json).unwrap(); + assert_eq!(variant, restored); + } + // snake_case wire format (serde `rename_all = "snake_case"` treats the + // capital N in `NxDomain` as a word boundary → `nx_domain`) + assert_eq!( + serde_json::to_string(&AdBlockResponse::ZeroAddress).unwrap(), + "\"zero_address\"" + ); + assert_eq!( + serde_json::to_string(&AdBlockResponse::NxDomain).unwrap(), + "\"nx_domain\"" + ); + } + + #[test] + fn test_ad_block_source_serde_skips_none_optionals() { + let source = AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: "StevenBlack".to_string(), + url: "https://example.com/list.txt".to_string(), + enabled: true, + response: AdBlockResponse::NxDomain, + last_fetched_at: None, + last_error: None, + rule_count: 0, + etag: None, + }; + let json = serde_json::to_string(&source).unwrap(); + assert!(!json.contains("last_fetched_at")); + assert!(!json.contains("last_error")); + assert!(!json.contains("etag")); + let restored: AdBlockSource = serde_json::from_str(&json).unwrap(); + assert_eq!(source, restored); + } + + #[test] + fn test_ad_block_source_serde_includes_some_optionals() { + let source = AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: "Test".to_string(), + url: "https://example.com/list.txt".to_string(), + enabled: false, + response: AdBlockResponse::ZeroAddress, + last_fetched_at: Some("2026-07-28T00:00:00Z".parse().unwrap()), + last_error: Some("timeout".to_string()), + rule_count: 42, + etag: Some("W/\"abc\"".to_string()), + }; + let json = serde_json::to_string(&source).unwrap(); + assert!(json.contains("last_fetched_at")); + assert!(json.contains("last_error")); + assert!(json.contains("etag")); + let restored: AdBlockSource = serde_json::from_str(&json).unwrap(); + assert_eq!(source, restored); + } + + #[test] + fn test_ad_block_state_default() { + let state = AdBlockState::default(); + assert!(!state.enabled); + assert!(state.sources.is_empty()); + assert!(state.whitelist.is_empty()); + assert!(state.auto_refresh_enabled); + assert_eq!(state.refresh_interval_hours, 24); + } + + #[test] + fn test_ad_block_state_serde_roundtrip() { + let state = AdBlockState { + enabled: true, + sources: vec![AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: "S1".to_string(), + url: "https://example.com/s1".to_string(), + enabled: true, + response: AdBlockResponse::NxDomain, + last_fetched_at: None, + last_error: None, + rule_count: 100, + etag: None, + }], + whitelist: vec!["trusted.example.com".to_string()], + auto_refresh_enabled: true, + refresh_interval_hours: 12, + }; + let json = serde_json::to_string(&state).unwrap(); + let restored: AdBlockState = serde_json::from_str(&json).unwrap(); + assert_eq!(state, restored); + } + // ----------------------------------------------------------------------- // ID type tests // ----------------------------------------------------------------------- diff --git a/src-tauri/crates/mhost-dns/src/adblock.rs b/src-tauri/crates/mhost-dns/src/adblock.rs new file mode 100644 index 0000000..6ba90de --- /dev/null +++ b/src-tauri/crates/mhost-dns/src/adblock.rs @@ -0,0 +1,450 @@ +//! Ad block engine for DNS mode (issue #130). +//! +//! Two sets of rules held independently so each source can choose its own +//! response strategy (0.0.0.0 A record vs NXDOMAIN): +//! +//! * `zero_addr` — sources configured with `AdBlockResponse::ZeroAddress`. +//! Hits return a 0.0.0.0 A record (`NoError` rcode, `A 0.0.0.0`). +//! * `nxdomain` — sources configured with `AdBlockResponse::NxDomain`. +//! Hits return `Rcode::NameError` so the client gives up immediately. +//! +//! Plus a whitelist: if a domain matches any whitelist entry (suffix-walked), +//! the ad block layer is bypassed entirely for that domain. Whitelist is +//! applied **before** the ad block engines — so whitelist wins over both +//! response variants. +//! +//! All three lookups use the shared [`crate::matcher::walk_parents`] helper +//! so `ad.example.com` matches a registered `example.com` (issue #79 fix). + +use parking_lot::RwLock; +use std::collections::{HashMap, HashSet}; +use std::net::IpAddr; +use std::sync::Arc; + +use crate::matcher::walk_parents; + +/// The action to take when an ad block rule matches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdBlockAction { + /// Return a 0.0.0.0 A record (with `NoError` rcode). + ZeroAddress(IpAddr), + /// Return `NXDOMAIN` (rcode = NameError). The `IpAddr` is unused; carried + /// as `()`-equivalent. Held in an `IpAddr` for symmetry with `ZeroAddress` + /// so both arms are `Copy + Eq`; downstream never reads the value. + NxDomain, +} + +/// An immutable snapshot of all ad-block rule sets, published as a single +/// `Arc` so a concurrent `check` either sees the old snapshot in full or the +/// new one in full — never a half-rebuilt mix (issue #132). +/// +/// This replaces the old three-`RwLock` + `AtomicUsize` short-circuit, which +/// updated three maps and a cached size in separate steps. On a 0→N rebuild +/// that ordering briefly let `check` short-circuit to `None` while the new +/// (loaded) maps were already in place, leaking ad-block hits through. A +/// single `Arc` swap removes the multi-step inconsistency entirely. +#[derive(Default)] +struct RulesSnapshot { + zero_addr: HashMap, + nxdomain: HashSet, + whitelist: HashSet, +} + +impl RulesSnapshot { + /// Whether this snapshot has any rules that could produce a block. + /// Whitelist is intentionally excluded: the master switch (`state.enabled`) + /// only gates zero_addr / nxdomain, while whitelist is always collected + /// regardless (review Medium #2). An empty `has_block_rules` means + /// `check()` can only return `None`, so callers can short-circuit the + /// parent-walk entirely. + #[inline] + fn has_block_rules(&self) -> bool { + !self.zero_addr.is_empty() || !self.nxdomain.is_empty() + } + + /// Total rule count for stats (`rule_count()`). Includes whitelist + /// because it's the externally observable number; short-circuit + /// decisions use [`has_block_rules`] instead. + fn total(&self) -> usize { + self.zero_addr.len() + self.nxdomain.len() + self.whitelist.len() + } +} + +/// DNS-mode ad block engine. Thread-safe; holds two rule sets + a whitelist. +/// +/// Hot-reload pattern: `rebuild(...)` builds a fresh [`RulesSnapshot`] and +/// swaps the whole `Arc` in under a single write lock — one atomic +/// publication step, so every concurrent `check` observes a coherent +/// snapshot (issue #132). `check` takes the read lock only long enough to +/// clone the `Arc` (refcount bump), then walks the immutable snapshot +/// lock-free. +pub struct AdBlockEngine { + current: RwLock>, +} + +impl AdBlockEngine { + pub fn new() -> Self { + Self { + current: RwLock::new(Arc::new(RulesSnapshot::default())), + } + } + + /// Atomically swap in new rule sets. + /// + /// Builds the three sets into one [`RulesSnapshot`] and replaces the + /// published `Arc` under a single write lock. This is the single point + /// of publication — there is no multi-step inconsistency window, so the + /// 0→N leak that the old `AtomicUsize` short-circuit had is gone + /// (issue #132). + /// + /// `zero_addr_rules` and `nxdomain_rules` come from parsing the cached + /// blocklist of each enabled source (per-source response strategy). + /// `whitelist` is the user-curated allow-list. + pub fn rebuild( + &self, + zero_addr_rules: HashMap, + nxdomain_rules: HashSet, + whitelist: HashSet, + ) { + let snapshot = Arc::new(RulesSnapshot { + zero_addr: zero_addr_rules, + nxdomain: nxdomain_rules, + whitelist, + }); + // Swap the Arc under one write lock, then drop the old snapshot + // OUTSIDE the lock. The old snapshot can hold 100k+ entries; letting + // its refcount hit zero and deallocate under the write lock would + // block every concurrent `check()` reader (review Medium #1). + let old = { + let mut g = self.current.write(); + std::mem::replace(&mut *g, snapshot) + }; + drop(old); + } + + /// Read the currently published snapshot. Takes the read lock only for + /// the `Arc::clone` (cheap — refcount bump), then releases it before any + /// domain walking, so concurrent rebuilds never block readers. + fn snapshot(&self) -> Arc { + Arc::clone(&self.current.read()) + } + + /// Decide what to do with a query. + /// + /// Returns `None` if the domain is whitelisted (fall through to the + /// regular rule engine / upstream) or not blocked at all. + pub fn check(&self, domain: &str) -> Option { + let snap = self.snapshot(); + // Fast-path: no block rules loaded → no possible hit. Avoids any + // domain walking for the common `state.enabled == false` case. + // Whitelist is excluded because it's collected regardless of the + // master switch (review Medium #2); an empty block-rule set means + // `check()` can only return `None`. Unlike the old `AtomicUsize` + // short-circuit this reads the very snapshot the walk below uses, + // so the empty-check can't disagree with the rule data (issue #132). + if !snap.has_block_rules() { + return None; + } + + // 1. whitelist (read once, then release) + if walk_parents(domain, |d| snap.whitelist.contains(d).then_some(())).is_some() { + return None; + } + // 2. NXDOMAIN sources first — more aggressive, save a hashmap lookup + if walk_parents(domain, |d| snap.nxdomain.contains(d).then_some(())).is_some() { + return Some(AdBlockAction::NxDomain); + } + // 3. zero-address sources + if let Some(ip) = walk_parents(domain, |d| snap.zero_addr.get(d).copied()) { + return Some(AdBlockAction::ZeroAddress(ip)); + } + None + } + + /// Total number of rules loaded across both action sets (whitelist + /// included — it influences `check` outcomes). + pub fn rule_count(&self) -> usize { + self.snapshot().total() + } + + pub fn whitelist_size(&self) -> usize { + self.snapshot().whitelist.len() + } + + pub fn zero_addr_count(&self) -> usize { + self.snapshot().zero_addr.len() + } + + pub fn nxdomain_count(&self) -> usize { + self.snapshot().nxdomain.len() + } +} + +impl Default for AdBlockEngine { + fn default() -> Self { + Self::new() + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use std::net::Ipv4Addr; + + fn za(domains: &[&str]) -> HashMap { + domains + .iter() + .map(|d| ((*d).to_string(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)))) + .collect() + } + + fn nx(domains: &[&str]) -> HashSet { + domains.iter().map(|d| (*d).to_string()).collect() + } + + fn wl(domains: &[&str]) -> HashSet { + domains.iter().map(|d| (*d).to_string()).collect() + } + + #[test] + fn empty_engine_returns_none() { + let engine = AdBlockEngine::new(); + assert_eq!(engine.check("anything.com"), None); + assert_eq!(engine.rule_count(), 0); + } + + /// PR #131 review finding 1.2: the fast-path short-circuit must bypass + /// the rule walking when no rules are loaded. With the `Arc::swap` + /// publication (issue #132) the empty-check reads the same snapshot the + /// walk would use, so we keep a dedicated assertion that an empty engine + /// returns `None` without touching rule data. + #[test] + fn empty_engine_short_circuits_before_locking() { + let engine = AdBlockEngine::new(); + // rule_count() now reads the published snapshot's total; on a fresh + // engine it's 0. We can't directly observe "no walk ran" from a + // public test, but a single-shot regression here locks in the + // contract that an empty engine returns None without state. + assert_eq!(engine.rule_count(), 0); + assert_eq!(engine.check("a.b.c.example.com"), None); + } + + #[test] + fn zero_address_hit_returns_zero_address() { + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["ad.example.com"]), nx(&[]), wl(&[])); + let action = engine.check("ad.example.com"); + assert_eq!( + action, + Some(AdBlockAction::ZeroAddress(IpAddr::V4(Ipv4Addr::new( + 0, 0, 0, 0 + )))) + ); + } + + #[test] + fn nxdomain_hit_returns_nxdomain() { + let engine = AdBlockEngine::new(); + engine.rebuild(za(&[]), nx(&["tracker.example.com"]), wl(&[])); + assert_eq!( + engine.check("tracker.example.com"), + Some(AdBlockAction::NxDomain) + ); + } + + #[test] + fn suffix_walk_matches_subdomains() { + // ad-blocker semantics: registering example.com hits *.example.com + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["example.com"]), nx(&[]), wl(&[])); + for d in ["example.com", "ad.example.com", "deep.ad.example.com"] { + assert!( + matches!(engine.check(d), Some(AdBlockAction::ZeroAddress(_))), + "{} should hit", + d + ); + } + assert_eq!(engine.check("example.org"), None); + } + + #[test] + fn nxdomain_consulted_before_zero_addr() { + // Per design (issue #130): the lookup order is whitelist → nxdomain + // → zero_addr. An NXDOMAIN rule on a parent domain blocks descendants + // before the more-specific zero_addr rule is reached. This is the + // intended Pi-hole-style semantic: NXDOMAIN is the more aggressive + // action and is consulted first to save a hashmap lookup. + let engine = AdBlockEngine::new(); + engine.rebuild( + za(&["specific.ad.example.com"]), + nx(&["example.com"]), + wl(&[]), + ); + // The parent NXDOMAIN wins because it's consulted first. + assert_eq!( + engine.check("specific.ad.example.com"), + Some(AdBlockAction::NxDomain) + ); + // But on a domain NOT covered by the parent nxdomain rule, the + // more-specific zero_addr rule still applies. + assert_eq!( + engine.check("specific.ad.other.com"), + None, + "other.com not under the nxdomain rule, and not registered in zero_addr" + ); + } + + #[test] + fn whitelist_overrides_everything() { + let engine = AdBlockEngine::new(); + // blocked on both engines; whitelisted → falls through. + engine.rebuild( + za(&["example.com"]), + nx(&["example.com"]), + wl(&["good.example.com"]), + ); + // whitelist exact hit + assert_eq!(engine.check("good.example.com"), None); + // whitelist suffix hit + assert_eq!(engine.check("api.good.example.com"), None); + // not whitelisted — still blocked + assert!(engine.check("ad.example.com").is_some()); + } + + #[test] + fn rebuild_replaces_state_atomically() { + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["a.com"]), nx(&["b.com"]), wl(&["c.com"])); + assert_eq!(engine.zero_addr_count(), 1); + assert_eq!(engine.nxdomain_count(), 1); + assert_eq!(engine.whitelist_size(), 1); + + engine.rebuild(za(&["d.com", "e.com"]), nx(&[]), wl(&[])); + assert_eq!(engine.zero_addr_count(), 2); + assert_eq!(engine.nxdomain_count(), 0); + assert_eq!(engine.whitelist_size(), 0); + // Old rules no longer hit + assert_eq!(engine.check("a.com"), None); + assert_eq!(engine.check("b.com"), None); + // New rule does hit + assert!(matches!( + engine.check("d.com"), + Some(AdBlockAction::ZeroAddress(_)) + )); + } + + #[test] + fn rule_count_sums_both_sets() { + let engine = AdBlockEngine::new(); + engine.rebuild( + za(&["a.com", "b.com"]), + nx(&["c.com", "d.com", "e.com"]), + wl(&[]), + ); + assert_eq!(engine.rule_count(), 5); + } + + /// PR #131 review finding 1.2: `rule_count` reflects the sum of + /// zero_addr, nxdomain and whitelist entries. Whitelist counts toward + /// this externally observable stat even though it no longer gates the + /// hot-path short-circuit — that's [`RulesSnapshot::has_block_rules`] + /// (PR #135 review, Medium #2). + #[test] + fn rule_count_includes_whitelist() { + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["a.com"]), nx(&[]), wl(&["w1", "w2", "w3"])); + assert_eq!(engine.rule_count(), 4); + } + + /// PR #135 review, Medium #2: a whitelist alone must not defeat the fast + /// path. `classify_rules` collects the whitelist regardless of the master + /// switch, so "ad block off + user has whitelist entries" would otherwise + /// make every DNS query walk parents for a result that can only ever be + /// `None`. + /// + /// This asserts on [`RulesSnapshot::has_block_rules`] **directly** on + /// purpose. The short-circuit is a pure optimisation — `check` returns + /// `None` either way — so a test that only drives the public API passes + /// against the un-fixed predicate too and guards nothing. + #[test] + fn whitelist_only_snapshot_has_no_block_rules() { + let whitelist_only = RulesSnapshot { + zero_addr: za(&[]), + nxdomain: nx(&[]), + whitelist: wl(&["trusted.com", "safe.com"]), + }; + assert!( + !whitelist_only.has_block_rules(), + "whitelist alone must not arm the hot path" + ); + assert_eq!( + whitelist_only.total(), + 2, + "...but it still counts toward the externally visible stat" + ); + + // Either block-rule set alone is enough to arm it. + for snap in [ + RulesSnapshot { + zero_addr: za(&["a.com"]), + nxdomain: nx(&[]), + whitelist: wl(&[]), + }, + RulesSnapshot { + zero_addr: za(&[]), + nxdomain: nx(&["b.com"]), + whitelist: wl(&[]), + }, + ] { + assert!(snap.has_block_rules()); + } + + // End-to-end tie-in: behaviour is unchanged by the optimisation. + let engine = AdBlockEngine::new(); + engine.rebuild(za(&[]), nx(&[]), wl(&["trusted.com"])); + assert_eq!(engine.check("trusted.com"), None); + assert_eq!(engine.rule_count(), 1); + } + + /// Rebuild publishes a fresh snapshot atomically, so the count reflects + /// the new total immediately — no cached-size window to fall out of sync + /// with the rule data (issue #132). + #[test] + fn rebuild_updates_cached_total() { + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["a.com"]), nx(&[]), wl(&[])); + assert_eq!(engine.rule_count(), 1); + + engine.rebuild(za(&[]), nx(&[]), wl(&[])); + assert_eq!( + engine.rule_count(), + 0, + "size must drop to 0 after empty rebuild" + ); + assert_eq!(engine.check("a.com"), None, "fast-path should now fire"); + } + + #[test] + fn tld_alone_matches_every_subdomain() { + // Pi-hole semantic: registering "com" blocks every *.com because + // walk_parents visits single-label parents once. This is intentional + // — users sometimes deliberately TLD-block (e.g. blocking the entire + // `.xyz` TLD used by abuse). + let engine = AdBlockEngine::new(); + engine.rebuild(za(&["com"]), nx(&[]), wl(&[])); + assert!(matches!( + engine.check("example.com"), + Some(AdBlockAction::ZeroAddress(_)) + )); + assert!(matches!( + engine.check("anything.anything.com"), + Some(AdBlockAction::ZeroAddress(_)) + )); + // A different TLD is untouched. + assert_eq!(engine.check("example.org"), None); + } +} diff --git a/src-tauri/crates/mhost-dns/src/lib.rs b/src-tauri/crates/mhost-dns/src/lib.rs index f9102d1..10765e3 100644 --- a/src-tauri/crates/mhost-dns/src/lib.rs +++ b/src-tauri/crates/mhost-dns/src/lib.rs @@ -1,9 +1,12 @@ +pub mod adblock; pub mod config; +pub mod matcher; pub mod platform; pub mod proxy; pub mod resolver; pub mod server; +pub use adblock::{AdBlockAction, AdBlockEngine}; pub use config::DnsConfig; pub use platform::UpstreamTier; pub use resolver::RuleEngine; diff --git a/src-tauri/crates/mhost-dns/src/matcher.rs b/src-tauri/crates/mhost-dns/src/matcher.rs new file mode 100644 index 0000000..02ef5b2 --- /dev/null +++ b/src-tauri/crates/mhost-dns/src/matcher.rs @@ -0,0 +1,105 @@ +//! Shared suffix-walking helper used by both [`crate::resolver::RuleEngine`] +//! (issue #79) and [`crate::adblock::AdBlockEngine`] (issue #130). +//! +//! Behaviour: +//! +//! ```text +//! walk_parents("a.b.c.example.com", predicate) +//! checks "a.b.c.example.com" → "b.c.example.com" → "c.example.com" +//! → "example.com" → "com" → stops (no dot left) +//! ``` +//! +//! First match wins. **Single-label parents are visited once** — so a +//! caller that registers `"com"` will match every `.com` query (Pi-hole +//! semantic). The walk only terminates when the current label has no +//! `.` in it AND predicate did not yield a value. + +/// Walk parent domains of `domain`, applying `predicate` to each candidate +/// (including `domain` itself). Returns the first value `predicate` yields +/// via `Some`, or `None` if the walk exhausts without a hit. +/// +/// `domain` is treated as already lowercased / canonicalised by the caller. +pub(crate) fn walk_parents(domain: &str, predicate: F) -> Option +where + F: Fn(&str) -> Option, +{ + let mut current = domain; + loop { + if let Some(v) = predicate(current) { + return Some(v); + } + let pos = current.find('.')?; + current = ¤t[pos + 1..]; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn walk_parents_first_match_wins() { + // "hits" exact domain first, then parents + let r = walk_parents("a.b.example.com", |d| match d { + "a.b.example.com" => Some(1), + "b.example.com" => Some(2), + "example.com" => Some(3), + _ => None, + }); + assert_eq!(r, Some(1)); + } + + #[test] + fn walk_parents_falls_through_to_parent() { + let r = walk_parents("a.b.example.com", |d| match d { + "example.com" => Some("hit"), + _ => None, + }); + assert_eq!(r, Some("hit")); + } + + #[test] + fn walk_parents_visits_single_label_parent_once() { + // Pi-hole semantic: registering "com" should block every *.com. + // Walk visits "example.com" once, then "com" once, then stops. + let r = walk_parents("example.com", |d| match d { + "com" => Some("hit"), + _ => None, + }); + assert_eq!(r, Some("hit")); + } + + #[test] + fn walk_parents_no_hit() { + let r = walk_parents("a.b.example.com", |_| None::<()>); + assert_eq!(r, None); + } + + /// **PR #154 review (P3)**: Pi-hole-style TLD blocking semantics. + /// Registering `"com"` in the engine should match every query ending + /// in `.com` (single-label ancestor walk is intentional). Verifies + /// the full `walk_parents` chain, not just the single-step case. + #[test] + fn walk_parents_registered_tld_matches_every_subdomain() { + let r = walk_parents("deeply.nested.subdomain.example.com", |d| match d { + "com" => Some("TLD-hit"), + _ => None, + }); + assert_eq!(r, Some("TLD-hit")); + + // And also the trivial single-label form. + let r = walk_parents("example.com", |d| match d { + "com" => Some("TLD-hit"), + _ => None, + }); + assert_eq!(r, Some("TLD-hit")); + + // But a query NOT under .com should miss (proves the hit + // wasn't a false positive from walk mechanics). + let r = walk_parents("example.org", |d| match d { + "com" => Some("TLD-hit"), + _ => None, + }); + assert_eq!(r, None); + } +} diff --git a/src-tauri/crates/mhost-dns/src/proxy.rs b/src-tauri/crates/mhost-dns/src/proxy.rs index c64d703..7376ffa 100644 --- a/src-tauri/crates/mhost-dns/src/proxy.rs +++ b/src-tauri/crates/mhost-dns/src/proxy.rs @@ -504,6 +504,14 @@ pub(crate) mod tests { pub(crate) static TEST_LOCK: Mutex<()> = Mutex::new(()); /// 持锁 guard,测试结束时自动 drop。 + /// + /// **lint 抑制**:`await_holding_lock` lint 会触发,但我们**故意** + /// 让 guard 跨越 `.await` —— 这些测试共享文件系统(runtime_dir、 + /// signal file),用 Mutex 串行化保证它们不会和并行运行的其他测试 + /// 相互覆盖。`.drop(test_lock())` 会破坏这个序列化(已验证:drop + /// 后 `test_check_shutdown_signal` + `test_read_original_dns_from_file` + /// 并行跑会偶发失败)。 + #[allow(clippy::await_holding_lock)] pub(crate) fn test_lock() -> std::sync::MutexGuard<'static, ()> { TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()) } @@ -543,6 +551,7 @@ pub(crate) mod tests { } #[tokio::test] + #[allow(clippy::await_holding_lock)] async fn test_proxy_concurrent_clients() { // 关键测试:两个 client 并发,proxy 不能把 response 交叉 let query_a = b"QUERY_A".to_vec(); @@ -624,9 +633,16 @@ pub(crate) mod tests { // 简化版集成测试:spawn proxy,**不**写 file signal,proxy // 应该持续运行(不主动退出)。验证 poll 不会让 proxy 误退出。 // 完整 shutdown 行为用 dev 模式手动验证。 + // + // **lint 选择**:lock 只用于序列化 setup(清理 signal file + 写 + // tempdir + 启动 proxy)。后续的 `sleep(1500ms)` 不再触及共享 + // 文件系统状态,提前 drop 避免 `await_holding_lock`。 let _lock = test_lock(); let _tmp = set_test_runtime_dir(); let _ = std::fs::remove_file(crate::platform::shutdown_signal_file()); + // 释放 lock:setup 已完成(tempdir + signal file + port), + // 接下来的 `UdpSocket::bind().await` + spawn 不需要再串行化。 + drop(_lock); let listen_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let listen_port = listen_socket.local_addr().unwrap().port(); @@ -634,8 +650,6 @@ pub(crate) mod tests { let mut proxy = DnsProxy::new(listen_port, 1053); let _ = proxy.take_shutdown_sender(); let proxy_handle = tokio::spawn(async move { proxy.run().await }); - drop(_lock); - // 等 1.5s(覆盖至少 1 个 poll tick)。proxy 不应该退出。 tokio::time::sleep(Duration::from_millis(1500)).await; assert!( @@ -667,7 +681,7 @@ pub(crate) mod tests { tokio::spawn(async move { let mut buf = vec![0u8; 4096]; // 仅 recv 不 reply,让每个 query 等待 5s 超时 - while let Ok(_) = upstream_socket.recv_from(&mut buf).await {} + while upstream_socket.recv_from(&mut buf).await.is_ok() {} }); let listen_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); @@ -715,6 +729,7 @@ pub(crate) mod tests { /// `available_permits()` 验证首批 N 个 query 正好占满了所有 permit, /// 后续 query 被丢弃。 #[tokio::test] + #[allow(clippy::await_holding_lock)] async fn test_proxy_semaphore_blocks_excess_spawns() { let _lock = test_lock(); let _tmp = set_test_runtime_dir(); @@ -727,7 +742,7 @@ pub(crate) mod tests { let upstream_port = upstream_socket.local_addr().unwrap().port(); tokio::spawn(async move { let mut buf = vec![0u8; 4096]; - while let Ok(_) = upstream_socket.recv_from(&mut buf).await {} + while upstream_socket.recv_from(&mut buf).await.is_ok() {} }); let listen_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); diff --git a/src-tauri/crates/mhost-dns/src/server.rs b/src-tauri/crates/mhost-dns/src/server.rs index 2c717cf..fa3adb1 100644 --- a/src-tauri/crates/mhost-dns/src/server.rs +++ b/src-tauri/crates/mhost-dns/src/server.rs @@ -17,6 +17,7 @@ use tokio::task::JoinHandle; use crate::config::DnsConfig; use crate::resolver::RuleEngine; +use crate::AdBlockEngine; /// DNS 服务错误。 #[derive(Debug, thiserror::Error)] @@ -88,6 +89,11 @@ pub struct DnsServer { /// - 锁从 `std::sync::Mutex` 换 `parking_lot::Mutex`:parking_lot 比 std /// Mutex 在非竞争路径更快、poison-free(这里我们不需要处理 poison)。 cache: Arc, CacheEntry>>>, + /// Ad block engine(issue #130)。Lookup 在 `handle_address_query` + /// 第一步执行(在本地 `rule_engine.resolve` 之前),拦截广告域名 + /// 返回 0.0.0.0 / NXDOMAIN。共享 `Arc` 让 spawn 任务零锁读 + 命令 + /// 层通过 `reload_ad_block_rules` 在外面 hot-reload。 + ad_block_engine: Arc, } impl DnsServer { @@ -114,6 +120,7 @@ impl DnsServer { refresh_handle: Mutex::new(None), refresh_shutdown: Arc::new(tokio::sync::Notify::new()), cache: Arc::new(PlMutex::new(LruCache::new(cache_size))), + ad_block_engine: Arc::new(AdBlockEngine::new()), }) } @@ -171,6 +178,7 @@ impl DnsServer { // 也不阻塞 refresh task 的 write 锁。 let resolver_slot = Arc::clone(&self.resolver); let cache = self.cache.clone(); + let ad_block_engine = self.ad_block_engine.clone(); let handle = tokio::spawn(async move { let mut buf = vec![0u8; UDP_BUF_SIZE]; @@ -187,6 +195,7 @@ impl DnsServer { let response_data = match handle_dns_request( request_data, &rule_engine, + &ad_block_engine, &resolver, &cache, ).await { @@ -335,6 +344,38 @@ impl DnsServer { self.cache.lock().clear(); } + /// 重新加载广告屏蔽规则(issue #130)。 + /// + /// 与 `reload_rules` 同等语义:rebuild 引擎后清空 LRU 缓存, + /// 否则 reload 前向上游查过并缓存的域名仍会返回 stale upstream IP, + /// 覆盖新的 ad-block 命中(issue #132 follow-up)。 + pub fn reload_ad_block_rules( + &self, + zero_addr_rules: std::collections::HashMap, + nxdomain_rules: std::collections::HashSet, + whitelist: std::collections::HashSet, + ) { + self.ad_block_engine + .rebuild(zero_addr_rules, nxdomain_rules, whitelist); + self.cache.lock().clear(); + } + + /// 广告屏蔽规则数量(含 nxdomain + zero_addr,**不含** whitelist)。 + pub fn ad_block_rule_count(&self) -> usize { + self.ad_block_engine.rule_count() + } + + /// 白名单条目数量。 + pub fn ad_block_whitelist_size(&self) -> usize { + self.ad_block_engine.whitelist_size() + } + + /// 测试用:直接拿到 AdBlockEngine。 + #[doc(hidden)] + pub fn ad_block_engine_for_test(&self) -> Arc { + Arc::clone(&self.ad_block_engine) + } + /// 是否正在运行。 pub fn is_running(&self) -> bool { self.running.load(Ordering::SeqCst) @@ -386,6 +427,7 @@ impl DnsServer { async fn handle_dns_request( request_data: &[u8], rule_engine: &RuleEngine, + ad_block_engine: &AdBlockEngine, resolver: &Arc, cache: &Arc, CacheEntry>>>, ) -> Option> { @@ -465,7 +507,16 @@ async fn handle_dns_request( return build_notimp_response(&request, &query); } - match handle_address_query(name_str, query.name(), record_type, rule_engine, resolver).await { + match handle_address_query( + name_str, + query.name(), + record_type, + rule_engine, + ad_block_engine, + resolver, + ) + .await + { QueryResult::Answer(record) => { let ttl = record.ttl(); // **fix (P-R2, issue #90)**: `*record.clone()` 之前是 `clone Box + deref` @@ -483,6 +534,7 @@ async fn handle_dns_request( response_bytes } QueryResult::NoError => build_noerror_response(&request, &query), + QueryResult::NxDomain => build_nxdomain_response(&request, &query), QueryResult::ServFail => build_servfail_response(&request, &query), } } @@ -564,6 +616,25 @@ fn build_servfail_response(request: &Message, query: &Query) -> Option> } } +/// 构造 NXDOMAIN 响应(issue #130 ad block NxDomain 响应类型专用)。 +fn build_nxdomain_response(request: &Message, query: &Query) -> Option> { + let mut header = Header::response_from_request(request.header()); + header.set_authoritative(false); + header.set_recursion_available(true); + header.set_response_code(ResponseCode::NXDomain); + let mut response = Message::new(); + response.set_header(header); + response.set_id(request.id()); + response.add_query(query.clone()); + match response.to_bytes() { + Ok(bytes) => Some(bytes), + Err(e) => { + tracing::warn!("Failed to encode NXDOMAIN response: {}", e); + None + } + } +} + /// 构造 NotImp 响应(不支持的查询类型)。 fn build_notimp_response(request: &Message, query: &Query) -> Option> { let mut header = Header::response_from_request(request.header()); @@ -586,6 +657,8 @@ fn build_notimp_response(request: &Message, query: &Query) -> Option> { enum QueryResult { Answer(Box), NoError, + /// NXDOMAIN 响应(ad block NxDomain 响应类型专用)。 + NxDomain, ServFail, } @@ -595,8 +668,45 @@ async fn handle_address_query( name: &Name, qtype: RecordType, rule_engine: &RuleEngine, + ad_block_engine: &AdBlockEngine, resolver: &Arc, ) -> QueryResult { + // 0. 广告屏蔽(issue #130):在常规规则之前拦截。 + if let Some(action) = ad_block_engine.check(name_str) { + match action { + crate::adblock::AdBlockAction::NxDomain => return QueryResult::NxDomain, + crate::adblock::AdBlockAction::ZeroAddress(ip) => { + let record = match (qtype, ip) { + (RecordType::A, IpAddr::V4(v4)) => Some(Record::from_rdata( + name.clone(), + LOCAL_RULE_TTL, + RData::A(A(v4)), + )), + (RecordType::AAAA, IpAddr::V6(v6)) => { + use hickory_proto::rr::rdata::AAAA; + Some(Record::from_rdata( + name.clone(), + LOCAL_RULE_TTL, + RData::AAAA(AAAA(v6)), + )) + } + // qtype 与规则 IP family 不匹配(如 AAAA 命中 IPv4 规则): + // 返回 NxDomain 而非 NoError。NoError + 空答案在 RFC 2308 + // 语义里是「name 存在但无此 type 记录」,等于放行让 + // 上游继续解析 —— 把广告屏蔽的意图完全绕开了。 + // (review feedback:PR #154 P1) + _ => return QueryResult::NxDomain, + }; + return match record { + Some(r) => QueryResult::Answer(Box::new(r)), + // 同样不可能到这里(family mismatch 已返回 NxDomain), + // 但保留分支以防御未来新增的 family 类型。 + None => QueryResult::NxDomain, + }; + } + } + } + // 1. 优先匹配本地规则 if let Some(ip) = rule_engine.resolve(name_str) { let record = match (qtype, ip) { @@ -1552,7 +1662,7 @@ mod tests { // 所以两次都是 cache miss(首次)。但本地规则对 AAAA 不匹配 → // 返回 NoError。如果 cache key 没区分 type,第二次 A 查询可能 // 拿到 AAAA 的 NoError 响应(错误)。 - async fn query_once(server: &DnsServer, port: u16, qtype: RecordType) -> Vec { + async fn query_once(_server: &DnsServer, port: u16, qtype: RecordType) -> Vec { let query_name = Name::from_utf8("typed.example.com.").unwrap(); let query = Query::query(query_name, qtype); let mut request = Message::new(); diff --git a/src-tauri/crates/mhost-hosts/src/parser.rs b/src-tauri/crates/mhost-hosts/src/parser.rs index 8256de5..5b0092b 100644 --- a/src-tauri/crates/mhost-hosts/src/parser.rs +++ b/src-tauri/crates/mhost-hosts/src/parser.rs @@ -446,7 +446,7 @@ mod tests { #[test] fn test_extract_managed_block_bytes_crlf() { let input = "# hdr\r\n# ---- mHost start ----\r\n127.0.0.1 x.com\r\n# ---- mHost end ----\r\n# tail\r\n"; - let (start, end) = Parser::extract_managed_block_bytes(input).expect("block exists"); + let (_start, end) = Parser::extract_managed_block_bytes(input).expect("block exists"); // byte_end must point just past the \r\n after end marker assert_eq!(input.as_bytes()[end - 2], b'\r'); diff --git a/src-tauri/crates/mhost-storage/src/adblock.rs b/src-tauri/crates/mhost-storage/src/adblock.rs new file mode 100644 index 0000000..971ea67 --- /dev/null +++ b/src-tauri/crates/mhost-storage/src/adblock.rs @@ -0,0 +1,458 @@ +//! Persistence layer for the DNS-mode ad block subsystem (issue #130). +//! +//! Layout under the storage root: +//! +//! ```text +//! {root}/ +//! adblock.json # AdBlockState (single small JSON document) +//! adblock-cache/{id}.txt # raw fetched blocklist per source +//! ``` +//! +//! All writes use `atomic_write` (defined in `storage.rs`) so a crash during +//! write never leaves a half-written config or cache file. Reads return +//! `AdBlockState::default()` when the JSON is missing — ad block is opt-in +//! so first run is fine with no file at all. + +use std::fmt; +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use chrono::Utc; +use mhost_core::{AdBlockSource, AdBlockState, SourceId}; + +use super::storage::atomic_write; + +// --------------------------------------------------------------------------- +// Paths +// --------------------------------------------------------------------------- + +const STATE_FILE: &str = "adblock.json"; +const CACHE_DIR: &str = "adblock-cache"; + +// --------------------------------------------------------------------------- +// State (adblock.json) +// --------------------------------------------------------------------------- + +/// Read the persisted ad block state. Returns `AdBlockState::default()` if +/// the file does not exist (first run, ad block never enabled). +pub fn read_state(root: &Path) -> io::Result { + let path = root.join(STATE_FILE); + match fs::read_to_string(&path) { + Ok(s) => serde_json::from_str(&s).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("adblock.json is corrupted: {}", e), + ) + }), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(AdBlockState::default()), + Err(e) => Err(e), + } +} + +/// Read adblock state with a safety net against silent data loss. +/// +/// If the file is missing, returns `AdBlockState::default()` (first run). +/// If the file is corrupted (invalid JSON), the file is **renamed** to +/// `adblock.json.corrupt-{YYYYMMDDhhmmssuuuuuu}` so the user/support can +/// recover their previous whitelist + source list, and the default state +/// is returned. Returning `AdBlockState::default()` unconditionally would +/// rewrite the corrupt file with an empty state on the next save, silently +/// wiping the user's configuration (PR #131 review, finding 0.2). +pub fn read_state_or_default_with_backup(root: &Path) -> AdBlockState { + let path = root.join(STATE_FILE); + let raw = match fs::read_to_string(&path) { + Ok(s) => s, + Err(e) if e.kind() == io::ErrorKind::NotFound => return AdBlockState::default(), + Err(e) => { + eprintln!("[mHost] adblock.json unreadable: {}; using empty state", e); + return AdBlockState::default(); + } + }; + match serde_json::from_str::(&raw) { + Ok(s) => s, + Err(parse_err) => { + // PR #131 self-review §3: avoid clobbering an existing backup. + // Microsecond-precision timestamps + two corruptions in the same + // microsecond (rare but observed in tight test loops) would + // otherwise replace the prior backup atomically (POSIX `rename`) + // or fail outright (Windows), losing the previous corruption's + // bytes. Pick the first non-existent filename with a counter. + let stamp: BackupStamp = BackupStamp(Utc::now()); + let mut backup = root.join(format!("adblock.json.corrupt-{}", stamp)); + let mut counter: u32 = 1; + while backup.exists() { + backup = root.join(format!("adblock.json.corrupt-{}-{}", stamp, counter)); + counter += 1; + // belt-and-suspenders: if we somehow spin without progress + // (read-only filesystem?), bail rather than spin forever. + if counter > 1024 { + eprintln!( + "[mHost] adblock.json corrupted ({}); could not find free backup \ + name after 1024 attempts. Skipping backup, falling back to empty state.", + parse_err + ); + return AdBlockState::default(); + } + } + match fs::rename(&path, &backup) { + Ok(_) => eprintln!( + "[mHost] adblock.json corrupted: {}. Backed up to {}; \ + falling back to empty state. Your whitelist and sources were lost — \ + re-add them via the Ad Block page (or restore from the backup file).", + parse_err, + backup.display() + ), + Err(rename_err) => eprintln!( + "[mHost] adblock.json corrupted ({}); backup rename also failed ({}). \ + Next save will overwrite — falling back to empty state.", + parse_err, rename_err + ), + } + AdBlockState::default() + } + } +} + +struct BackupStamp(chrono::DateTime); +impl fmt::Display for BackupStamp { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Microsecond precision so two corruptions in the same second + // still get distinct filenames. + write!(f, "{}", self.0.format("%Y%m%d%H%M%S%6f")) + } +} + +/// Atomically write the ad block state. +pub fn write_state(root: &Path, state: &AdBlockState) -> io::Result<()> { + let path = root.join(STATE_FILE); + let json = serde_json::to_string_pretty(state) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + atomic_write(&path, json.as_bytes()) +} + +// --------------------------------------------------------------------------- +// Source cache (adblock-cache/{id}.txt) +// --------------------------------------------------------------------------- + +fn cache_dir(root: &Path) -> PathBuf { + root.join(CACHE_DIR) +} + +fn cache_path(root: &Path, source_id: &SourceId) -> PathBuf { + let id_str = source_id.to_string(); + // SourceId wraps a UUID; the rendered form is hex + dashes, so neither + // `/` nor `\` can appear. The assert is belt-and-suspenders against any + // future identifier type that might allow path-significant characters. + debug_assert!( + !id_str.contains('/') && !id_str.contains('\\'), + "SourceId rendered as `{}` — must not contain path separators", + id_str + ); + cache_dir(root).join(format!("{}.txt", id_str)) +} + +/// Write the raw fetched blocklist content for a source. Atomic — partial +/// writes never leave a torn cache file. +pub fn write_cache(root: &Path, source_id: &SourceId, content: &[u8]) -> io::Result<()> { + fs::create_dir_all(cache_dir(root))?; + let path = cache_path(root, source_id); + atomic_write(&path, content) +} + +/// Read the raw blocklist content for a source. Returns `None` if no cache +/// has ever been written for this source. +pub fn read_cache(root: &Path, source_id: &SourceId) -> io::Result> { + let path = cache_path(root, source_id); + match fs::read_to_string(&path) { + Ok(s) => Ok(Some(s)), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None), + Err(e) => Err(e), + } +} + +/// Remove the cache file for a source. Idempotent — a missing file is OK. +pub fn delete_cache(root: &Path, source_id: &SourceId) -> io::Result<()> { + let path = cache_path(root, source_id); + match fs::remove_file(&path) { + Ok(()) => Ok(()), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e), + } +} + +// --------------------------------------------------------------------------- +// Bulk delete helpers (used when a source is removed) +// --------------------------------------------------------------------------- + +/// Remove both the source's cache file and ensure the source's id is no +/// longer present in `state.sources`. Caller is responsible for `write_state` +/// after this call. +pub fn purge_source(root: &Path, state: &mut AdBlockState, source_id: &SourceId) -> io::Result<()> { + state.sources.retain(|s| &s.source_id != source_id); + delete_cache(root, source_id) +} + +// --------------------------------------------------------------------------- +// Convenience: list sources from state (re-exported to avoid churn) +// --------------------------------------------------------------------------- + +/// Find a source by id. Linear scan; expected N is small (single digits). +pub fn find_source<'a>(state: &'a AdBlockState, id: &SourceId) -> Option<&'a AdBlockSource> { + state.sources.iter().find(|s| &s.source_id == id) +} + +/// Mutable equivalent of [`find_source`]. +pub fn find_source_mut<'a>( + state: &'a mut AdBlockState, + id: &SourceId, +) -> Option<&'a mut AdBlockSource> { + state.sources.iter_mut().find(|s| &s.source_id == id) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use mhost_core::AdBlockResponse; + use tempfile::TempDir; + + fn sample_source(name: &str) -> AdBlockSource { + AdBlockSource { + source_id: SourceId(uuid::Uuid::new_v4()), + name: name.to_string(), + url: format!("https://example.com/{}.txt", name), + enabled: true, + response: AdBlockResponse::ZeroAddress, + last_fetched_at: None, + last_error: None, + rule_count: 0, + etag: None, + } + } + + #[test] + fn read_state_returns_default_when_missing() { + let temp = TempDir::new().unwrap(); + let state = read_state(temp.path()).unwrap(); + assert_eq!(state, AdBlockState::default()); + } + + #[test] + fn write_then_read_state_roundtrip() { + let temp = TempDir::new().unwrap(); + let mut state = AdBlockState { + enabled: true, + ..AdBlockState::default() + }; + state.sources.push(sample_source("a")); + state.whitelist.push("trusted.example.com".to_string()); + + write_state(temp.path(), &state).unwrap(); + let restored = read_state(temp.path()).unwrap(); + assert_eq!(state, restored); + } + + #[test] + fn write_state_is_atomic_no_tmp_files_leaked() { + let temp = TempDir::new().unwrap(); + let state = AdBlockState::default(); + write_state(temp.path(), &state).unwrap(); + // atomic_write uses tempfile::NamedTempFile which cleans up on drop; + // assert no stray .tmp file remains at root. + let stray: Vec<_> = fs::read_dir(temp.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|e| { + e.path() + .extension() + .map(|ext| ext == "tmp") + .unwrap_or(false) + }) + .collect(); + assert!(stray.is_empty(), "found stray tmp files: {:?}", stray); + } + + #[test] + fn cache_write_read_delete_roundtrip() { + let temp = TempDir::new().unwrap(); + let id = SourceId(uuid::Uuid::new_v4()); + + // Missing → None + assert!(read_cache(temp.path(), &id).unwrap().is_none()); + + // Write + read + write_cache(temp.path(), &id, b"0.0.0.0 ad.example\n").unwrap(); + let content = read_cache(temp.path(), &id).unwrap().unwrap(); + assert_eq!(content, "0.0.0.0 ad.example\n"); + + // Delete + delete_cache(temp.path(), &id).unwrap(); + assert!(read_cache(temp.path(), &id).unwrap().is_none()); + + // Delete again is idempotent + delete_cache(temp.path(), &id).unwrap(); + } + + #[test] + fn purge_source_removes_cache_and_listing() { + let temp = TempDir::new().unwrap(); + let mut state = AdBlockState::default(); + let s1 = sample_source("keep"); + let s2 = sample_source("drop"); + let keep_id = s1.source_id.clone(); + let drop_id = s2.source_id.clone(); + state.sources.push(s1); + state.sources.push(s2); + write_cache(temp.path(), &keep_id, b"keep").unwrap(); + write_cache(temp.path(), &drop_id, b"drop").unwrap(); + + purge_source(temp.path(), &mut state, &drop_id).unwrap(); + + assert!(read_cache(temp.path(), &drop_id).unwrap().is_none()); + assert_eq!(state.sources.len(), 1); + assert_eq!(state.sources[0].source_id, keep_id); + // Keep cache intact + assert!(read_cache(temp.path(), &keep_id).unwrap().is_some()); + } + + #[test] + fn read_state_corrupted_json_returns_error() { + let temp = TempDir::new().unwrap(); + fs::write(temp.path().join(STATE_FILE), b"{not valid json").unwrap(); + let err = read_state(temp.path()).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn write_cache_creates_cache_dir() { + let temp = TempDir::new().unwrap(); + let id = SourceId(uuid::Uuid::new_v4()); + write_cache(temp.path(), &id, b"x").unwrap(); + assert!(cache_dir(temp.path()).is_dir()); + } + + #[test] + fn read_state_or_default_missing_returns_default() { + let temp = TempDir::new().unwrap(); + let state = read_state_or_default_with_backup(temp.path()); + assert_eq!(state, AdBlockState::default()); + } + + /// PR #131 review finding 0.2: a corrupted adblock.json must not be + /// silently thrown away. `read_state_or_default_with_backup` renames + /// the bad file aside and returns the default state, so the user can + /// recover manually if needed. + #[test] + fn read_state_or_default_corrupt_backs_up_file_and_returns_default() { + let temp = TempDir::new().unwrap(); + // Seed a deliberately broken file. Also inject real content so + // the user can tell what they lost. + let original = b"{not valid json"; + fs::write(temp.path().join(STATE_FILE), original).unwrap(); + + let recovered = read_state_or_default_with_backup(temp.path()); + assert_eq!(recovered, AdBlockState::default()); + + // The corrupted file must no longer be at the canonical path — + // otherwise the next save() would silently overwrite it. + let canonical = temp.path().join(STATE_FILE); + assert!( + !canonical.exists(), + "corrupt adblock.json should have been renamed away" + ); + + // ...and renamed to a `corrupt-{timestamp}` file (we don't pin the + // timestamp, just confirm at least one exists and carries the + // original bytes). + let backups: Vec<_> = fs::read_dir(temp.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|e| { + e.file_name() + .to_string_lossy() + .starts_with("adblock.json.corrupt-") + }) + .collect(); + assert_eq!(backups.len(), 1, "expected exactly one backup file"); + let backup_contents = fs::read(backups[0].path()).unwrap(); + assert_eq!(backup_contents, original, "backup preserves original bytes"); + } + + #[test] + fn read_state_or_default_valid_unchanged() { + let temp = TempDir::new().unwrap(); + let mut state = AdBlockState { + enabled: true, + ..AdBlockState::default() + }; + state.whitelist.push("a.com".to_string()); + write_state(temp.path(), &state).unwrap(); + + let restored = read_state_or_default_with_backup(temp.path()); + assert_eq!(restored, state); + // Sanity: no backup files created on a successful read. + let stray: Vec<_> = fs::read_dir(temp.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|e| { + e.file_name() + .to_string_lossy() + .starts_with("adblock.json.corrupt-") + }) + .collect(); + assert!(stray.is_empty(), "valid file must not be backed up"); + } + + /// PR #131 self-review §3: even when two corruptions happen in the same + /// microsecond, the second backup must not clobber the first. The + /// `find_unique_backup_name` helper picks a counter-suffixed name when + /// the timestamped target already exists. + #[test] + fn read_state_or_default_collision_counter_kicks_in() { + let temp = TempDir::new().unwrap(); + // Pre-seed a file at the exact name our function would pick on the + // next call: we don't know the exact timestamp, but we don't have + // to — the loop in the helper keeps incrementing until it finds a + // free name. Seeding a *single* matching file is impossible without + // reproducing the helper's timestamp format; instead, force the + // collision deterministically by first creating any `corrupt-*` + // file in the directory, then asserting the function's chosen + // backup name is distinct. + // + // Cheaper: directly construct two distinct backups via two + // back-to-back reads. With microsecond timestamps, two reads in + // succession can produce the same stamp; if the helper works the + // backup set ends up with two files (not one). + let original = b"{not valid json"; + fs::write(temp.path().join(STATE_FILE), original).unwrap(); + // Replace the file before the second read (the first read renamed + // it aside as a backup). + let _ = read_state_or_default_with_backup(temp.path()); + // First read should have backed up the file; restore the corrupt + // canonical so a second read can also back up (different timestamp). + fs::write(temp.path().join(STATE_FILE), original).unwrap(); + let _ = read_state_or_default_with_backup(temp.path()); + + let backups: Vec<_> = fs::read_dir(temp.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|e| { + e.file_name() + .to_string_lossy() + .starts_with("adblock.json.corrupt-") + }) + .collect(); + assert!( + !backups.is_empty(), + "two corruptions must produce at least one backup file (got {})", + backups.len() + ); + // The interesting case (2 backups with same timestamp) is rare and + // timing-dependent; we accept either 1 or 2 backup files here — + // correctness is by inspection of the counter loop. + } +} diff --git a/src-tauri/crates/mhost-storage/src/lib.rs b/src-tauri/crates/mhost-storage/src/lib.rs index 82a846f..fd0a391 100644 --- a/src-tauri/crates/mhost-storage/src/lib.rs +++ b/src-tauri/crates/mhost-storage/src/lib.rs @@ -1,3 +1,4 @@ +pub mod adblock; pub mod manifest; pub mod migration; pub mod storage; diff --git a/src-tauri/crates/mhost-storage/src/storage.rs b/src-tauri/crates/mhost-storage/src/storage.rs index 5828922..708fc3a 100644 --- a/src-tauri/crates/mhost-storage/src/storage.rs +++ b/src-tauri/crates/mhost-storage/src/storage.rs @@ -331,7 +331,7 @@ impl Storage for FileStorage { /// /// 使用 `NamedTempFile` 避免并发写入时的固定临时文件名竞态条件。 /// 如果写入过程中发生错误,临时文件会自动清理。 -fn atomic_write(path: &Path, content: &[u8]) -> io::Result<()> { +pub(crate) fn atomic_write(path: &Path, content: &[u8]) -> io::Result<()> { let parent = path.parent().unwrap_or_else(|| Path::new(".")); let mut temp_file = tempfile::NamedTempFile::new_in(parent)?; temp_file.write_all(content)?; diff --git a/src-tauri/src/commands/adblock.rs b/src-tauri/src/commands/adblock.rs new file mode 100644 index 0000000..4d45f44 --- /dev/null +++ b/src-tauri/src/commands/adblock.rs @@ -0,0 +1,1308 @@ +//! DNS-mode ad block IPC commands (issue #130). +//! +//! 12 commands: state CRUD, source management, refresh control, whitelist. +//! Storage layout is defined in [`mhost_storage::adblock`]. The +//! in-memory `state.ad_block_state` is the source of truth for hot-reload; +//! changes go through [`persist_and_reload`] which keeps file + memory + +//! `DnsServer.ad_block_engine` in sync atomically. + +use std::collections::{HashMap, HashSet}; +use std::net::IpAddr; +use std::sync::atomic::Ordering; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use chrono::Utc; +use mhost_core::{AdBlockResponse, AdBlockSource, AdBlockState, MhostError, SourceId}; +use mhost_hosts::Parser; +use mhost_storage::adblock as adblock_store; +use tauri::State; +use uuid::Uuid; + +use crate::state::{lock_or_recover, AppState}; + +// --------------------------------------------------------------------------- +// Constants +// --------------------------------------------------------------------------- + +/// Hard upper bound on rules per source. The classic anti-pattern (issue #130 +/// background) was writing 100k+ hosts entries into `/etc/hosts`; we keep the +/// same ceiling at the parser layer so a misconfigured source can't OOM the +/// process. +const MAX_RULES_PER_SOURCE: usize = 100_000; + +/// HTTP fetch timeout. Blocklist refresh shouldn't block the UI thread; +/// 30 s is generous for typical hosts-format payloads. +const FETCH_TIMEOUT_SECS: u64 = 30; + +/// Maximum response body size for an ad-block source (PR #131 review +/// finding 1.8 — a malicious or unbounded-misconfigured source can still +/// run for `FETCH_TIMEOUT_SECS` and start streaming bytes; cap the bytes). +/// `MAX_RULES_PER_SOURCE × ~30 bytes ≈ 3 MB`; 16 MB headroom is enough for +/// legitimate (well-annotated) lists. +const MAX_RESPONSE_BYTES: usize = 16 * 1024 * 1024; + +/// Maximum length of a source URL. URLs longer than this are rejected +/// to prevent IPC-level memory abuse (PR #154 review P3). +const MAX_URL_LEN: usize = 2048; + +/// Maximum length of a whitelist domain entry (RFC 1035 §3.1: each +/// label ≤ 63 chars, full domain ≤ 253 chars). PR #154 review P3. +const MAX_DOMAIN_LEN: usize = 253; + +/// Concurrency cap for `refresh_all_ad_block_sources` and the periodic +/// background refresh (PR #131 review finding 1.4 — refresh was a serial +/// loop, blocking the UI for up to N × FETCH_TIMEOUT_SECS). +pub(crate) const REFRESH_CONCURRENCY: usize = 4; + +const USER_AGENT: &str = "mHost-Desktop/1.0"; + +// --------------------------------------------------------------------------- +// Shared HTTP client (PR #131 review findings 1.8 + 1.9) +// --------------------------------------------------------------------------- + +/// Process-wide shared `reqwest::Client`. Building a client is non-trivial +/// (TLS keylog, DNS resolver, connection pool) and we were doing it on every +/// fetch + every background refresh tick. Reusing one client also means +/// HTTP keep-alive across fetches and a bounded connection pool. +/// +/// `OnceLock::get_or_init` runs the closure synchronously on the first +/// call; all subsequent calls return the same `&'static` handle. Safe +/// because `reqwest::Client::build()` is sync. +static HTTP_CLIENT: OnceLock = OnceLock::new(); + +fn http_client() -> &'static reqwest::Client { + HTTP_CLIENT.get_or_init(|| { + reqwest::Client::builder() + .user_agent(USER_AGENT) + .timeout(Duration::from_secs(FETCH_TIMEOUT_SECS)) + // reqwest 0.13 lacks `body_limit`; we enforce MAX_RESPONSE_BYTES + // explicitly via `fetch_source_content` (early reject via the + // Content-Length header + post-read size check). + .build() + .expect("reqwest client build must succeed with static config") + }) +} + +/// Fetch `url` via the shared client. Returns the raw body bytes plus the +/// `ETag` header (if any), and rejects anything larger than +/// `MAX_RESPONSE_BYTES`. Two-stage guard: +/// +/// 1. Server-advertised `Content-Length` → reject without downloading. +/// 2. Post-read size check → catches servers that lie about length. +/// +/// PR #131 review finding 1.8 — a malicious or misconfigured source can +/// otherwise stream arbitrary bytes for `FETCH_TIMEOUT_SECS` before our +/// parser sees them. +async fn fetch_source(url: &str) -> Result<(Vec, Option), MhostError> { + let resp = http_client() + .get(url) + .send() + .await + .map_err(|e| MhostError::Network(format!("network error: {}", e)))?; + if !resp.status().is_success() { + return Err(MhostError::ExternalApi(format!( + "fetch {} failed: HTTP {}", + url, + resp.status() + ))); + } + let etag = resp + .headers() + .get(reqwest::header::ETAG) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + if let Some(len) = resp.content_length() { + if len > MAX_RESPONSE_BYTES as u64 { + return Err(MhostError::InvalidInput(format!( + "source body length {} exceeds limit {}", + len, MAX_RESPONSE_BYTES + ))); + } + } + let body = resp + .bytes() + .await + .map_err(|e| MhostError::Network(format!("read body error: {}", e)))?; + if body.len() > MAX_RESPONSE_BYTES { + return Err(MhostError::InvalidInput(format!( + "source body received {} bytes, exceeds limit {}", + body.len(), + MAX_RESPONSE_BYTES + ))); + } + Ok((body.to_vec(), etag)) +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/// Persist the in-memory state to disk and hot-reload the running DNS server's +/// ad block engine. Used by every state-mutating command so the on-disk file, +/// in-memory copy, and resolver engine never diverge. +/// +/// Must be called from a tokio context (uses `.await`). Acquires the state +/// write lock briefly to clone out, then releases before touching DNS server +/// to keep lock-hold time minimal. +/// +/// **Issue #138:** the `spawn_blocking` closure below self-checks +/// the cancel token immediately before calling `reload_ad_block_rules`. +/// This protects against the race where the disable path runs +/// mid-`classify_rules` (which is sync, not cancellable): the closure +/// finishes classifying, sees the token is set, and bails before +/// mutating a `DnsServer` that's already been stopped. `write_state` is +/// intentionally NOT gated on the token — persisting in-memory state to +/// disk is the safe thing to do regardless of DNS-mode state. +pub(crate) async fn persist_and_reload(state: &AppState) -> Result<(), MhostError> { + // Clone out under the lock, then drop the guard before DNS work. + let snapshot: AdBlockState = { + let guard = state.ad_block_state.read().await; + guard.clone() + }; + + // Wrap write_state + classify_rules + reload in a single + // spawn_blocking so none of the sync file IO or parsing blocks a + // tokio worker thread (issue #133 — parsing 100k+ domain blocklists + // on the reload path starved concurrent DNS queries). + let root = state.storage.root().to_path_buf(); + let dns_enabled = state.dns_enabled.load(Ordering::Relaxed); + let dns_server = Arc::clone(&state.dns_server); + // Clone the token out from its Mutex slot before crossing the + // spawn_blocking boundary. Issue #138 follow-up: the field is a + // `Mutex` (not a bare token) so the refresh + // task can swap in a fresh, uncancelled token on every spawn — + // persist_and_reload reads whatever is currently in the slot, + // which is the token bound to the latest spawned task. + let cancel = lock_or_recover(&state.ad_block_refresh_cancel).clone(); + tokio::task::spawn_blocking(move || -> Result<(), MhostError> { + adblock_store::write_state(&root, &snapshot) + .map_err(|e| MhostError::InvalidInput(format!("write_state: {}", e)))?; + if dns_enabled && !cancel.is_cancelled() { + // We deliberately run classify_rules even if the pre-check + // just succeeded: it's the long sync step (100k+ domain + // parsing) and is exactly where cancel is most likely to + // land. The post-classify check below is the only + // authoritative one for the reload decision; the pre-check + // exists only to skip the work entirely when we know up + // front that we'll bail. + let (zero_addr, nxdomain, whitelist) = classify_rules(&snapshot, &root); + // Re-check after classify_rules: it's the long sync step and + // is exactly where cancel is most likely to have landed. + // (See issue #138: spawn_blocking cannot be aborted.) + if !cancel.is_cancelled() { + if let Some(server) = lock_or_recover(&dns_server).as_ref() { + server.reload_ad_block_rules(zero_addr, nxdomain, whitelist); + } + } + } + Ok(()) + }) + .await + .map_err(|e| MhostError::InvalidInput(format!("persist task failed: {}", e)))? +} + +/// Reduce `AdBlockState` into the three rule sets consumed by the engine. +/// Reads each enabled source's cache file synchronously — only invoked from +/// `persist_and_reload`, which is in turn called from a tokio task; the IO +/// is fast (small files, no parsing needed here). +pub(crate) fn classify_rules( + state: &AdBlockState, + root: &std::path::Path, +) -> (HashMap, HashSet, HashSet) { + let mut zero_addr: HashMap = HashMap::new(); + let mut nxdomain: HashSet = HashSet::new(); + + if state.enabled { + // 仅 master switch 开启时才下发规则到引擎;关闭时引擎收到空集, + // 自然 fallback 到原始规则 / 上游。 + for source in &state.sources { + if !source.enabled { + continue; + } + let domains = domains_for_source(root, source); + match source.response { + AdBlockResponse::ZeroAddress => { + let ip = IpAddr::from([0, 0, 0, 0]); + for d in domains { + zero_addr.entry(d).or_insert(ip); + } + } + AdBlockResponse::NxDomain => { + for d in domains { + nxdomain.insert(d); + } + } + } + } + } + + let whitelist: HashSet = state.whitelist.iter().cloned().collect(); + + (zero_addr, nxdomain, whitelist) +} + +/// Load cached parsed domains for a single source. Returns an empty Vec if +/// the cache file is missing or fails to parse (caller logs and continues). +pub(crate) fn domains_for_source(root: &std::path::Path, source: &AdBlockSource) -> Vec { + match adblock_store::read_cache(root, &source.source_id) { + Ok(Some(content)) => parse_blocklist_domains(&content), + Ok(None) => Vec::new(), + Err(e) => { + eprintln!( + "[adblock] failed to read cache for source {}: {}", + source.name, e + ); + Vec::new() + } + } +} + +/// Validate a whitelist entry. Returns the canonical form (trimmed + +/// lowercased) on success, or an error message describing why the +/// input is invalid. +/// +/// **PR #154 review (P2):** the original code only checked for empty +/// input, so entries like `*.example.com`, `example.com/path`, or +/// `not a domain at all` were persisted silently and never matched in +/// `walk_parents` (it does literal `HashSet::contains`). They also +/// didn't surface in `last_error`, so the user had no signal that the +/// entry was broken. +/// +/// Rules enforced: +/// - non-empty after trim +/// - no whitespace anywhere (`*.example.com` etc. → reject) +/// - no path separator (`example.com/path` → reject) +/// - no leading dot (`.example.com` — engines don't expect this; the +/// suffix-walk covers the case anyway) +/// - no wildcard chars `*` (suffix-walk handles hierarchical match) +/// - only ASCII letters / digits / `-` / `.` +fn validate_whitelist_domain(raw: &str) -> Result { + let trimmed = raw.trim().to_lowercase(); + if trimmed.is_empty() { + return Err("whitelist entry is empty".to_string()); + } + if trimmed.len() > MAX_DOMAIN_LEN { + return Err(format!( + "whitelist entry length {} exceeds limit {}", + trimmed.len(), + MAX_DOMAIN_LEN + )); + } + if trimmed.contains(char::is_whitespace) { + return Err(format!("whitelist entry contains whitespace: {:?}", raw)); + } + if trimmed.contains('/') || trimmed.contains('\\') { + return Err(format!("whitelist entry looks like a URL/path: {:?}", raw)); + } + if trimmed.starts_with('.') { + return Err(format!( + "whitelist entry must not start with '.': {:?}", + raw + )); + } + if trimmed.contains('*') { + return Err(format!( + "whitelist entry must not contain '*' (suffix-walk matches subdomains): {:?}", + raw + )); + } + for ch in trimmed.chars() { + if !(ch.is_ascii_alphanumeric() || ch == '-' || ch == '.') { + return Err(format!( + "whitelist entry has invalid character {:?}: {:?}", + ch, raw + )); + } + } + Ok(trimmed) +} + +/// Parse hosts-format blocklist content into a flat list of domains. +/// Comments (`#`) and empty lines are filtered out by `Parser::parse_line`. +/// +/// **PR #154 review (P2)**: no-op — after analysis, the original +/// `d.to_lowercase()` is correct and the only allocation we can avoid +/// here is for already-lowercase strings (the common case for +/// well-formed blocklists). The `eq_ignore_ascii_case` / +/// `to_ascii_uppercase` shortcut doesn't actually save allocations +/// (`to_ascii_uppercase` allocates a String) and breaks the +/// `MiXed.ExAmPlE.com → mixed.example.com` semantic that the +/// `parse_blocklist_lowercases` test relies on. Sticking with the +/// straightforward `to_lowercase()` — the work runs in +/// `spawn_blocking` (PR #131 P1-2 + issue #133), so DNS queries +/// aren't blocked during the parse. +fn parse_blocklist_domains(content: &str) -> Vec { + let result = Parser::parse(content); + let mut domains: Vec = Vec::new(); + for rule in result.rules { + if !rule.enabled { + continue; + } + for d in rule.domains { + domains.push(d.to_lowercase()); + } + } + domains +} + +/// Fetch a remote blocklist over HTTP(S), validate, and persist the raw +/// content + parsed-domain count back into `state.sources[i]`. The hot-reload +/// is the caller's responsibility (use `persist_and_reload` after). +pub(crate) async fn fetch_and_cache_source( + state: &AppState, + source_id: &SourceId, +) -> Result<(), MhostError> { + // 1. Read the source record under the read lock. + let source_clone = { + let guard = state.ad_block_state.read().await; + match adblock_store::find_source(&guard, source_id) { + Some(s) => s.clone(), + None => { + return Err(MhostError::InvalidInput(format!( + "ad block source not found: {}", + source_id + ))) + } + } + }; + + // 2. Fetch raw bytes via the shared reqwest client (PR #131 review + // finding 1.9). The static client is built once; size enforcement + // happens inside `fetch_source` (PR #131 review finding 1.8). + // + // PR #131 re-review P1-2: record a fetch failure on the source's + // `last_error` before propagating — the parse-failure branch below + // already did this, but a network/size failure returned via `?` with no + // record, so the UI badge and the persisted state both stayed stale. + let url = source_clone.url.clone(); + let body_and_etag = fetch_source(&url).await; + let (body, etag) = match body_and_etag { + Ok(v) => v, + Err(e) => { + let msg = e.to_string(); + let _ = record_fetch_error(state, source_id, &msg).await; + return Err(e); + } + }; + let content_str = std::str::from_utf8(&body) + .map_err(|e| MhostError::InvalidInput(format!("response is not valid UTF-8: {}", e)))?; + + // 3. Parse + enforce hard limit. Use spawn_blocking because the parser + // is sync and the input can be large. + let content_owned = content_str.to_string(); + let root = state.storage.root().to_path_buf(); + let id_owned = source_id.clone(); + let parse_result: Result<(usize, Vec), MhostError> = + tokio::task::spawn_blocking(move || { + let domains = parse_blocklist_domains(&content_owned); + if domains.len() > MAX_RULES_PER_SOURCE { + return Err(MhostError::InvalidInput(format!( + "source produced {} rules (limit: {})", + domains.len(), + MAX_RULES_PER_SOURCE + ))); + } + // Re-serialize as canonical hosts text so the cache is always + // valid hosts format (drops comments the original may have). + let canon = domains + .iter() + .map(|d| format!("0.0.0.0 {}", d)) + .collect::>() + .join("\n"); + adblock_store::write_cache(&root, &id_owned, canon.as_bytes())?; + Ok((domains.len(), Vec::new())) + }) + .await + .map_err(|e| MhostError::InvalidInput(format!("parse task failed: {}", e)))?; + + let rule_count = match parse_result { + Ok((count, _)) => count, + Err(e) => { + // Persist the failure on the source so the UI can show it, + // but keep the previous cache intact for DNS to keep working. + record_fetch_error(state, source_id, &e.to_string()).await?; + return Err(e); + } + }; + + // 4. Update source record: clear error, set fetched_at, rule_count, etag. + { + let mut guard = state.ad_block_state.write().await; + if let Some(s) = adblock_store::find_source_mut(&mut guard, source_id) { + s.last_error = None; + s.last_fetched_at = Some(Utc::now()); + s.rule_count = rule_count; + s.etag = etag; + } + } + Ok(()) +} + +/// Persist an error string onto a source's `last_error` field. Does NOT +/// touch `last_fetched_at` or `rule_count` — those reflect the last +/// successful fetch and should be preserved on failure. +pub(crate) async fn record_fetch_error( + state: &AppState, + source_id: &SourceId, + err: &str, +) -> Result<(), MhostError> { + record_fetch_error_internal(&state.ad_block_state, source_id, err).await +} + +/// `AppState`-free variant for the background refresh task (which clones +/// just the Arcs it needs at spawn time). +pub(crate) async fn record_fetch_error_internal( + ad_block_state: &Arc>, + source_id: &SourceId, + err: &str, +) -> Result<(), MhostError> { + let mut guard = ad_block_state.write().await; + if let Some(s) = adblock_store::find_source_mut(&mut guard, source_id) { + s.last_error = Some(err.to_string()); + } + Ok(()) +} + +/// `AppState`-free variant of `fetch_and_cache_source` for the background +/// refresh task. Skips the proxy detection of "is DNS still on" — that's +/// checked at the call site in `dns.rs` before invoking this. +pub(crate) async fn fetch_and_cache_source_internal( + storage: &Arc, + ad_block_state: &Arc>, + source_id: &SourceId, +) -> Result<(), MhostError> { + let source_clone = { + let guard = ad_block_state.read().await; + match adblock_store::find_source(&guard, source_id) { + Some(s) => s.clone(), + None => { + return Err(MhostError::InvalidInput(format!( + "ad block source not found: {}", + source_id + ))) + } + } + }; + + let url = source_clone.url.clone(); + let (body, etag) = fetch_source(&url).await?; + let content_str = std::str::from_utf8(&body) + .map_err(|e| MhostError::InvalidInput(format!("response is not valid UTF-8: {}", e)))?; + + let content_owned = content_str.to_string(); + let root = storage.root().to_path_buf(); + let id_owned = source_id.clone(); + let parse_result: Result = tokio::task::spawn_blocking(move || { + let domains = parse_blocklist_domains(&content_owned); + if domains.len() > MAX_RULES_PER_SOURCE { + return Err(MhostError::InvalidInput(format!( + "source produced {} rules (limit: {})", + domains.len(), + MAX_RULES_PER_SOURCE + ))); + } + let canon = domains + .iter() + .map(|d| format!("0.0.0.0 {}", d)) + .collect::>() + .join("\n"); + adblock_store::write_cache(&root, &id_owned, canon.as_bytes())?; + Ok(domains.len()) + }) + .await + .map_err(|e| MhostError::InvalidInput(format!("parse task failed: {}", e)))?; + + let rule_count = match parse_result { + Ok(n) => n, + Err(e) => { + record_fetch_error_internal(ad_block_state, source_id, &e.to_string()).await?; + return Err(e); + } + }; + + { + let mut guard = ad_block_state.write().await; + if let Some(s) = adblock_store::find_source_mut(&mut guard, source_id) { + s.last_error = None; + s.last_fetched_at = Some(Utc::now()); + s.rule_count = rule_count; + s.etag = etag; + } + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// State +// --------------------------------------------------------------------------- + +/// Return the full ad block state (sources + whitelist + meta). +#[tauri::command] +pub async fn get_ad_block_state(state: State<'_, AppState>) -> Result { + Ok(state.ad_block_state.read().await.clone()) +} + +/// Master switch. Disabling also clears the engine's rule sets via +/// `persist_and_reload` (which classifies with `enabled=false` → empty). +#[tauri::command] +pub async fn set_ad_block_enabled( + enabled: bool, + state: State<'_, AppState>, +) -> Result<(), MhostError> { + { + let mut guard = state.ad_block_state.write().await; + guard.enabled = enabled; + } + persist_and_reload(&state).await +} + +/// Change the auto-refresh interval in hours. `0` disables background +/// refresh (frontend shows a hint to refresh manually). +#[tauri::command] +pub async fn set_ad_block_refresh_interval( + hours: u32, + state: State<'_, AppState>, +) -> Result<(), MhostError> { + // 软上限:1h .. 7d。低于 1h 太频繁伤上游;超过 7d 几乎失去"自动"意义。 + let clamped = hours.clamp(0, 24 * 7); + { + let mut guard = state.ad_block_state.write().await; + guard.refresh_interval_hours = clamped; + } + persist_and_reload(&state).await +} + +// --------------------------------------------------------------------------- +// Source management +// --------------------------------------------------------------------------- + +#[tauri::command] +pub async fn list_ad_block_sources( + state: State<'_, AppState>, +) -> Result, MhostError> { + Ok(state.ad_block_state.read().await.sources.clone()) +} + +/// Add a new source, fetch it immediately, then persist. Returns the source +/// record (with `last_fetched_at`, `rule_count`, possibly `last_error`). +#[tauri::command] +pub async fn add_ad_block_source( + name: String, + url: String, + response: AdBlockResponse, + state: State<'_, AppState>, +) -> Result { + add_ad_block_source_impl(&state, name, url, response).await +} + +/// `AppState`-by-ref impl so the persistence-on-fetch-failure contract +/// (PR #131 re-review P1-2) can be unit-tested without a Tauri `State`. +pub(crate) async fn add_ad_block_source_impl( + state: &AppState, + name: String, + url: String, + response: AdBlockResponse, +) -> Result { + if name.trim().is_empty() { + return Err(MhostError::InvalidInput("source name is empty".into())); + } + if url.len() > MAX_URL_LEN { + return Err(MhostError::InvalidInput(format!( + "source url length {} exceeds limit {}", + url.len(), + MAX_URL_LEN + ))); + } + if !(url.starts_with("http://") || url.starts_with("https://")) { + return Err(MhostError::InvalidInput(format!( + "url must be http(s); got '{}'", + url + ))); + } + + let new_source = AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name, + url, + enabled: true, + response, + last_fetched_at: None, + last_error: None, + rule_count: 0, + etag: None, + }; + let new_id = new_source.source_id.clone(); + + { + let mut guard = state.ad_block_state.write().await; + guard.sources.push(new_source.clone()); + } + + // Fetch + propagate errors. PR #131 review finding 1.5: the previous + // `let _ = …` discarded the error, leaving the UI with a successful + // toast and a "fetch failed" badge next to a brand-new source (the + // UX was confusing). Surface the error to the frontend toast; the + // source is still in `state.sources` with `last_error` populated so + // a later "Refresh" works as expected. + // + // PR #131 re-review P1-2: `?` here skipped `persist_and_reload` on + // fetch failure, so the source existed only in memory and was lost + // on restart. Persist unconditionally first (capturing `last_error` + // too), then surface the fetch error to the toast. + let fetch_result = fetch_and_cache_source(state, &new_id).await; + persist_and_reload(state).await?; + fetch_result?; + + // Return the freshly-fetched source record to the UI. + let snap = state.ad_block_state.read().await; + let stored = adblock_store::find_source(&snap, &new_id) + .cloned() + .unwrap_or(new_source); + drop(snap); + Ok(stored) +} + +#[tauri::command] +pub async fn remove_ad_block_source( + source_id: SourceId, + state: State<'_, AppState>, +) -> Result<(), MhostError> { + let root = state.storage.root().to_path_buf(); + { + let mut guard = state.ad_block_state.write().await; + adblock_store::purge_source(&root, &mut guard, &source_id) + .map_err(|e| MhostError::InvalidInput(format!("purge_source: {}", e)))?; + } + persist_and_reload(&state).await +} + +#[tauri::command] +pub async fn set_ad_block_source_enabled( + source_id: SourceId, + enabled: bool, + state: State<'_, AppState>, +) -> Result { + { + let mut guard = state.ad_block_state.write().await; + let s = adblock_store::find_source_mut(&mut guard, &source_id) + .ok_or_else(|| MhostError::InvalidInput(format!("source not found: {}", source_id)))?; + s.enabled = enabled; + } + persist_and_reload(&state).await?; + let snap = state.ad_block_state.read().await; + Ok(adblock_store::find_source(&snap, &source_id) + .cloned() + .expect("source just updated")) +} + +#[tauri::command] +pub async fn set_ad_block_source_response( + source_id: SourceId, + response: AdBlockResponse, + state: State<'_, AppState>, +) -> Result { + { + let mut guard = state.ad_block_state.write().await; + let s = adblock_store::find_source_mut(&mut guard, &source_id) + .ok_or_else(|| MhostError::InvalidInput(format!("source not found: {}", source_id)))?; + s.response = response; + } + persist_and_reload(&state).await?; + let snap = state.ad_block_state.read().await; + Ok(adblock_store::find_source(&snap, &source_id) + .cloned() + .expect("source just updated")) +} + +// --------------------------------------------------------------------------- +// Refresh (concurrent) +// --------------------------------------------------------------------------- + +/// Fan out `fetch_and_cache_source_internal` over `source_ids` with bounded +/// concurrency. Per-source errors are recorded on `last_error` (preserving +/// the existing semantics) so the caller doesn't have to propagate. +/// +/// PR #131 review finding 1.4: the previous serial loop could block the UI +/// for `N × FETCH_TIMEOUT_SECS` while `N` sources sequentially hit the +/// network. With this helper, a typical 4-source list finishes in ~one +/// timeout instead of four, and `isLoadingAtom` no longer lingers. +pub(crate) async fn fetch_sources_concurrent( + storage: &Arc, + ad_block_state: &Arc>, + source_ids: &[SourceId], + concurrency: usize, +) { + if source_ids.is_empty() { + return; + } + let sem = Arc::new(tokio::sync::Semaphore::new(concurrency.max(1))); + let mut handles = Vec::with_capacity(source_ids.len()); + for id in source_ids { + let permit = Arc::clone(&sem) + .acquire_owned() + .await + .expect("semaphore starts with positive permits and is never closed"); + let storage = storage.clone(); + let ad_block_state = ad_block_state.clone(); + let id = id.clone(); + handles.push(tokio::spawn(async move { + // Permit drops at end of task → slot released regardless of + // success/failure. + let _permit = permit; + if let Err(e) = fetch_and_cache_source_internal(&storage, &ad_block_state, &id).await { + let _ = record_fetch_error_internal(&ad_block_state, &id, &e.to_string()).await; + eprintln!("[adblock] concurrent refresh source {} failed: {}", id, e); + } + })); + } + // Drain in submission order so a slow source doesn't keep its permit + // forever if the user disables / deletes it mid-flight. We don't use + // the results; the spawn task already did the write. + for h in handles { + let _ = h.await; + } +} + +// --------------------------------------------------------------------------- +// Refresh +// --------------------------------------------------------------------------- + +#[tauri::command] +pub async fn refresh_ad_block_source( + source_id: SourceId, + state: State<'_, AppState>, +) -> Result { + // PR #131 re-review P1-2 (same pattern as add_ad_block_source): persist + // unconditionally so `last_error` is captured on disk, then surface the + // fetch error. The source already exists on disk here, so this is about + // not losing the error state rather than not losing the source. + let fetch_result = fetch_and_cache_source(&state, &source_id).await; + persist_and_reload(&state).await?; + fetch_result?; + let snap = state.ad_block_state.read().await; + Ok(adblock_store::find_source(&snap, &source_id) + .cloned() + .expect("source just fetched")) +} + +#[tauri::command] +pub async fn refresh_all_ad_block_sources( + state: State<'_, AppState>, +) -> Result, MhostError> { + // Snapshot IDs up-front to avoid holding the lock across await. + let ids: Vec = { + let snap = state.ad_block_state.read().await; + snap.sources + .iter() + .filter(|s| s.enabled) + .map(|s| s.source_id.clone()) + .collect() + }; + // Concurrent fetch — bounded at REFRESH_CONCURRENCY. Per-source + // failures are recorded on `last_error` via the helper. + fetch_sources_concurrent( + &state.storage, + &state.ad_block_state, + &ids, + REFRESH_CONCURRENCY, + ) + .await; + persist_and_reload(&state).await?; + Ok(state.ad_block_state.read().await.sources.clone()) +} + +// --------------------------------------------------------------------------- +// Whitelist +// --------------------------------------------------------------------------- + +#[tauri::command] +pub async fn list_ad_block_whitelist( + state: State<'_, AppState>, +) -> Result, MhostError> { + Ok(state.ad_block_state.read().await.whitelist.clone()) +} + +#[tauri::command] +pub async fn add_ad_block_whitelist( + domain: String, + state: State<'_, AppState>, +) -> Result, MhostError> { + let normalized = validate_whitelist_domain(&domain).map_err(MhostError::InvalidInput)?; + { + let mut guard = state.ad_block_state.write().await; + if !guard.whitelist.contains(&normalized) { + guard.whitelist.push(normalized); + } + } + persist_and_reload(&state).await?; + Ok(state.ad_block_state.read().await.whitelist.clone()) +} + +#[tauri::command] +pub async fn remove_ad_block_whitelist( + domain: String, + state: State<'_, AppState>, +) -> Result, MhostError> { + // Removal tolerates the same input the user typed when adding (i.e. + // no validation — silently no-op on missing). This matches the + // contract of "remove what matches; ignore the rest". + let normalized = domain.trim().to_lowercase(); + { + let mut guard = state.ad_block_state.write().await; + guard.whitelist.retain(|d| d != &normalized); + } + persist_and_reload(&state).await?; + Ok(state.ad_block_state.read().await.whitelist.clone()) +} + +// --------------------------------------------------------------------------- +// Unit tests (helpers only — IPC commands themselves covered by +// `commands/integration_tests.rs`-style tests in a follow-up). +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classify_rules_disabled_master_yields_empty() { + let temp = tempfile::TempDir::new().unwrap(); + let mut state = AdBlockState { + enabled: false, + ..Default::default() + }; + state.sources.push(AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: "s".into(), + url: "https://x".into(), + enabled: true, + response: AdBlockResponse::ZeroAddress, + last_fetched_at: None, + last_error: None, + rule_count: 1, + etag: None, + }); + let (z, n, w) = classify_rules(&state, temp.path()); + assert!(z.is_empty()); + assert!(n.is_empty()); + assert!(w.is_empty()); + } + + #[test] + fn classify_rules_partitions_by_response() { + let temp = tempfile::TempDir::new().unwrap(); + let mk = |name: &str, response: AdBlockResponse, enabled: bool| AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: name.into(), + url: "https://x".into(), + enabled, + response, + last_fetched_at: None, + last_error: None, + rule_count: 0, + etag: None, + }; + let za_source = mk("za", AdBlockResponse::ZeroAddress, true); + let nx_source = mk("nx", AdBlockResponse::NxDomain, true); + let off_source = mk("off", AdBlockResponse::ZeroAddress, false); + + // Seed cache files so the zero_addr / nxdomain partitions are + // non-empty (issue #134 — previously only `w.len()==1` was + // asserted, leaving the partition logic untested). + mhost_storage::adblock::write_cache( + temp.path(), + &za_source.source_id, + b"0.0.0.0 ads.example.com\n0.0.0.0 tracker.example.com\n", + ) + .unwrap(); + mhost_storage::adblock::write_cache( + temp.path(), + &nx_source.source_id, + b"0.0.0.0 blocked.example.com\n", + ) + .unwrap(); + // The disabled source also has a cache file — its domains must + // NOT appear in any partition (enabled=false short-circuits it). + mhost_storage::adblock::write_cache( + temp.path(), + &off_source.source_id, + b"0.0.0.0 should-not-appear.com\n", + ) + .unwrap(); + + let state = AdBlockState { + enabled: true, + sources: vec![za_source, nx_source, off_source], + whitelist: vec!["trusted.com".to_string()], + ..Default::default() + }; + let (z, n, w) = classify_rules(&state, temp.path()); + + // zero_addr partition: domains from the enabled ZeroAddress source, + // mapped to 0.0.0.0. + assert_eq!(z.len(), 2, "zero_addr seeded from enabled za source"); + assert!(z.contains_key("ads.example.com")); + assert!(z.contains_key("tracker.example.com")); + assert_eq!( + z.get("ads.example.com").copied(), + Some(IpAddr::from([0u8, 0, 0, 0])), + "ZeroAddress domains must map to 0.0.0.0" + ); + + // nxdomain partition: domains from the enabled NxDomain source. + assert_eq!(n.len(), 1, "nxdomain seeded from enabled nx source"); + assert!(n.contains("blocked.example.com")); + + // whitelist partition. + assert_eq!(w.len(), 1); + assert!(w.contains("trusted.com")); + + // The disabled source's domain must not leak into any partition. + assert!( + !z.contains_key("should-not-appear.com"), + "disabled source must not contribute to zero_addr" + ); + assert!( + !n.contains("should-not-appear.com"), + "disabled source must not contribute to nxdomain" + ); + } + + #[test] + fn parse_blocklist_extracts_domains() { + let text = "\ +# ad-block test +0.0.0.0 ad.example.com +0.0.0.0 tracker.example.com +127.0.0.1 also.example.com + +# comment +"; + let domains = parse_blocklist_domains(text); + assert!(domains.contains(&"ad.example.com".to_string())); + assert!(domains.contains(&"tracker.example.com".to_string())); + assert!(domains.contains(&"also.example.com".to_string())); + // comments and blanks are filtered by the parser + } + + #[test] + fn parse_blocklist_lowercases() { + let text = "0.0.0.0 MiXed.ExAmPlE.com\n"; + let domains = parse_blocklist_domains(text); + assert_eq!(domains, vec!["mixed.example.com".to_string()]); + } + + // ----------------------------------------------------------------- + // PR #154 review (P2): validate_whitelist_domain test coverage. + // Each rejection branch + the happy path + the MAX_DOMAIN_LEN + // guard. These are pure sync tests — no AppState / DnsServer + // needed. + // ----------------------------------------------------------------- + + #[test] + fn validate_whitelist_domain_happy_path_lowercases_and_trims() { + assert_eq!( + validate_whitelist_domain(" Example.COM ").unwrap(), + "example.com" + ); + assert_eq!( + validate_whitelist_domain("foo.example.com").unwrap(), + "foo.example.com" + ); + } + + #[test] + fn validate_whitelist_domain_rejects_empty_or_whitespace_only() { + assert!(validate_whitelist_domain("").is_err()); + assert!(validate_whitelist_domain(" ").is_err()); + let err = validate_whitelist_domain("").unwrap_err(); + assert!(err.contains("empty"), "unexpected error: {}", err); + } + + #[test] + fn validate_whitelist_domain_rejects_whitespace_inside() { + assert!(validate_whitelist_domain("foo bar.com").is_err()); + assert!(validate_whitelist_domain("foo\tbar.com").is_err()); + } + + #[test] + fn validate_whitelist_domain_rejects_path_separator() { + assert!(validate_whitelist_domain("example.com/path").is_err()); + assert!(validate_whitelist_domain("example.com\\path").is_err()); + let err = validate_whitelist_domain("example.com/path").unwrap_err(); + assert!(err.contains("URL/path"), "unexpected error: {}", err); + } + + #[test] + fn validate_whitelist_domain_rejects_wildcard() { + assert!(validate_whitelist_domain("*.example.com").is_err()); + let err = validate_whitelist_domain("*.example.com").unwrap_err(); + assert!(err.contains("*"), "unexpected error: {}", err); + } + + #[test] + fn validate_whitelist_domain_rejects_leading_dot() { + assert!(validate_whitelist_domain(".example.com").is_err()); + } + + #[test] + fn validate_whitelist_domain_rejects_unicode() { + assert!(validate_whitelist_domain("例え.com").is_err()); + assert!(validate_whitelist_domain("café.example.com").is_err()); + } + + #[test] + fn validate_whitelist_domain_rejects_oversize() { + // 254 chars — exceeds RFC 1035 max of 253. + let huge = "a".repeat(254); + assert!(validate_whitelist_domain(&huge).is_err()); + let err = validate_whitelist_domain(&huge).unwrap_err(); + assert!(err.contains("exceeds limit"), "unexpected error: {}", err); + // 253 chars — exactly at the boundary, should pass. + let at_limit = "a".repeat(253); + assert!(validate_whitelist_domain(&at_limit).is_ok()); + } + + // ----------------------------------------------------------------- + // PR #131 re-review P1-1: the cold-start fix in `set_dns_mode_enable` + // and `AppState::new` relies on `classify_rules` turning a source's + // cached blocklist into non-empty rule sets, then `reload_ad_block_rules` + // populating the engine. This locks that building block so a refactor + // can't silently empty the engine on DNS enable. + // ----------------------------------------------------------------- + #[test] + fn classify_rules_populates_from_cached_source() { + let temp = tempfile::TempDir::new().unwrap(); + let mk = |name: &str, response: AdBlockResponse| AdBlockSource { + source_id: SourceId(Uuid::new_v4()), + name: name.into(), + url: "https://x".into(), + enabled: true, + response, + last_fetched_at: None, + last_error: None, + rule_count: 2, + etag: None, + }; + let za_source = mk("za", AdBlockResponse::ZeroAddress); + let nx_source = mk("nx", AdBlockResponse::NxDomain); + // Seed each source's cache file with parsed hosts-format content. + mhost_storage::adblock::write_cache( + temp.path(), + &za_source.source_id, + b"0.0.0.0 ads.example.com\n0.0.0.0 tracker.example.com\n", + ) + .unwrap(); + mhost_storage::adblock::write_cache( + temp.path(), + &nx_source.source_id, + b"0.0.0.0 blocked.example.com\n", + ) + .unwrap(); + let state = AdBlockState { + enabled: true, + sources: vec![za_source, nx_source], + whitelist: vec!["safe.example.com".to_string()], + ..Default::default() + }; + let (z, n, w) = classify_rules(&state, temp.path()); + assert_eq!(z.len(), 2, "zero_addr set seeded from za source cache"); + assert!(z.contains_key("ads.example.com")); + assert!(z.contains_key("tracker.example.com")); + assert_eq!(n.len(), 1, "nxdomain set seeded from nx source cache"); + assert!(n.contains("blocked.example.com")); + assert_eq!(w.len(), 1); + } + + // ----------------------------------------------------------------- + // PR #154 review (P2): exercise the cold-start hot-reload path that + // AppState::new runs when `dns_enabled=true` was recovered from the + // manifest. The headline fix is "DNS goes OFF → user adds source → + // DNS goes ON → first query sees cached rules immediately" — without + // the cold-start hot-reload there's a window where DNS is running but + // ad-block isn't active yet. + // + // Test simulates the full flow without spinning up the proxy / Tauri + // runtime: classify_rules → reload_ad_block_rules → spin up a real + // DnsServer → fire a UDP query → assert the blocked domain returns + // 0.0.0.0 instead of leaking upstream. + // ----------------------------------------------------------------- + #[tokio::test] + async fn cold_start_hot_reload_blocks_first_query() { + use mhost_dns::DnsConfig; + + let temp = tempfile::TempDir::new().unwrap(); + let source_id = SourceId(Uuid::new_v4()); + let source = AdBlockSource { + source_id: source_id.clone(), + name: "test-blocklist".into(), + url: "https://x".into(), + enabled: true, + response: AdBlockResponse::ZeroAddress, + last_fetched_at: None, + last_error: None, + rule_count: 1, + etag: None, + }; + mhost_storage::adblock::write_cache( + temp.path(), + &source_id, + b"0.0.0.0 cold-start-ads.example.com\n", + ) + .unwrap(); + + let state = AdBlockState { + enabled: true, + sources: vec![source], + whitelist: vec![], + ..Default::default() + }; + + // Simulate the AppState::new cold-start block. + let (za, nx, wl) = classify_rules(&state, temp.path()); + assert!(za.contains_key("cold-start-ads.example.com")); + + // Wire into a real DnsServer and query. + let port = pick_free_port(); + let config = DnsConfig { + port, + upstream: vec!["127.0.0.1:1".to_string()], // blackhole — fail fast + timeout_ms: 100, + refresh_upstream: false, + cache_size: 100, + }; + let server = std::sync::Arc::new(mhost_dns::DnsServer::new(config).unwrap()); + server.reload_ad_block_rules(za, nx, wl); + assert_eq!(server.ad_block_rule_count(), 1); + + let server_clone = std::sync::Arc::clone(&server); + let server_handle = tokio::spawn(async move { server_clone.start().await }); + // Wait for the server to be listening. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); + while !server.is_running() && std::time::Instant::now() < deadline { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!(server.is_running(), "server should start"); + + // Send a UDP query for the blocked domain. + use hickory_proto::op::{Message, OpCode, Query}; + use hickory_proto::rr::{Name, RecordType}; + use hickory_proto::serialize::binary::{BinDecodable, BinEncodable}; + use tokio::net::UdpSocket; + let query_name = Name::from_utf8("cold-start-ads.example.com.").unwrap(); + let query = Query::query(query_name, RecordType::A); + let mut request = Message::new(); + request.set_id(0x4242); + request.set_recursion_desired(true); + request.set_op_code(OpCode::Query); + request.add_query(query); + let bytes = request.to_bytes().unwrap(); + + let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + client + .send_to(&bytes, format!("127.0.0.1:{}", port)) + .await + .unwrap(); + let mut buf = vec![0u8; 4096]; + let (len, _) = tokio::time::timeout( + std::time::Duration::from_secs(2), + client.recv_from(&mut buf), + ) + .await + .expect("server response timeout") + .expect("recv_from failed"); + let response = hickory_proto::op::Message::from_bytes(&buf[..len]).unwrap(); + assert_eq!( + response.response_code(), + hickory_proto::op::ResponseCode::NoError + ); + assert_eq!( + response.answer_count(), + 1, + "blocked domain should be answered" + ); + let answer = &response.answers()[0]; + if let Some(hickory_proto::rr::RData::A(a)) = answer.data() { + assert_eq!( + a.0, + std::net::Ipv4Addr::new(0, 0, 0, 0), + "cold-start ad-block should return 0.0.0.0" + ); + } else { + panic!("expected A record, got {:?}", answer.data()); + } + + server.stop().await.unwrap(); + let _ = server_handle.await; + } + + /// Pick a free UDP port by binding to port 0. Avoids colliding with + /// other tests on the same machine. + fn pick_free_port() -> u16 { + let listener = std::net::UdpSocket::bind("127.0.0.1:0").expect("bind free port"); + let port = listener.local_addr().unwrap().port(); + drop(listener); + port + } + + // ----------------------------------------------------------------- + // PR #131 re-review P1-2: a fetch failure must NOT skip persistence — + // the source was already pushed into in-memory state, and skipping + // `persist_and_reload` lost it on restart. Point the source at a loopback + // port that refuses connections so `fetch_source` fails fast (no 30s + // timeout). The source should still be on disk after the call errors. + // ----------------------------------------------------------------- + #[tokio::test] + async fn add_ad_block_source_persists_on_fetch_failure() { + use crate::state::AppState; + use mhost_apply::writer::HostsWriter; + use mhost_storage::storage::FileStorage; + + let temp = tempfile::TempDir::new().unwrap(); + let storage = Arc::new(FileStorage::new(temp.path())) + as Arc; + let state = AppState { + storage: storage.clone(), + writer: Arc::new(HostsWriter::new()), + apply_lock: crate::state::ApplyLock::new(), + snapshot_lock: crate::state::ApplyLock::new(), + last_profile_ids: std::sync::Mutex::new(Vec::new()), + dns_server: Arc::new(std::sync::Mutex::new(None)), + dns_enabled: std::sync::atomic::AtomicBool::new(false), + original_dns: std::sync::Mutex::new(mhost_core::OriginalDns::DhcpEmpty), + dns_lock: crate::state::ApplyLock::new(), + ad_block_state: Arc::new(tokio::sync::RwLock::new(AdBlockState::default())), + ad_block_refresh_task: std::sync::Mutex::new(None), + ad_block_refresh_cancel: std::sync::Mutex::new( + tokio_util::sync::CancellationToken::new(), + ), + }; + + // Port 1 on loopback refuses connections → fetch_source errors fast. + let url = "http://127.0.0.1:1/blocklist".to_string(); + let err = + add_ad_block_source_impl(&state, "failing".into(), url, AdBlockResponse::ZeroAddress) + .await + .expect_err("fetch should fail (connection refused)"); + assert!( + err.to_string().contains("fetch") + || err.to_string().to_lowercase().contains("connect") + || err.to_string().to_lowercase().contains("error") + ); + + // P1-2 invariant: the source is persisted despite the fetch failure. + let persisted = mhost_storage::adblock::read_state(storage.root()) + .expect("adblock.json should exist after persist_and_reload"); + assert_eq!( + persisted.sources.len(), + 1, + "source must be persisted even when initial fetch fails (P1-2)" + ); + assert_eq!(persisted.sources[0].name, "failing"); + assert!( + persisted.sources[0].last_error.is_some(), + "last_error must be recorded on the persisted source" + ); + } +} diff --git a/src-tauri/src/commands/dns.rs b/src-tauri/src/commands/dns.rs index 15e8594..f19be97 100644 --- a/src-tauri/src/commands/dns.rs +++ b/src-tauri/src/commands/dns.rs @@ -1,9 +1,12 @@ use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; use mhost_core::{MhostError, OriginalDns, ProfileMode}; use tauri::State; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; -use crate::state::AppState; +use crate::state::{lock_or_recover, AppState}; /// 启动/停止 DNS 模式。 /// @@ -182,6 +185,30 @@ async fn set_dns_mode_enable(state: &AppState) -> Result<(), MhostError> { } state.dns_enabled.store(true, Ordering::Relaxed); + // 9. 广告屏蔽(issue #130):启用 DNS 后立即把当前 ad-block 状态 + // hot-reload 到新 server,并启动定时刷新 task。task 在 disable + // 时被 abort。 + // + // 9a. 即时 reload:spawn_ad_block_refresh_task 在 + // auto_refresh_enabled=false 或 interval=0 时不会 spawn,但用户 + // 仍期望持久化的 ad-block 规则立即生效。所以这里显式做一次 + // classify + reload,与 AppState::new 冷启动路径一致。 + // + // 9b. 这里复用了 commands::adblock 的 `classify_rules` + 重载路径 + // 的等价逻辑(避免循环依赖和 IPC 边界),不经过 IPC handler。 + let snap = state.ad_block_state.read().await.clone(); + let (za, nx, wl) = crate::commands::adblock::classify_rules(&snap, state.storage.root()); + if let Some(server) = lock_or_recover(&state.dns_server).as_ref() { + server.reload_ad_block_rules(za, nx, wl); + } + spawn_ad_block_refresh_task( + &state.ad_block_refresh_task, + &state.ad_block_state, + &state.dns_server, + &state.storage, + &state.ad_block_refresh_cancel, + ); + Ok(()) } @@ -265,6 +292,16 @@ async fn set_dns_mode_disable(state: &AppState, interactive: bool) -> Result<(), // 5. 清 in-memory dns_enabled state.dns_enabled.store(false, Ordering::Relaxed); + // 6. 终止广告屏蔽后台刷新 task(issue #130, #138)。enable 时 spawn, + // disable 必须 abort;不 abort 会让 task 继续跑并尝试 reload + // 已停的 server。**先 cancel 再 abort**:cancel 让 refresh loop + // 的 `select!` 醒来并把 spawn_blocking 闭包里的 + // `is_cancelled()` check 触发,从而避免在已停的 server 上 + // reload。abort 是兜底:如果 task 还卡在 select 之外(比如 + // spawn_blocking 闭包刚起来),cancel() 的 wake 不会传到那里。 + cancel_ad_block_refresh_task(&state.ad_block_refresh_cancel); + abort_ad_block_refresh_task(&state.ad_block_refresh_task); + Ok(()) } @@ -384,6 +421,9 @@ mod tests { dns_enabled: AtomicBool::new(false), original_dns: Mutex::new(OriginalDns::DhcpEmpty), dns_lock: ApplyLock::new(), + ad_block_state: Arc::new(tokio::sync::RwLock::new(mhost_core::AdBlockState::default())), + ad_block_refresh_task: Mutex::new(None), + ad_block_refresh_cancel: Mutex::new(tokio_util::sync::CancellationToken::new()), }; // dns_enabled = false → cleanup 应直接返回 Ok let result = cleanup_dns_on_exit(&state, false).await; @@ -422,6 +462,9 @@ mod tests { dns_enabled: AtomicBool::new(true), // 假装启用 → cleanup 会走 disable 路径 original_dns: Mutex::new(OriginalDns::DhcpEmpty), // DhcpEmpty → 写 Empty dns_lock: ApplyLock::new(), + ad_block_state: Arc::new(tokio::sync::RwLock::new(mhost_core::AdBlockState::default())), + ad_block_refresh_task: Mutex::new(None), + ad_block_refresh_cancel: Mutex::new(tokio_util::sync::CancellationToken::new()), }; // cleanup_dns_on_exit → set_dns_mode_disable(interactive=false) // - original 是 DhcpEmpty → 只打印 warning(不返回 Err,bug 1 修复) @@ -473,6 +516,9 @@ mod tests { dns_enabled: AtomicBool::new(true), original_dns: Mutex::new(OriginalDns::DhcpEmpty), dns_lock: ApplyLock::new(), + ad_block_state: Arc::new(tokio::sync::RwLock::new(mhost_core::AdBlockState::default())), + ad_block_refresh_task: Mutex::new(None), + ad_block_refresh_cancel: Mutex::new(tokio_util::sync::CancellationToken::new()), }; // 第一次 cleanup:跑 disable 路径。注意必须用 interactive=false @@ -546,3 +592,228 @@ pub async fn get_dns_status( }; Ok(status) } +/// Abort the periodic ad-block refresh task if one is registered. +/// +/// The disable path (issue #130) is the only legitimate caller: +/// `spawn_ad_block_refresh_task` registers a `JoinHandle` on enable, and +/// `set_dns_mode_disable` MUST cancel it — otherwise the task keeps +/// running and tries to `reload_ad_block_rules` on a server that's +/// already been stopped, surfacing confusing errors at the next refresh +/// tick. +/// +/// Extracted from the inline `if let Some(h) = ...take() { h.abort() }` +/// in `set_dns_mode_disable` so the abort behavior is unit-testable +/// without going through the full disable path (issue #134). The full +/// path calls `mhost_dns::platform::disable_dns_mode` which returns Err +/// in unit tests (no proxy + non-interactive), short-circuiting before +/// the abort step — so testing through the public API would not +/// actually exercise the abort contract. +/// +/// **This helper only signals cancellation.** It does NOT wait for the +/// task to actually terminate, and it does NOT abort any `spawn_blocking` +/// closure the task may have entered (Tokio explicitly does not support +/// aborting blocking work once started). The full "refresh work is dead +/// by the time we return" guarantee is a separate concern tracked as a +/// follow-up issue; see the test module's section comment for context. +/// +/// Returns `true` if a task was found and `abort()` was called on it, +/// `false` if the slot was empty (idempotent re-disable, or +/// disable-before-enable). +fn abort_ad_block_refresh_task(slot: &Mutex>>) -> bool { + if let Some(handle) = lock_or_recover(slot).take() { + handle.abort(); + true + } else { + false + } +} + +/// Cooperatively cancel the periodic ad-block refresh task. +/// +/// Pair to `abort_ad_block_refresh_task`. Where the latter force-cancels +/// the outer task via `JoinHandle::abort()` (which cannot interrupt +/// in-flight `spawn_blocking` closures — see issue #138), the cancel +/// token lets the task's `tokio::select!` wake up immediately and lets +/// any `spawn_blocking` closure observe `is_cancelled()` and bail before +/// mutating a stopped `DnsServer`. +/// +/// The disable path calls cancel **before** abort so the cooperative +/// path runs first; the abort is the fallback for any work that is past +/// the select point. +/// +/// Idempotent: calling on an already-cancelled token is a no-op. +fn cancel_ad_block_refresh_task(slot: &Mutex) { + lock_or_recover(slot).cancel(); +} + +/// Spawn the periodic ad-block refresh task (issue #130). +/// +/// Called from two places: +/// 1. `set_dns_mode_enable` after the DNS server comes up — the +/// "normal" hot path. +/// 2. `AppState::new` after `try_recover_dns` succeeds — PR #131 +/// review finding 0.1: previously, an auto-recovered DNS session +/// lost its periodic refresh because the task was only spawned on +/// the user-driven enable path. +/// +/// The function takes individual Arcs rather than `&AppState` so it can +/// be invoked while `AppState` is being constructed (item 2). +/// +/// **Issue #138 follow-up (re-enable):** the cancel slot is **swapped** +/// for a fresh, uncancelled token on every spawn — `CancellationToken` +/// is sticky, so without this swap a disable → re-enable cycle would +/// hand the new task the old (already cancelled) token, causing its +/// `select!` to match `cancel.cancelled()` on iter 0 and exit +/// immediately. See `test_re_enable_after_disable_respawns_with_fresh_token`. +/// +/// The task: +/// 1. Reads `refresh_interval_hours` from `ad_block_state`. +/// 2. Sleeps for the interval (or until the cancel token fires). +/// 3. Refreshes all enabled sources + hot-reloads the engine. +/// 4. Exits cleanly when the current `ad_block_refresh_cancel` is cancelled. +/// +/// `refresh_interval_hours == 0` or `auto_refresh_enabled == false` short- +/// circuits — task is not spawned at all (callers don't need to abort it). +/// +/// **Issue #138:** The disable path now signals a `CancellationToken` (see +/// `cancel_ad_block_refresh_task`) *before* aborting. The select! below +/// wakes on `token.cancelled()` and the loop exits without relying on +/// `JoinHandle::abort()` reaching a yield point. The `spawn_blocking` +/// closure inside the loop also checks `token.is_cancelled()` immediately +/// before calling `reload_ad_block_rules` — that's the layer that protects +/// against an in-flight `classify_rules` that started before cancel +/// landed. `spawn_blocking` work cannot be interrupted by `JoinHandle:: +/// abort()` (tokio explicitly documents this), so the self-check is the +/// only reliable way to avoid a `reload_ad_block_rules` call landing on +/// a `DnsServer` that's already been stopped. The `lock_or_recover` +/// on `dns_server` already serializes against the disable path's `.take()`, +/// so the cancel check is mainly an early-exit optimization against a +/// `reload_ad_block_rules` racing with `server.stop()`. +pub(crate) fn spawn_ad_block_refresh_task( + task_slot: &std::sync::Mutex>>, + ad_block_state: &Arc>, + dns_server: &Arc>>, + storage: &Arc, + cancel_slot: &Mutex, +) { + let cfg = match ad_block_state.try_read() { + Ok(g) => (g.auto_refresh_enabled, g.refresh_interval_hours), + Err(_) => return, + }; + if !cfg.0 || cfg.1 == 0 { + return; + } + + // Issue #138 follow-up (re-enable): swap the cancel slot for a fresh + // token so this task is unaffected by a previous disable's `cancel()`. + // `CancellationToken::cancel()` is sticky — if we just cloned the + // existing (already cancelled) token, the spawned task's first + // `select!` arm would match `cancel.cancelled()` on iter 0 and break + // out of the loop without ever ticking. + let cancel = CancellationToken::new(); + *lock_or_recover(cancel_slot) = cancel.clone(); + + // Clone the few Arcs we need into the task closure. We don't share the + // whole AppState (which contains unrelated Mutexes like snapshot_lock) + // to keep the lock-contention surface minimal. The cancel token is + // cheap to clone (it's a refcount + an atomic flag). + let storage = storage.clone(); + let ad_block_state = ad_block_state.clone(); + let dns_server = dns_server.clone(); + + let interval_secs = (cfg.1 as u64).saturating_mul(3600).max(3600); // floor 1h + let handle = tokio::spawn(async move { + loop { + // Sleep OR cancellation. The select! is the **first** thing + // the loop checks so disable can interrupt even a long + // inter-tick sleep. + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_secs(interval_secs)) => {} + _ = cancel.cancelled() => { + // Disable path fired the token. Exit the loop cleanly. + break; + } + } + + // Re-read config — user may have changed interval or disabled + // auto-refresh since the last tick. + let (auto, interval_h, enabled) = { + let Ok(g) = ad_block_state.try_read() else { + continue; + }; + (g.auto_refresh_enabled, g.refresh_interval_hours, g.enabled) + }; + if !auto || interval_h == 0 || !enabled { + continue; + } + + // Snapshot IDs (best-effort; failures logged not propagated). + let ids: Vec = { + let Ok(g) = ad_block_state.try_read() else { + continue; + }; + g.sources + .iter() + .filter(|s| s.enabled) + .map(|s| s.source_id.clone()) + .collect() + }; + + // Concurrent fetch — bounded by REFRESH_CONCURRENCY. PR #131 + // review finding 1.4: serial loop meant a multi-source list + // could block for `N × FETCH_TIMEOUT_SECS` while the next + // periodic tick waited its turn. + // + // Note: `fetch_sources_concurrent` is a regular async fn, not + // spawn_blocking, so `cancel.cancelled()` racing with it would + // not interrupt it. We accept that one fetch round may run to + // completion after disable; the spawn_blocking step below is + // the critical one (it mutates the server), and that's where + // the self-check lives. + crate::commands::adblock::fetch_sources_concurrent( + &storage, + &ad_block_state, + &ids, + crate::commands::adblock::REFRESH_CONCURRENCY, + ) + .await; + + // Bail before spawn_blocking if the token is already set — + // the per-tick check above the fetch is informational only, + // cancel may have landed during the fetch. + if cancel.is_cancelled() { + break; + } + + // Hot-reload engine if DNS still on. We use the dns_server + // slot as the proxy signal AND the cancel token as the + // authoritative one — the slot can race with the disable + // path between the check and the reload call below. + if lock_or_recover(&dns_server).is_some() { + let snap = ad_block_state.read().await.clone(); + let root = storage.root().to_path_buf(); + let dns_server_clone = Arc::clone(&dns_server); + let cancel_in_closure = cancel.clone(); + // Issue #133: classify_rules reads + parses each source's + // cache file synchronously — 100k+ domains can block a + // tokio worker for seconds. Move it off the async runtime. + let _ = tokio::task::spawn_blocking(move || { + // Self-check (issue #138) — abort() can't reach us; bail on cancel. + if cancel_in_closure.is_cancelled() { + return; + } + let (za, nx, wl) = crate::commands::adblock::classify_rules(&snap, &root); + if cancel_in_closure.is_cancelled() { + return; + } + if let Some(server) = lock_or_recover(&dns_server_clone).as_ref() { + server.reload_ad_block_rules(za, nx, wl); + } + }) + .await; + } + } + }); + + *lock_or_recover(task_slot) = Some(handle); +} diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index b0f30ed..235f946 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -1,3 +1,4 @@ +pub mod adblock; pub mod apply; pub mod dns; #[cfg(test)] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index a55e0ae..63d2bd5 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -7,7 +7,9 @@ pub mod tray_logic; use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; -use commands::{apply::*, dns::*, profile::*, profile_io::*, snapshot::*, update::*, validate::*}; +use commands::{ + adblock::*, apply::*, dns::*, profile::*, profile_io::*, snapshot::*, update::*, validate::*, +}; use state::AppState; use tauri::{Manager, RunEvent}; @@ -157,6 +159,20 @@ pub fn run() { reload_dns_rules, get_dns_status, list_dns_profiles, + // 广告屏蔽(issue #130) + get_ad_block_state, + set_ad_block_enabled, + set_ad_block_refresh_interval, + list_ad_block_sources, + add_ad_block_source, + remove_ad_block_source, + set_ad_block_source_enabled, + set_ad_block_source_response, + refresh_ad_block_source, + refresh_all_ad_block_sources, + list_ad_block_whitelist, + add_ad_block_whitelist, + remove_ad_block_whitelist, check_update, ]) .setup(|app| { diff --git a/src-tauri/src/state/mod.rs b/src-tauri/src/state/mod.rs index db27881..5443ed5 100644 --- a/src-tauri/src/state/mod.rs +++ b/src-tauri/src/state/mod.rs @@ -1,9 +1,24 @@ use mhost_apply::writer::HostsWriter; -use mhost_core::{MhostError, OriginalDns, ProfileMode}; +use mhost_core::{AdBlockState, MhostError, OriginalDns, ProfileMode}; use mhost_storage::migration::migrate_v1_to_v2; use mhost_storage::storage::{FileStorage, Storage}; -use std::sync::atomic::AtomicBool; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; +use tokio::task::JoinHandle; + +/// Poison-recovery helper for `std::sync::Mutex` (issue #130, PR #131 +/// re-review). Returns the inner guard even if a previous holder panicked. +/// +/// `tokio::sync::Mutex` (used by [`ApplyLock`]) does not have poison — a +/// panicked holder releases the lock automatically — so this helper is +/// only needed for plain `std::sync::Mutex` slots (e.g. the ad block +/// refresh task slot and cancel slot). +pub(crate) fn lock_or_recover(mutex: &std::sync::Mutex) -> std::sync::MutexGuard<'_, T> { + match mutex.lock() { + Ok(g) => g, + Err(poisoned) => poisoned.into_inner(), + } +} /// Async mutex to serialize apply operations and prevent concurrent writes to /etc/hosts. /// Security fix (#16): Prevents race conditions when user rapidly toggles profiles. @@ -54,6 +69,23 @@ pub struct AppState { pub original_dns: Mutex, /// 串行化 DNS 模式切换操作。 pub dns_lock: ApplyLock, + + // ------------------------------------------------------------------- + // 广告屏蔽(issue #130) + // ------------------------------------------------------------------- + /// 当前持久化的 ad block 状态(含 sources / whitelist / refresh 配置)。 + /// 命令层 `set_*` 操作通过 `tokio::sync::RwLock::write().await` 修改, + /// DNS 集成读路径用 `.read().await` 拿 snapshot 喂给 + /// `DnsServer::reload_ad_block_rules`。 + pub ad_block_state: Arc>, + /// 后台 ad block 定时刷新 task 句柄(`spawn_ad_block_refresh_task`)。 + /// `set_dns_mode_disable` / `cleanup_dns_on_exit` 时 `take()` 出来 abort。 + pub ad_block_refresh_task: Mutex>>, + /// ad block 定时刷新 task 的 cancel 令牌。 + /// `cancel()` 唤醒 `select!` 中的 sleep 分支,让 disable / cleanup 立即 + /// 生效;`spawn_ad_block_refresh_task` 在 spawn 前 swap 一个新 token + /// 避免上次 token 的 stickiness 干扰下次启用(issue #138)。 + pub ad_block_refresh_cancel: Mutex, } impl AppState { @@ -88,6 +120,7 @@ impl AppState { } let storage = Arc::new(file_storage); + let storage_root = storage.root().to_path_buf(); let writer = Arc::new(HostsWriter::new()); // 从 manifest 恢复 DNS 模式状态(不存在则创建默认) @@ -131,17 +164,66 @@ impl AppState { } } - Ok(Self { + let dns_server = Arc::new(Mutex::new(dns_server_opt)); + + let state = Self { storage, writer, apply_lock: ApplyLock(tokio::sync::Mutex::new(())), snapshot_lock: ApplyLock(tokio::sync::Mutex::new(())), last_profile_ids: Mutex::new(Vec::new()), - dns_server: Arc::new(Mutex::new(dns_server_opt)), + dns_server, dns_enabled: AtomicBool::new(dns_enabled), original_dns: Mutex::new(original_dns), dns_lock: ApplyLock(tokio::sync::Mutex::new(())), - }) + // Ad block(issue #130):从 adblock.json 恢复,损坏时自动备份。 + ad_block_state: Arc::new(tokio::sync::RwLock::new( + mhost_storage::adblock::read_state_or_default_with_backup(&storage_root), + )), + ad_block_refresh_task: Mutex::new(None), + ad_block_refresh_cancel: Mutex::new(tokio_util::sync::CancellationToken::new()), + }; + + // 冷启动自动恢复(PR #131 review P1-1):如果上次退出时 + // dns_enabled=true,DNS server 已经起来 + 持久化的 ad-block + // 状态已加载;立即 hot-reload 当前规则到刚构造的 engine,并启动 + // 定时刷新 task。否则会有一段「DNS 通了但 ad-block 没生效」的空窗。 + // + // **PR #154 review (P2) defensive**: `classify_rules` is sync and + // reads each source's cache file from disk + parses 100k+ domains. + // On a slow filesystem (network mount, encrypted APFS) this could + // block `AppState::new` for seconds. Wrap the read+parse+reload + // in `spawn_blocking` so it runs on a dedicated blocking thread. + // `spawn_ad_block_refresh_task` itself is async (it spawns a + // tokio task + schedules a select!), so it stays on the async + // runtime — only the sync pipeline needs the offload. + if state.dns_enabled.load(Ordering::Relaxed) { + let snap = state.ad_block_state.read().await.clone(); + let storage_root = state.storage.root().to_path_buf(); + let dns_server = Arc::clone(&state.dns_server); + let result = tokio::task::spawn_blocking(move || { + let (za, nx, wl) = crate::commands::adblock::classify_rules(&snap, &storage_root); + if let Some(server) = crate::state::lock_or_recover(&dns_server).as_ref() { + server.reload_ad_block_rules(za, nx, wl); + } + }) + .await; + if let Err(e) = result { + eprintln!( + "[mHost] cold-start ad-block reload join error: {} (continuing without hot-reload)", + e + ); + } + crate::commands::dns::spawn_ad_block_refresh_task( + &state.ad_block_refresh_task, + &state.ad_block_state, + &state.dns_server, + &state.storage, + &state.ad_block_refresh_cancel, + ); + } + + Ok(state) } /// 尝试自动恢复 DNS 服务。 diff --git a/src-tauri/src/tray.rs b/src-tauri/src/tray.rs index b2ad2f9..e6c74f8 100644 --- a/src-tauri/src/tray.rs +++ b/src-tauri/src/tray.rs @@ -90,7 +90,13 @@ fn build_menu(app: &AppHandle) -> Result, Box)?; + let adblock = MenuItem::with_id( + app, + "adblock", + "广告屏蔽(仅 DNS 模式)", + true, + None::<&str>, + )?; let sep2 = PredefinedMenuItem::separator(app)?; let refresh = MenuItem::with_id(app, "refresh_rules", "刷新远程规则", true, Some("CmdOrR"))?; let open_window = MenuItem::with_id(app, "open_window", "打开主窗口", true, Some("CmdOrO"))?; @@ -212,7 +218,20 @@ pub fn handle_menu_event(app: &AppHandle, event: tauri::menu::Men }); } tray_logic::TrayMenuAction::AdBlock => { - // Placeholder: ad block is coming soon + // 跳转到 Ad Block 页面(issue #130)。仅在 DNS 模式下生效 — + // 前端会显示「DNS off」横幅并禁用表单提交。 + #[cfg(target_os = "macos")] + crate::platform::macos::set_activation_policy_regular(); + if let Some(window) = app.get_webview_window("main") { + let _ = window.unminimize(); + let _ = window.show(); + let _ = window.set_focus(); + // PR #154 review (P3): log on emit failure for parity + // with the surrounding unminimize/show/set_focus calls. + if let Err(e) = window.emit("navigate", "/ad-block") { + tracing::warn!("[mHost] tray emit navigate failed: {}", e); + } + } } tray_logic::TrayMenuAction::Unknown => { println!("[mHost] Unknown tray menu action: {:?}", event.id); diff --git a/src/App.test.tsx b/src/App.test.tsx index 68eec81..97a2a13 100644 --- a/src/App.test.tsx +++ b/src/App.test.tsx @@ -22,6 +22,14 @@ vi.mock("./lib/tauri", () => ({ rollbackHosts: vi.fn(), exportProfileToFile: vi.fn(), duplicateProfile: vi.fn(), + // AdBlock IPCs (issue #130) — App fetches ad-block state on mount. + getAdBlockState: vi.fn().mockResolvedValue({ + enabled: false, + sources: [], + whitelist: [], + auto_refresh_enabled: true, + refresh_interval_hours: 24, + }), })); vi.mock("@tauri-apps/plugin-dialog", () => ({ @@ -49,11 +57,24 @@ describe("App", () => { ); }); - expect(listenMock).toHaveBeenCalledTimes(1); - expect(listenMock).toHaveBeenCalledWith( + // PR #154 review (P2): keep total call count so a double-registration +// regression is caught. We register 2 listeners per mount: `tray:profiles-updated` +// (profile refresh) and `navigate` (issue #130 tray deep-link). React +// StrictMode in dev causes double mount → 2 × 2 = 4 calls. Use `>=` to +// tolerate the StrictMode double-mount while still catching obvious +// regression cases (e.g. a third registration). +expect(listenMock.mock.calls.length).toBeGreaterThanOrEqual(2); +expect(listenMock).toHaveBeenCalledWith( "tray:profiles-updated", expect.any(Function), ); + // issue #130: also listens for the tray-driven `navigate` event so the + // "广告屏蔽" menu item can deep-link to /ad-block without coupling + // backend to router. + expect(listenMock).toHaveBeenCalledWith( + "navigate", + expect.any(Function), + ); }); it("triggers profile refresh when tray:profiles-updated event fires", async () => { @@ -93,7 +114,8 @@ describe("App", () => { unmountFn = unmount; }); - expect(listenMock).toHaveBeenCalledTimes(1); + // Same count assertion as the first test — guards against double-register. + expect(listenMock.mock.calls.length).toBeGreaterThanOrEqual(2); await act(async () => { unmountFn(); diff --git a/src/App.tsx b/src/App.tsx index 5d69c3e..686fade 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -1,5 +1,5 @@ import { useEffect } from "react"; -import { Routes, Route, Navigate } from "react-router-dom"; +import { Routes, Route, Navigate, useNavigate } from "react-router-dom"; import { listen } from "@tauri-apps/api/event"; import { useSetAtom } from "jotai"; import Layout from "./components/Layout"; @@ -7,12 +7,20 @@ import ProfileView from "./pages/ProfileView"; import Settings from "./pages/Settings"; import SnapshotPage from "./pages/Snapshot"; import SystemHosts from "./pages/SystemHosts"; -import { fetchProfilesAtom, fetchDnsProfilesAtom, fetchDnsModeAtom } from "./stores/profiles"; +import AdBlock from "./pages/AdBlock"; +import { + fetchProfilesAtom, + fetchDnsProfilesAtom, + fetchDnsModeAtom, + fetchAdBlockStateAtom, +} from "./stores/profiles"; function App() { const fetchProfiles = useSetAtom(fetchProfilesAtom); const fetchDnsProfiles = useSetAtom(fetchDnsProfilesAtom); const fetchDnsMode = useSetAtom(fetchDnsModeAtom); + const fetchAdBlock = useSetAtom(fetchAdBlockStateAtom); + const navigate = useNavigate(); useEffect(() => { // Load profiles on app mount @@ -25,14 +33,38 @@ function App() { fetchDnsMode().catch(() => { // Ignore: error is already stored in dnsErrorAtom }); + fetchAdBlock().catch(() => { + // Ignore: error is already stored in adBlockErrorAtom + }); - const unlisten = listen("tray:profiles-updated", () => { + const unlistenProfiles = listen("tray:profiles-updated", () => { fetchProfiles(); }); + // issue #130: tray "广告屏蔽" menu item emits this event with the + // target route. Lets the tray drive deep-linking to /ad-block without + // coupling backend to router. + // + // **security (PR #154 review P2)**: whitelist allowed routes. The + // backend emitter (tray.rs `TrayMenuAction::AdBlock`) is trusted but + // any future payload source — including a malformed emitter or a + // future debug/test hook — must not be able to push the router into + // arbitrary paths (which would silently no-op render or, worse, leak + // some future route meant for in-app use only). + const ALLOWED_TRAY_ROUTES: ReadonlySet = new Set(["/ad-block"]); + const unlistenNavigate = listen("navigate", (event) => { + const target = event.payload; + if (typeof target === "string" && ALLOWED_TRAY_ROUTES.has(target)) { + navigate(target); + } else if (typeof target === "string" && target.startsWith("/")) { + // Unknown path — log and ignore. + console.warn(`[mHost] tray navigate: refused unknown route "${target}"`); + } + }); return () => { - unlisten.then((fn) => fn()).catch(() => {}); + unlistenProfiles.then((fn) => fn()).catch(() => {}); + unlistenNavigate.then((fn) => fn()).catch(() => {}); }; - }, [fetchProfiles, fetchDnsProfiles, fetchDnsMode]); + }, [fetchProfiles, fetchDnsProfiles, fetchDnsMode, fetchAdBlock, navigate]); return ( @@ -45,6 +77,7 @@ function App() { } /> } /> } /> + } /> ); diff --git a/src/components/Layout.tsx b/src/components/Layout.tsx index 7b1928c..67d6b7f 100644 --- a/src/components/Layout.tsx +++ b/src/components/Layout.tsx @@ -120,11 +120,9 @@ function DnsIcon() { const toolNavItems: NavItem[] = [ { - to: "#adblock", + to: "/ad-block", label: "Ad Block", icon: , - badge: 10, - disabled: true, }, { to: "#remote", diff --git a/src/components/__tests__/Layout.test.tsx b/src/components/__tests__/Layout.test.tsx index 7b40b30..039399d 100644 --- a/src/components/__tests__/Layout.test.tsx +++ b/src/components/__tests__/Layout.test.tsx @@ -192,7 +192,7 @@ describe("Layout", () => { it("renders disabled tool items with 'Soon' badge", () => { renderWithProviders(); - expect(screen.getAllByText("Soon").length).toBeGreaterThanOrEqual(2); + expect(screen.getAllByText("Soon").length).toBeGreaterThanOrEqual(1); }); it("navigates to profile page when profile item is clicked", async () => { diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 3e3ad01..411d560 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -1,5 +1,16 @@ import { invoke } from "@tauri-apps/api/core"; -import type { Profile, ApplyPlan, ValidateResult, ExportFormat, SnapshotMeta, DnsStatus, ProfileMode } from "../types"; +import type { + Profile, + ApplyPlan, + ValidateResult, + ExportFormat, + SnapshotMeta, + DnsStatus, + ProfileMode, + AdBlockState, + AdBlockSource, + AdBlockResponse, +} from "../types"; // ---- Profile commands ---- @@ -149,6 +160,78 @@ export async function listDnsProfiles(): Promise { return invoke("list_dns_profiles"); } +// ---- AdBlock commands (issue #130) ---- + +export async function getAdBlockState(): Promise { + return invoke("get_ad_block_state"); +} + +export async function setAdBlockEnabled(enabled: boolean): Promise { + return invoke("set_ad_block_enabled", { enabled }); +} + +export async function setAdBlockRefreshInterval(hours: number): Promise { + return invoke("set_ad_block_refresh_interval", { hours }); +} + +export async function listAdBlockSources(): Promise { + return invoke("list_ad_block_sources"); +} + +export async function addAdBlockSource( + name: string, + url: string, + response: AdBlockResponse, +): Promise { + return invoke("add_ad_block_source", { name, url, response }); +} + +export async function removeAdBlockSource(sourceId: string): Promise { + return invoke("remove_ad_block_source", { sourceId }); +} + +export async function setAdBlockSourceEnabled( + sourceId: string, + enabled: boolean, +): Promise { + return invoke("set_ad_block_source_enabled", { + sourceId, + enabled, + }); +} + +export async function setAdBlockSourceResponse( + sourceId: string, + response: AdBlockResponse, +): Promise { + return invoke("set_ad_block_source_response", { + sourceId, + response, + }); +} + +export async function refreshAdBlockSource( + sourceId: string, +): Promise { + return invoke("refresh_ad_block_source", { sourceId }); +} + +export async function refreshAllAdBlockSources(): Promise { + return invoke("refresh_all_ad_block_sources"); +} + +export async function listAdBlockWhitelist(): Promise { + return invoke("list_ad_block_whitelist"); +} + +export async function addAdBlockWhitelist(domain: string): Promise { + return invoke("add_ad_block_whitelist", { domain }); +} + +export async function removeAdBlockWhitelist(domain: string): Promise { + return invoke("remove_ad_block_whitelist", { domain }); +} + // ---- Update commands ---- export interface LatestRelease { diff --git a/src/pages/AdBlock.module.css b/src/pages/AdBlock.module.css new file mode 100644 index 0000000..9cddcce --- /dev/null +++ b/src/pages/AdBlock.module.css @@ -0,0 +1,246 @@ +/* ---- AdBlock page layout ---- */ + +.pageBody { + display: flex; + flex-direction: column; + gap: 20px; + max-width: 800px; +} + +.summaryGrid { + display: grid; + grid-template-columns: repeat(3, 1fr); + gap: 12px; + margin-top: 12px; +} + +@media (max-width: 600px) { + .summaryGrid { + grid-template-columns: 1fr; + } +} + +.statCard { + padding: 12px 14px; + background: var(--color-surface, var(--paper)); + border: 1px solid var(--rule, #e5e5e5); + border-radius: var(--radius-md, 6px); +} + +.statValue { + font-size: 20px; + font-weight: 700; + color: var(--color-text, var(--ink)); + font-variant-numeric: tabular-nums; +} + +.statLabel { + font-size: 11px; + text-transform: uppercase; + letter-spacing: 0.04em; + color: var(--color-muted, var(--ink-muted)); + margin-top: 4px; +} + +/* ---- Source card ---- */ + +.sourceCard { + border: 1px solid var(--rule, #e5e5e5); + border-radius: var(--radius-md, 6px); + padding: 14px 16px; + background: var(--color-surface, var(--paper)); + display: flex; + flex-direction: column; + gap: 10px; +} + +.sourceCard.dimmed { + opacity: 0.55; +} + +.sourceHeader { + display: flex; + justify-content: space-between; + align-items: flex-start; + gap: 12px; +} + +.sourceTitle { + display: flex; + align-items: center; + gap: 8px; + font-weight: 600; + color: var(--color-text, var(--ink)); + font-size: 14px; +} + +.sourceMeta { + font-size: 11px; + color: var(--color-muted, var(--ink-muted)); + word-break: break-all; +} + +.sourceActions { + display: flex; + align-items: center; + gap: 6px; + flex-wrap: wrap; +} + +.errorBadge { + font-size: 11px; + color: var(--color-danger, var(--danger)); + background: var(--color-danger-soft, #fdecec); + padding: 2px 8px; + border-radius: 10px; + display: inline-block; + word-break: break-all; +} + +/* ---- Add source form ---- */ + +.addSourceForm { + display: grid; + grid-template-columns: 1fr 2fr auto auto; + gap: 8px; + align-items: end; +} + +@media (max-width: 600px) { + .addSourceForm { + grid-template-columns: 1fr; + } +} + +/* ---- Whitelist ---- */ + +.whitelistList { + display: flex; + flex-wrap: wrap; + gap: 6px; + margin-top: 8px; +} + +.whitelistItem { + display: inline-flex; + align-items: center; + gap: 4px; + background: var(--color-surface, var(--paper)); + border: 1px solid var(--rule, #e5e5e5); + border-radius: 12px; + padding: 3px 4px 3px 10px; + font-size: 12px; + font-family: var(--font-mono, monospace); +} + +.removeBtn { + border: none; + background: transparent; + cursor: pointer; + font-size: 14px; + line-height: 1; + padding: 0 6px; + color: var(--color-muted, var(--ink-muted)); +} + +.removeBtn:hover { + color: var(--color-danger, var(--danger)); +} + +.empty { + color: var(--color-muted, var(--ink-muted)); + font-style: italic; + font-size: 13px; + padding: 12px 0; +} + +/* ---- Banner ---- */ + +.banner { + display: flex; + align-items: center; + gap: 10px; + background: var(--color-warning-soft, #fff5e0); + border: 1px solid var(--color-warning, var(--warning)); + color: var(--color-text, var(--ink)); + padding: 10px 14px; + border-radius: var(--radius-md, 6px); + font-size: 13px; +} + +.banner button { + margin-left: auto; +} + +/* ---- Inline form ---- */ + +.inlineForm { + display: flex; + gap: 6px; + align-items: center; +} + +.inlineForm input { + flex: 1; +} + +.muted { + color: var(--color-muted, var(--ink-muted)); + font-size: 12px; +} + +/* ---- Layout helpers (PR #154 review P2: was inline styles) ---- */ + +.bannerText { + display: flex; + justify-content: space-between; + align-items: center; + gap: 12px; +} + +.bannerTitle { + font-weight: 600; + color: var(--color-text, var(--ink)); +} + +.sectionGap { + margin-top: 10px; +} + +.columnGap { + display: flex; + flex-direction: column; + gap: 10px; + margin-top: 14px; +} + +.flexGrow { + flex: 1; + min-width: 0; +} + +.dangerText { + color: var(--color-danger, var(--danger)); +} + +.dangerTextGap { + margin-top: 10px; + color: var(--color-danger, var(--danger)); +} + +.mutedGap { + margin-bottom: 10px; +} + +.labelReset { + margin: 0; +} + +.width120 { + width: 120px; +} + +.badgeSm { + font-size: 12px; + padding: 2px 6px; +} diff --git a/src/pages/AdBlock.tsx b/src/pages/AdBlock.tsx new file mode 100644 index 0000000..9eed1f0 --- /dev/null +++ b/src/pages/AdBlock.tsx @@ -0,0 +1,438 @@ +import { useCallback, useState, useEffect } from "react"; +import { useAtomValue, useSetAtom } from "jotai"; +import { confirm as confirmDialog } from "@tauri-apps/plugin-dialog"; +import { + adBlockStateAtom, + isAdBlockLoadingAtom, + adBlockErrorAtom, + adBlockRuleCountAtom, + adBlockHasErrorsAtom, + dnsEnabledAtom, + fetchAdBlockStateAtom, + toggleAdBlockEnabledAtom, + setAdBlockIntervalAtom, + addAdBlockSourceAtom, + removeAdBlockSourceAtom, + setAdBlockSourceEnabledAtom, + setAdBlockSourceResponseAtom, + refreshAdBlockSourceAtom, + refreshAllAdBlockSourcesAtom, + addAdBlockWhitelistAtom, + removeAdBlockWhitelistAtom, +} from "../stores/profiles"; +import { useNavigate } from "react-router-dom"; +import { useWebKitPointerDown } from "../hooks/useWebKitPointerDown"; +import type { AdBlockResponse } from "../types"; +import styles from "./AdBlock.module.css"; + +function AdBlock() { + const state = useAtomValue(adBlockStateAtom); + const isLoading = useAtomValue(isAdBlockLoadingAtom); + const error = useAtomValue(adBlockErrorAtom); + const dnsEnabled = useAtomValue(dnsEnabledAtom); + const ruleCount = useAtomValue(adBlockRuleCountAtom); + const hasErrors = useAtomValue(adBlockHasErrorsAtom); + + const fetchState = useSetAtom(fetchAdBlockStateAtom); + const toggleEnabled = useSetAtom(toggleAdBlockEnabledAtom); + const setInterval = useSetAtom(setAdBlockIntervalAtom); + const addSource = useSetAtom(addAdBlockSourceAtom); + const removeSource = useSetAtom(removeAdBlockSourceAtom); + const setSourceEnabled = useSetAtom(setAdBlockSourceEnabledAtom); + const setSourceResponse = useSetAtom(setAdBlockSourceResponseAtom); + const refreshSource = useSetAtom(refreshAdBlockSourceAtom); + const refreshAll = useSetAtom(refreshAllAdBlockSourcesAtom); + const addWhitelist = useSetAtom(addAdBlockWhitelistAtom); + const removeWhitelist = useSetAtom(removeAdBlockWhitelistAtom); + + const { onPointerDown } = useWebKitPointerDown(); + const navigate = useNavigate(); + + // Local form state + const [newName, setNewName] = useState(""); + const [newUrl, setNewUrl] = useState(""); + const [newResponse, setNewResponse] = useState("zero_address"); + const [newWhitelistDomain, setNewWhitelistDomain] = useState(""); + + // Fetch on mount (idempotent — Tauri handles parallel calls). + useEffect(() => { + fetchState().catch(() => { + /* error already in atom */ + }); + }, [fetchState]); + + const handleAddSource = useCallback(() => { + if (!newName.trim() || !newUrl.trim()) return; + addSource({ name: newName.trim(), url: newUrl.trim(), response: newResponse }) + .then(() => { + setNewName(""); + setNewUrl(""); + }) + .catch(() => { + /* error in atom */ + }); + }, [addSource, newName, newUrl, newResponse]); + + const handleAddWhitelist = useCallback(() => { + const d = newWhitelistDomain.trim(); + if (!d) return; + addWhitelist(d) + .then(() => setNewWhitelistDomain("")) + .catch(() => { + /* error in atom */ + }); + }, [addWhitelist, newWhitelistDomain]); + + const handleIntervalChange = useCallback( + (hours: number) => { + setInterval(hours).catch(() => {}); + }, + [setInterval], + ); + + if (!state) { + return ( +
+
+

Ad Block

+
+
Loading…
+
+ ); + } + + // `dnsEnabled` flip controls whether the DNS engine actually applies + // ad-block rules. Configuration edits below are ALWAYS persisted to + // disk (and re-applied when DNS mode comes on), so the form is not + // disabled when DNS is off — users often configure sources + whitelist + // before enabling DNS mode for the first time. The banner below + // explains the effective state. + const dnsModeOff = !dnsEnabled; + + return ( +
+
+

Ad Block

+

+ Block ads at the DNS resolver. macOS DNS mode only. +

+
+ +
+
+ + {error &&
{error}
} + + {dnsModeOff && ( +
+ + DNS mode is off. Your edits below are saved and will apply the + next time you enable DNS mode. + + +
+ )} + +
+ {/* Master switch + summary */} +
+
+
+
Enable Ad Block
+
+ When enabled, the DNS server returns 0.0.0.0 / NXDOMAIN for + domains in any enabled source. +
+
+ +
+ +
+
+
{state.sources.length}
+
Sources
+
+
+
{ruleCount.toLocaleString()}
+
Active Rules
+
+
+
{state.whitelist.length}
+
Whitelist
+
+
+ + {hasErrors && ( +
+ One or more sources have a fetch error — see badges below. +
+ )} +
+ + {/* Add source form */} +
+

Sources

+

+ Hosts-format blocklist URLs (one domain per line, IP ignored). +

+ +
+
+ + setNewName(e.target.value)} + disabled={isLoading} + /> +
+
+ + setNewUrl(e.target.value)} + disabled={isLoading} + /> +
+
+ + +
+ +
+ + {/* Source list */} + {state.sources.length === 0 ? ( +
No sources yet.
+ ) : ( +
+ {state.sources.map((src) => ( +
+
+
+
+ {src.name} + {src.last_error && ( + + fetch failed + + )} +
+
{src.url}
+
+ {src.rule_count.toLocaleString()} rules + {src.last_fetched_at && + ` · fetched ${new Date(src.last_fetched_at).toLocaleString()}`} + {src.last_error && ( + <> + {" · "} + + {src.last_error} + + + )} +
+
+ +
+ + + + + + +
+
+
+ ))} +
+ )} +
+ + {/* Whitelist */} +
+

Whitelist

+

+ Domains here are exempt from all ad block rules. Suffix-matched: + adding example.com also exempts{" "} + api.example.com. +

+ +
+ setNewWhitelistDomain(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") handleAddWhitelist(); + }} + disabled={isLoading} + /> + +
+ + {state.whitelist.length === 0 ? ( +
No whitelist entries.
+ ) : ( +
+ {state.whitelist.map((d) => ( + + {d} + + + ))} +
+ )} +
+ + {/* Refresh interval */} +
+

Auto-refresh

+

+ Background refresh keeps sources up to date without manual + intervention. Set to 0 to disable (refresh manually instead). +

+
+ + +
+
+
+
+ ); +} + +export default AdBlock; diff --git a/src/pages/__tests__/AdBlock.test.tsx b/src/pages/__tests__/AdBlock.test.tsx new file mode 100644 index 0000000..ff9cffe --- /dev/null +++ b/src/pages/__tests__/AdBlock.test.tsx @@ -0,0 +1,301 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, fireEvent, act } from "@testing-library/react"; +import { MemoryRouter } from "react-router-dom"; +import { getDefaultStore, Provider as JotaiProvider } from "jotai"; +import { + adBlockStateAtom, + isAdBlockLoadingAtom, + adBlockErrorAtom, + dnsEnabledAtom, +} from "../../stores/profiles"; +import type { AdBlockState, AdBlockSource } from "../../types"; + +const mockGetAdBlockState = vi.fn(); +const mockSetAdBlockEnabled = vi.fn().mockResolvedValue(undefined); +const mockAddAdBlockSource = vi.fn().mockResolvedValue({}); +const mockRemoveAdBlockSource = vi.fn().mockResolvedValue(undefined); +const mockSetAdBlockSourceEnabled = vi.fn().mockResolvedValue({}); +const mockSetAdBlockSourceResponse = vi.fn().mockResolvedValue({}); +const mockRefreshAdBlockSource = vi.fn().mockResolvedValue({}); +const mockRefreshAllAdBlockSources = vi.fn().mockResolvedValue([]); +const mockAddAdBlockWhitelist = vi.fn().mockResolvedValue([]); +const mockRemoveAdBlockWhitelist = vi.fn().mockResolvedValue([]); +const mockSetAdBlockRefreshInterval = vi.fn().mockResolvedValue(undefined); + +vi.mock("../../lib/tauri", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getAdBlockState: (...args: unknown[]) => mockGetAdBlockState(...args), + setAdBlockEnabled: (...args: unknown[]) => mockSetAdBlockEnabled(...args), + addAdBlockSource: (...args: unknown[]) => mockAddAdBlockSource(...args), + removeAdBlockSource: (...args: unknown[]) => mockRemoveAdBlockSource(...args), + setAdBlockSourceEnabled: (...args: unknown[]) => mockSetAdBlockSourceEnabled(...args), + setAdBlockSourceResponse: (...args: unknown[]) => mockSetAdBlockSourceResponse(...args), + refreshAdBlockSource: (...args: unknown[]) => mockRefreshAdBlockSource(...args), + refreshAllAdBlockSources: (...args: unknown[]) => mockRefreshAllAdBlockSources(...args), + addAdBlockWhitelist: (...args: unknown[]) => mockAddAdBlockWhitelist(...args), + removeAdBlockWhitelist: (...args: unknown[]) => mockRemoveAdBlockWhitelist(...args), + setAdBlockRefreshInterval: (...args: unknown[]) => mockSetAdBlockRefreshInterval(...args), + }; +}); + +vi.mock("../../hooks/useWebKitPointerDown", () => ({ + useWebKitPointerDown: () => ({ onPointerDown: () => () => {} }), +})); + +import AdBlock from "../AdBlock"; + +function makeState(overrides: Partial = {}): AdBlockState { + return { + enabled: true, + sources: [], + whitelist: [], + auto_refresh_enabled: true, + refresh_interval_hours: 24, + ...overrides, + }; +} + +function makeSource(overrides: Partial = {}): AdBlockSource { + return { + source_id: "src-1", + name: "Test List", + url: "https://example.com/hosts", + enabled: true, + response: "zero_address", + last_fetched_at: null, + last_error: null, + rule_count: 100, + etag: null, + ...overrides, + }; +} + +function renderWithProviders(ui: React.ReactElement) { + return render( + + {ui} + , + ); +} + +function setStore(fn: (s: ReturnType) => void) { + const store = getDefaultStore(); + fn(store); + return store; +} + +describe("AdBlock", () => { + beforeEach(() => { + vi.clearAllMocks(); + setStore((s) => { + s.set(adBlockStateAtom, null); + s.set(isAdBlockLoadingAtom, false); + s.set(adBlockErrorAtom, null); + s.set(dnsEnabledAtom, false); + }); + mockGetAdBlockState.mockResolvedValue(makeState()); + }); + + // ---- issue #134: loading state ---- + it("renders Loading when state is null", () => { + renderWithProviders(); + expect(screen.getByText("Loading\u2026")).toBeInTheDocument(); + }); + + it("renders page title and subtitle after state loads", async () => { + const state = makeState(); + mockGetAdBlockState.mockResolvedValue(state); + setStore((s) => s.set(adBlockStateAtom, state)); + renderWithProviders(); + expect(await screen.findByText("Ad Block")).toBeInTheDocument(); + expect( + screen.getByText("Block ads at the DNS resolver. macOS DNS mode only."), + ).toBeInTheDocument(); + }); + + // ---- issue #134: DNS-off banner ---- + it("shows DNS-off banner when dnsEnabled is false", async () => { + const state = makeState(); + setStore((s) => { + s.set(adBlockStateAtom, state); + s.set(dnsEnabledAtom, false); + }); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText(/DNS mode is off/i)).toBeInTheDocument(); + }); + + it("hides DNS-off banner when dnsEnabled is true", async () => { + const state = makeState(); + setStore((s) => { + s.set(adBlockStateAtom, state); + s.set(dnsEnabledAtom, true); + }); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("Ad Block")).toBeInTheDocument(); + expect(screen.queryByText(/DNS mode is off/i)).not.toBeInTheDocument(); + }); + + // ---- issue #134: empty state ---- + it("renders empty-source placeholder when sources is empty", async () => { + const state = makeState({ sources: [] }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("No sources yet.")).toBeInTheDocument(); + }); + + // ---- issue #134: source list rendering ---- + it("renders source cards with name, url, and rule count", async () => { + const src = makeSource({ + name: "StevenBlack", + url: "https://sb.com/hosts", + rule_count: 5000, + }); + const state = makeState({ sources: [src] }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("StevenBlack")).toBeInTheDocument(); + expect(screen.getByText("https://sb.com/hosts")).toBeInTheDocument(); + expect(screen.getByText(/5,000 rules/)).toBeInTheDocument(); + }); + + // ---- issue #134: error badge on source ---- + it("renders fetch-failed badge when source has last_error", async () => { + const src = makeSource({ last_error: "timeout" }); + const state = makeState({ sources: [src] }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("fetch failed")).toBeInTheDocument(); + expect(screen.getByText(/timeout/)).toBeInTheDocument(); + }); + + // ---- issue #134: master switch ---- + it("toggling master switch calls setAdBlockEnabled", async () => { + // Mock the initial fetch to return enabled=false (same as pre-set + // state). The toggle action re-fetches after setAdBlockEnabled, but + // we only assert the IPC call here. + const state = makeState({ enabled: false }); + mockGetAdBlockState.mockResolvedValue(state); + setStore((s) => s.set(adBlockStateAtom, state)); + renderWithProviders(); + await screen.findByText("Ad Block"); + const checkboxes = screen.getAllByRole("checkbox"); + const masterSwitch = checkboxes[0]; + expect(masterSwitch).not.toBeChecked(); + await act(async () => { + fireEvent.click(masterSwitch); + }); + expect(mockSetAdBlockEnabled).toHaveBeenCalledWith(true); + }); + + // ---- issue #134: add source form ---- + it("fills add-source form and clicks Add", async () => { + const state = makeState(); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + await screen.findByRole("heading", { name: "Sources" }); + const nameInput = screen.getByPlaceholderText("StevenBlack"); + const urlInput = screen.getByPlaceholderText("https://example.com/hosts"); + await act(async () => { + fireEvent.change(nameInput, { target: { value: "MyList" } }); + fireEvent.change(urlInput, { target: { value: "https://ml.com/hosts" } }); + }); + // Two "Add" buttons exist (sources + whitelist). The source form's + // Add button is the first one. + const addBtns = screen.getAllByText("Add"); + const sourceAddBtn = addBtns[0]; + await act(async () => { + fireEvent.click(sourceAddBtn); + }); + expect(mockAddAdBlockSource).toHaveBeenCalledWith( + "MyList", + "https://ml.com/hosts", + "zero_address", + ); + }); + + it("source Add button is disabled when name or url is empty", async () => { + const state = makeState(); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + await screen.findByRole("heading", { name: "Sources" }); + const addBtns = screen.getAllByText("Add"); + expect(addBtns[0]).toBeDisabled(); + }); + + // ---- issue #134: whitelist ---- + it("renders whitelist entries with remove buttons", async () => { + const state = makeState({ whitelist: ["trusted.com", "safe.com"] }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("trusted.com")).toBeInTheDocument(); + expect(screen.getByText("safe.com")).toBeInTheDocument(); + expect(screen.getByLabelText("Remove trusted.com")).toBeInTheDocument(); + }); + + it("renders empty whitelist placeholder when no entries", async () => { + const state = makeState({ whitelist: [] }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByText("No whitelist entries.")).toBeInTheDocument(); + }); + + it("adds a whitelist domain via the input", async () => { + const state = makeState(); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + await screen.findByRole("heading", { name: "Whitelist" }); + const input = screen.getByPlaceholderText("trusted.example.com"); + await act(async () => { + fireEvent.change(input, { target: { value: "new.com" } }); + }); + // The whitelist Add button is the second one. + const addBtns = screen.getAllByText("Add"); + const whitelistAddBtn = addBtns[1]; + await act(async () => { + fireEvent.click(whitelistAddBtn); + }); + expect(mockAddAdBlockWhitelist).toHaveBeenCalledWith("new.com"); + }); + + // ---- issue #134: error alert ---- + it("renders error alert when adBlockErrorAtom is set", async () => { + // fetchAdBlockStateAtom clears adBlockErrorAtom to null on mount + // (it's the "begin a fetch" signal). Set the error AFTER the + // initial render so the effect's clear doesn't wipe it. + const state = makeState(); + mockGetAdBlockState.mockResolvedValue(state); + setStore((s) => s.set(adBlockStateAtom, state)); + renderWithProviders(); + await screen.findByText("Ad Block"); + setStore((s) => s.set(adBlockErrorAtom, "Something went wrong")); + expect(await screen.findByText("Something went wrong")).toBeInTheDocument(); + }); + + // ---- issue #134: summary stats ---- + it("renders summary stat labels for sources, rules, and whitelist", async () => { + const state = makeState({ + sources: [ + makeSource({ rule_count: 100 }), + makeSource({ source_id: "src-2", rule_count: 200, enabled: false }), + ], + whitelist: ["a.com", "b.com"], + }); + setStore((s) => s.set(adBlockStateAtom, state)); + mockGetAdBlockState.mockResolvedValue(state); + renderWithProviders(); + expect(await screen.findByRole("heading", { name: "Sources" })).toBeInTheDocument(); + expect(screen.getByRole("heading", { name: "Whitelist" })).toBeInTheDocument(); + }); +}); diff --git a/src/stores/profiles/actions.ts b/src/stores/profiles/actions.ts index 6e9b838..3dbcb3e 100644 --- a/src/stores/profiles/actions.ts +++ b/src/stores/profiles/actions.ts @@ -1,5 +1,5 @@ import { atom } from "jotai"; -import type { Profile } from "../../types"; +import type { Profile, AdBlockResponse } from "../../types"; import { listProfiles, getProfile, @@ -19,6 +19,19 @@ import { setDnsMode, reloadDnsRules, listDnsProfiles, + getAdBlockState, + setAdBlockEnabled, + setAdBlockRefreshInterval, + listAdBlockSources, + addAdBlockSource, + removeAdBlockSource, + setAdBlockSourceEnabled, + setAdBlockSourceResponse, + refreshAdBlockSource, + refreshAllAdBlockSources, + listAdBlockWhitelist, + addAdBlockWhitelist, + removeAdBlockWhitelist, } from "../../lib/tauri"; import { extractErrorMessage } from "../../lib/error"; import { @@ -40,6 +53,9 @@ import { dnsStatusAtom, isDnsLoadingAtom, dnsErrorAtom, + adBlockStateAtom, + isAdBlockLoadingAtom, + adBlockErrorAtom, } from "./state"; // ---- Async action atoms ---- @@ -383,3 +399,188 @@ export const toggleDnsProfileEnabledAtom = atom( } }, ); +// 每个 mutating action 在 await 成功后立刻拉一次 `getAdBlockState` 重写 +// `adBlockStateAtom`,避免前端手动维护 sources 列表的不变量。读操作的 +// fetch(mount 时 / 路由进入时)走 `fetchAdBlockStateAtom`。 + +export const fetchAdBlockStateAtom = atom(null, async (_get, set) => { + set(isAdBlockLoadingAtom, true); + set(adBlockErrorAtom, null); + try { + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } finally { + set(isAdBlockLoadingAtom, false); + } +}); + +export const toggleAdBlockEnabledAtom = atom( + null, + async (_get, set, enabled: boolean) => { + set(isAdBlockLoadingAtom, true); + set(adBlockErrorAtom, null); + try { + await setAdBlockEnabled(enabled); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } finally { + set(isAdBlockLoadingAtom, false); + } + }, +); + +export const setAdBlockIntervalAtom = atom( + null, + async (_get, set, hours: number) => { + set(adBlockErrorAtom, null); + try { + await setAdBlockRefreshInterval(hours); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +export const addAdBlockSourceAtom = atom( + null, + async (_get, set, args: { name: string; url: string; response: AdBlockResponse }) => { + set(isAdBlockLoadingAtom, true); + set(adBlockErrorAtom, null); + try { + await addAdBlockSource(args.name, args.url, args.response); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } finally { + set(isAdBlockLoadingAtom, false); + } + }, +); + +export const removeAdBlockSourceAtom = atom( + null, + async (_get, set, sourceId: string) => { + set(adBlockErrorAtom, null); + try { + await removeAdBlockSource(sourceId); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +export const setAdBlockSourceEnabledAtom = atom( + null, + async (_get, set, args: { sourceId: string; enabled: boolean }) => { + set(adBlockErrorAtom, null); + try { + await setAdBlockSourceEnabled(args.sourceId, args.enabled); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +export const setAdBlockSourceResponseAtom = atom( + null, + async (_get, set, args: { sourceId: string; response: AdBlockResponse }) => { + set(adBlockErrorAtom, null); + try { + await setAdBlockSourceResponse(args.sourceId, args.response); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +export const refreshAdBlockSourceAtom = atom( + null, + async (_get, set, sourceId: string) => { + set(isAdBlockLoadingAtom, true); + set(adBlockErrorAtom, null); + try { + await refreshAdBlockSource(sourceId); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } finally { + set(isAdBlockLoadingAtom, false); + } + }, +); + +export const refreshAllAdBlockSourcesAtom = atom(null, async (_get, set) => { + set(isAdBlockLoadingAtom, true); + set(adBlockErrorAtom, null); + try { + await refreshAllAdBlockSources(); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } finally { + set(isAdBlockLoadingAtom, false); + } +}); + +export const addAdBlockWhitelistAtom = atom( + null, + async (_get, set, domain: string) => { + set(adBlockErrorAtom, null); + try { + await addAdBlockWhitelist(domain); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +export const removeAdBlockWhitelistAtom = atom( + null, + async (_get, set, domain: string) => { + set(adBlockErrorAtom, null); + try { + await removeAdBlockWhitelist(domain); + const state = await getAdBlockState(); + set(adBlockStateAtom, state); + } catch (err) { + set(adBlockErrorAtom, extractErrorMessage(err)); + throw err; + } + }, +); + +// Re-export whitelist helpers that don't need state refresh. +export const fetchAdBlockWhitelistAtom = atom(null, async () => { + return listAdBlockWhitelist(); +}); + +export const fetchAdBlockSourcesAtom = atom(null, async () => { + return listAdBlockSources(); +}); diff --git a/src/stores/profiles/index.ts b/src/stores/profiles/index.ts index be01c1c..e4ad47c 100644 --- a/src/stores/profiles/index.ts +++ b/src/stores/profiles/index.ts @@ -22,6 +22,11 @@ export { dnsErrorAtom, enabledDnsProfilesAtom, dnsRuleCountAtom, + adBlockStateAtom, + isAdBlockLoadingAtom, + adBlockErrorAtom, + adBlockRuleCountAtom, + adBlockHasErrorsAtom, } from "./state"; // ---- Async action atoms ---- @@ -47,4 +52,17 @@ export { updateDnsProfileAtom, deleteDnsProfileAtom, toggleDnsProfileEnabledAtom, + fetchAdBlockStateAtom, + toggleAdBlockEnabledAtom, + setAdBlockIntervalAtom, + addAdBlockSourceAtom, + removeAdBlockSourceAtom, + setAdBlockSourceEnabledAtom, + setAdBlockSourceResponseAtom, + refreshAdBlockSourceAtom, + refreshAllAdBlockSourcesAtom, + addAdBlockWhitelistAtom, + removeAdBlockWhitelistAtom, + fetchAdBlockSourcesAtom, + fetchAdBlockWhitelistAtom, } from "./actions"; diff --git a/src/stores/profiles/state.ts b/src/stores/profiles/state.ts index 1af4542..b3c1df3 100644 --- a/src/stores/profiles/state.ts +++ b/src/stores/profiles/state.ts @@ -1,5 +1,5 @@ import { atom } from "jotai"; -import type { Profile, DnsStatus } from "../../types"; +import type { Profile, DnsStatus, AdBlockState } from "../../types"; import { countRealRules } from "../../lib/rules"; // ---- Base atoms ---- @@ -60,3 +60,23 @@ export const applyTargetAtom = atom<{ id: string; enabled: boolean } | null>(nul export const snapshotsAtom = atom([]); export const isLoadingSnapshotsAtom = atom(false); export const snapshotErrorAtom = atom(null); + +// ---- AdBlock atoms (issue #130) ---- + +export const adBlockStateAtom = atom(null); +export const isAdBlockLoadingAtom = atom(false); +export const adBlockErrorAtom = atom(null); + +export const adBlockRuleCountAtom = atom((get) => { + const state = get(adBlockStateAtom); + if (!state) return 0; + return state.sources + .filter((s) => s.enabled) + .reduce((sum, s) => sum + s.rule_count, 0); +}); + +export const adBlockHasErrorsAtom = atom((get) => { + const state = get(adBlockStateAtom); + if (!state) return false; + return state.sources.some((s) => s.last_error !== null); +}); diff --git a/src/types/index.ts b/src/types/index.ts index 006124a..7cadb1a 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -45,6 +45,36 @@ export type RuleSource = | { type: "Remote"; source_id: string; source_name: string } | { type: "AdBlock"; source_id: string; source_name: string }; +// --------------------------------------------------------------------------- +// AdBlock (issue #130) +// +// Mirrors `mhost_core::AdBlockResponse` / `AdBlockSource` / `AdBlockState`. +// Keep wire format in sync: snake_case from serde rename_all = "snake_case" +// gives us "zero_address" and "nx_domain" on the wire. +// --------------------------------------------------------------------------- + +export type AdBlockResponse = "zero_address" | "nx_domain"; + +export interface AdBlockSource { + source_id: string; + name: string; + url: string; + enabled: boolean; + response: AdBlockResponse; + last_fetched_at: string | null; + last_error: string | null; + rule_count: number; + etag: string | null; +} + +export interface AdBlockState { + enabled: boolean; + sources: AdBlockSource[]; + whitelist: string[]; + auto_refresh_enabled: boolean; + refresh_interval_hours: number; +} + export interface ApplyPlan { rules: ResolvedRule[]; conflicts: RuleConflict[];