| use adblock::{Engine, FilterSet};
|
| use adblock::lists::ParseOptions;
|
| use adblock::request::Request;
|
| use std::collections::HashSet;
|
| use std::sync::RwLock;
|
| use std::sync::atomic::{AtomicU64, Ordering};
|
| use serde::Serialize;
|
|
|
|
|
| pub struct AdBlockState {
|
| pub engine: RwLock<Engine>,
|
| pub stats: AdBlockStats,
|
| pub allowlist: RwLock<HashSet<String>>,
|
| pub rule_count: AtomicU64,
|
| }
|
|
|
| #[derive(Default)]
|
| pub struct AdBlockStats {
|
| pub blocked_requests: AtomicU64,
|
| pub blocked_cosmetic: AtomicU64,
|
| pub https_upgrades: AtomicU64,
|
| }
|
|
|
| #[derive(Debug, Clone, Serialize)]
|
| pub struct ShieldReport {
|
| pub blocked_requests: u64,
|
| pub blocked_cosmetic: u64,
|
| pub https_upgrades: u64,
|
| pub engine_rules: usize,
|
| pub allowlisted_domains: usize,
|
| }
|
|
|
| impl AdBlockState {
|
| pub fn new() -> Self {
|
| let (engine, count) = build_engine_from_bundled();
|
| Self {
|
| engine: RwLock::new(engine),
|
| stats: AdBlockStats::default(),
|
| allowlist: RwLock::new(HashSet::new()),
|
| rule_count: AtomicU64::new(count as u64),
|
| }
|
| }
|
|
|
| pub fn should_block(&self, url: &str, source_url: &str, request_type: &str) -> bool {
|
| if let Some(domain) = extract_domain(source_url) {
|
| if let Ok(allowlist) = self.allowlist.read() {
|
| if allowlist.contains(&domain) {
|
| return false;
|
| }
|
| }
|
| }
|
|
|
| let engine = match self.engine.read() {
|
| Ok(e) => e,
|
| Err(_) => return false,
|
| };
|
|
|
| match Request::new(url, source_url, request_type) {
|
| Ok(req) => {
|
| let result = engine.check_network_request(&req);
|
| if result.matched && result.exception.is_none() {
|
| self.stats.blocked_requests.fetch_add(1, Ordering::Relaxed);
|
| true
|
| } else {
|
| false
|
| }
|
| }
|
| Err(_) => false,
|
| }
|
| }
|
|
|
| pub fn get_cosmetic_css(&self, url: &str) -> String {
|
| let engine = match self.engine.read() {
|
| Ok(e) => e,
|
| Err(_) => return String::new(),
|
| };
|
|
|
| let resources = engine.url_cosmetic_resources(url);
|
| if resources.hide_selectors.is_empty() {
|
| return String::new();
|
| }
|
|
|
| let selectors: Vec<&String> = resources.hide_selectors.iter().collect();
|
| let css = selectors.iter()
|
| .map(|s| format!("{s}{{display:none!important}}"))
|
| .collect::<Vec<_>>()
|
| .join("\n");
|
|
|
| self.stats.blocked_cosmetic.fetch_add(selectors.len() as u64, Ordering::Relaxed);
|
| css
|
| }
|
|
|
|
|
|
|
| pub fn get_injected_script(&self, url: &str) -> String {
|
| let engine = match self.engine.read() {
|
| Ok(e) => e,
|
| Err(_) => return String::new(),
|
| };
|
| let resources = engine.url_cosmetic_resources(url);
|
| resources.injected_script
|
| }
|
|
|
| pub fn report(&self) -> ShieldReport {
|
| let engine_rules = self.rule_count.load(Ordering::Relaxed) as usize;
|
| let allowlisted = self.allowlist.read().map(|a| a.len()).unwrap_or(0);
|
| ShieldReport {
|
| blocked_requests: self.stats.blocked_requests.load(Ordering::Relaxed),
|
| blocked_cosmetic: self.stats.blocked_cosmetic.load(Ordering::Relaxed),
|
| https_upgrades: self.stats.https_upgrades.load(Ordering::Relaxed),
|
| engine_rules,
|
| allowlisted_domains: allowlisted,
|
| }
|
| }
|
|
|
| pub fn reload_engine(&self, new_engine: Engine, rule_count: usize) {
|
| if let Ok(mut engine) = self.engine.write() {
|
| *engine = new_engine;
|
| }
|
| self.rule_count.store(rule_count as u64, Ordering::Relaxed);
|
| }
|
| }
|
|
|
| impl Default for AdBlockState {
|
| fn default() -> Self {
|
| Self::new()
|
| }
|
| }
|
|
|
| pub fn build_engine_from_bundled() -> (Engine, usize) {
|
| let mut filter_set = FilterSet::new(false);
|
|
|
| let lists: &[&str] = &[
|
| include_str!("../../resources/filters/easylist_mini.txt"),
|
| include_str!("../../resources/filters/easyprivacy_mini.txt"),
|
| include_str!("../../resources/filters/annoyances_mini.txt"),
|
| ];
|
|
|
| let mut count = 0usize;
|
| for list in lists {
|
| let _meta = filter_set.add_filter_list(list, ParseOptions::default());
|
| count += list.lines().filter(|l| !l.starts_with('!') && !l.trim().is_empty()).count();
|
| }
|
|
|
| (Engine::from_filter_set(filter_set, true), count)
|
| }
|
|
|
| fn extract_domain(url: &str) -> Option<String> {
|
| url.split("//")
|
| .nth(1)
|
| .and_then(|s| s.split('/').next())
|
| .map(|s| s.split(':').next().unwrap_or(s).to_lowercase())
|
| .map(|s| s.trim_start_matches("www.").to_string())
|
| }
|
|
|
