//! Свой DNS-резолвер на 127.0.0.1:53.
//!
//! Файл hosts не умеет ни масок, ни журнала: `discord.com` пришлось бы
//! перечислять вместе со всеми поддоменами, а узнать, что ребёнок ломился в
//! запрещённое, было бы неоткуда. Поэтому имена разрешает служба: адаптеры
//! смотрят на 127.0.0.1, запрещённое отсекается здесь же, остальное уходит
//! наверх.
//!
//! Резолвер стоит на пути каждого соединения в системе, так что цена ошибки —
//! не «медленно», а «интернета нет вообще». Отсюда три решения, определяющие
//! устройство файла:
//!
//! * Узкое место здесь не пропускная способность (домашняя машина даёт десятки
//!   запросов в секунду), а задержка. Поэтому запрос к вышестоящему серверу
//!   никогда никого не ждёт: один сокет на всех, свой идентификатор в каждом
//!   запросе и таблица незакрытых — как в dnsmasq. Задачи на запрос не
//!   заводятся, сокеты на запрос не открываются.
//! * Запрещённое имя и попадание в кэш обслуживаются без единой аллокации и без
//!   выхода в сеть — весь ответ собирается прямо в приёмном буфере.
//! * Всё, что резолвер не смог разобрать, уходит наверх как есть. Строгий
//!   разбор — причина отказать в блокировке, но не причина сломать разрешение
//!   имён.

use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::{
        atomic::{AtomicU16, AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use serde::Serialize;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream, UdpSocket},
    sync::watch,
};

use crate::Shared;

/// Слушаем только петлю: наружу резолвер отдавать нечего, а открытый DNS в
/// локальной сети — приглашение использовать машину как усилитель атак.
pub const BIND_ADDR: &str = "127.0.0.1:53";
/// Тот же резолвер по IPv6. Оставить шестую версию без присмотра нельзя:
/// Windows предпочитает её адрес сервера, и запросы ушли бы мимо нас — то есть
/// блокировки бы просто не было.
pub const BIND_ADDR_V6: &str = "[::1]:53";
/// Предельный размер сообщения DNS — столько адресует поле длины у варианта
/// поверх TCP.
///
/// Буфер сделан по этому пределу, а не по обычным для EDNS0 четырём килобайтам,
/// из-за поведения приёма датаграмм: то, что не влезло, отбрасывается молча, и
/// признака обрезки не остаётся. Клиенту ушёл бы обломок пакета — то есть имя
/// переставало бы разрешаться без всякого объяснения. Буфер выделяется один раз
/// на задачу, поэтому запас ничего не стоит.
const BUF: usize = 65535;
/// Заголовок DNS: идентификатор, флаги и четыре счётчика секций.
const HEADER: usize = 12;
/// Предел длины имени по RFC 1035 вместе с точками.
const MAX_NAME: usize = 255;
/// Сколько места нужно под приписываемый ответ: указатель сжатия, тип, класс,
/// TTL, длина и шестнадцать байт адреса IPv6.
const ANSWER_MAX: usize = 2 + 2 + 2 + 4 + 2 + 16;

const TYPE_A: u16 = 1;
const TYPE_AAAA: u16 = 28;
const CLASS_IN: u16 = 1;

/// TTL ответа-заглушки. Ноль означал бы «не кэшировать», и клиент спрашивал бы
/// заново на каждое соединение; минута же не даёт разблокировке залежаться.
const BLOCK_TTL: u32 = 60;
/// Сколько ждать вышестоящий сервер, прежде чем считать запрос потерянным.
const UPSTREAM_TIMEOUT: Duration = Duration::from_secs(3);
/// Потолок незакрытых запросов: без него шторм из localhost раздул бы таблицу.
const MAX_INFLIGHT: usize = 2048;
const CACHE_CAP: usize = 4096;
/// Свой кэш живёт недолго: TTL записей мы не пересчитываем, а отдаём как есть,
/// поэтому долгое хранение возвращало бы клиенту заведомо неверный остаток
/// жизни. Тридцати секунд хватает, чтобы погасить всплеск от загрузки страницы.
const CACHE_TTL: Duration = Duration::from_secs(30);
/// Отрицательные ответы держим ещё меньше: домен мог только что появиться.
const NEG_TTL: Duration = Duration::from_secs(10);
/// Как часто чистить таблицу незакрытых и сбрасывать журнал попыток.
const SWEEP_INTERVAL: Duration = Duration::from_secs(1);
const REPORT_INTERVAL: Duration = Duration::from_secs(30);
/// После скольких подряд молчаний переходить на запасной сервер.
const FAILOVER_AFTER: u64 = 5;
/// Сколько доменов помнить для панели.
const STATS_CAP: usize = 256;
/// Сколько последних запросов держать для панели. Двести — это примерно
/// полминуты жизни машины под нагрузкой и несколько часов простоя, то есть
/// ровно то окно, в котором ещё имеет смысл спрашивать «а что сейчас было».
const RECENT_CAP: usize = 200;
/// Потолок одновременных TCP-запросов. По TCP ходят единицы клиентов, но
/// принимать их без счёта — способ остаться без дескрипторов.
const MAX_TCP: usize = 32;
/// Оборванное TCP-соединение не должно занимать слот вечно.
const TCP_IDLE: Duration = Duration::from_secs(10);

/* ── список запрещённого ───────────────────────────────────── */

/// Домены в нижнем регистре, без точки на конце.
///
/// Совпадение ищется по суффиксу: запись `discord.com` закрывает и сам домен, и
/// любые поддомены — иначе список пришлось бы вести вручную, а Discord ходит
/// через десяток имён вида `gateway.discord.gg`.
#[derive(Default)]
pub struct Blocklist {
    domains: std::collections::HashSet<Box<str>>,
}

impl Blocklist {
    pub fn new<I: IntoIterator<Item = String>>(domains: I) -> Self {
        Self {
            domains: domains
                .into_iter()
                .filter_map(|d| normalize_domain(&d))
                .map(String::into_boxed_str)
                .collect(),
        }
    }

    /// Имя уже в нижнем регистре — так его кладёт `parse_question`.
    fn matches(&self, name: &str) -> bool {
        if self.domains.is_empty() {
            return false;
        }
        let mut tail = name;
        loop {
            if self.domains.contains(tail) {
                return true;
            }
            match tail.split_once('.') {
                // До последней метки не доходим намеренно: запись вида `com`
                // (опечатка или лишняя строка) иначе отрезала бы пол-интернета.
                Some((_, rest)) if rest.contains('.') => tail = rest,
                _ => return false,
            }
        }
    }
}

/// Приводит то, что человек ввёл в панели, к имени домена: люди вставляют
/// адрес целиком, вместе со схемой, путём и `www`.
pub fn normalize_domain(raw: &str) -> Option<String> {
    let lower = raw.trim().to_ascii_lowercase();
    let host = lower.split_once("://").map_or(lower.as_str(), |(_, r)| r);
    let host = host.split(['/', '?', '#']).next()?;
    // Имя пользователя перед адресом встречается в скопированных ссылках.
    let host = host.rsplit_once('@').map_or(host, |(_, h)| h);
    let host = host.split(':').next()?.trim_matches('.');
    // `www` отбрасываем: запись `www.discord.com` иначе не закрыла бы сам
    // `discord.com`, хотя человек имел в виду именно его.
    let host = host.strip_prefix("www.").unwrap_or(host);

    let plausible = host.contains('.')
        && host.len() <= MAX_NAME
        && host.chars().any(|c| c.is_ascii_alphabetic())
        && host
            .split('.')
            .all(|label| !label.is_empty() && label.len() <= 63)
        && host
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'.' || b == b'-');

    plausible.then(|| host.to_string())
}

/* ── разбор пакета ─────────────────────────────────────────── */

struct Question {
    /// Длина имени, записанного в переданный вызовом буфер.
    len: usize,
    qtype: u16,
    /// Класс запроса. Почти всегда `IN`, но различать обязаны: без него записи
    /// разных классов с одним именем схлопывались бы в кэше в одну.
    qclass: u16,
    /// Смещение сразу за секцией вопроса — оттуда дописывается ответ.
    end: usize,
}

/// Достаёт имя и тип из запроса. Имя кладётся в `out` в нижнем регистре.
///
/// Разбираем только то, что нужно для решения «блокировать или нет»: имя, тип и
/// границу вопроса. Всё остальное — секции ответа, EDNS0 — не трогаем, и запрос
/// уходит наверх байт в байт, вместе с объявленным размером буфера клиента.
fn parse_question(pkt: &[u8], out: &mut [u8; MAX_NAME]) -> Option<Question> {
    if pkt.len() < HEADER {
        return None;
    }
    // Ответ на входе означает либо чужой пакет, либо попытку нас запутать.
    if pkt[2] & 0x80 != 0 {
        return None;
    }
    // Стандартный запрос (opcode 0) ровно с одним вопросом. Всё прочее —
    // обновления зон, запросы состояния — пусть разбирает вышестоящий сервер.
    if pkt[2] & 0x78 != 0 || u16::from_be_bytes([pkt[4], pkt[5]]) != 1 {
        return None;
    }

    let mut pos = HEADER;
    let mut len = 0usize;
    loop {
        let label = *pkt.get(pos)? as usize;
        pos += 1;
        if label == 0 {
            break;
        }
        // Сжатие в вопросе запрещено: указатель здесь либо от сломанного
        // клиента, либо приглашение зациклить разбор.
        if label > 63 {
            return None;
        }
        let need = if len == 0 { label } else { label + 1 };
        if len + need > MAX_NAME {
            return None;
        }
        let bytes = pkt.get(pos..pos + label)?;
        if len > 0 {
            out[len] = b'.';
            len += 1;
        }
        for &b in bytes {
            out[len] = b.to_ascii_lowercase();
            len += 1;
        }
        pos += label;
    }

    let qtype = u16::from_be_bytes([*pkt.get(pos)?, *pkt.get(pos + 1)?]);
    let qclass = u16::from_be_bytes([*pkt.get(pos + 2)?, *pkt.get(pos + 3)?]);
    let end = pos + 4;
    (end <= pkt.len() && len > 0).then_some(Question {
        len,
        qtype,
        qclass,
        end,
    })
}

/// Переделывает запрос в ответ «такого адреса нет» прямо в буфере запроса и
/// возвращает новую длину пакета.
///
/// Отдаём 0.0.0.0, а не NXDOMAIN: подключение к нему обрывается сразу и браузер
/// показывает понятную ошибку, тогда как на «домена не существует» часть
/// программ уходит переспрашивать по кругу.
fn write_blocked(pkt: &mut [u8], q: &Question) -> Option<usize> {
    if q.end + ANSWER_MAX > pkt.len() {
        return None;
    }

    pkt[2] |= 0x80; // QR: это ответ
    pkt[2] &= !0x06; // снимаем AA и TC — ответ наш и целиком помещается
    pkt[3] = 0x80; // рекурсия доступна, код ответа нулевой

    let rdata: &[u8] = match (q.qclass, q.qtype) {
        (CLASS_IN, TYPE_A) => &[0, 0, 0, 0],
        (CLASS_IN, TYPE_AAAA) => &[0; 16],
        // На прочие типы отвечаем пустым успехом: подделывать MX или TXT незачем,
        // а отказ отправил бы клиента спрашивать снова. Чужой класс сюда же —
        // выдавать адрес из класса IN в ответ на вопрос про другой было бы
        // прямой неправдой.
        _ => &[],
    };

    let mut n = q.end;
    if !rdata.is_empty() {
        // Имя не повторяем: указатель сжатия на вопрос, он всегда со смещения 12.
        pkt[n..n + 2].copy_from_slice(&[0xc0, 0x0c]);
        pkt[n + 2..n + 4].copy_from_slice(&q.qtype.to_be_bytes());
        pkt[n + 4..n + 6].copy_from_slice(&q.qclass.to_be_bytes());
        pkt[n + 6..n + 10].copy_from_slice(&BLOCK_TTL.to_be_bytes());
        pkt[n + 10..n + 12].copy_from_slice(&(rdata.len() as u16).to_be_bytes());
        pkt[n + 12..n + 12 + rdata.len()].copy_from_slice(rdata);
        n += 12 + rdata.len();
    }

    pkt[6..8].copy_from_slice(&(if rdata.is_empty() { 0u16 } else { 1 }).to_be_bytes());
    // Секции полномочий и дополнений отбрасываем вместе с EDNS0 клиента:
    // заглушке они не нужны, а пересобирать их — лишняя работа на горячем пути.
    pkt[8..12].copy_from_slice(&[0, 0, 0, 0]);
    Some(n)
}

fn packet_id(pkt: &[u8]) -> u16 {
    u16::from_be_bytes([pkt[0], pkt[1]])
}

fn set_packet_id(pkt: &mut [u8], id: u16) {
    pkt[0..2].copy_from_slice(&id.to_be_bytes());
}

fn rcode(pkt: &[u8]) -> u8 {
    pkt[3] & 0x0f
}

fn truncated(pkt: &[u8]) -> bool {
    pkt[2] & 0x02 != 0
}

/* ── кэш ───────────────────────────────────────────────────── */

struct Entry {
    /// Готовый ответ целиком. Идентификатор в первых двух байтах у каждого
    /// клиента свой и подставляется при отправке.
    packet: Box<[u8]>,
    expires: Instant,
}

/// Ключ — имя вместе с типом и классом запроса, собранный в стековом буфере:
/// искать в карте можно по срезу, не выделяя памяти на каждый запрос.
fn cache_key(name: &[u8], qtype: u16, qclass: u16, out: &mut [u8; MAX_NAME + 4]) -> usize {
    let end = name.len();
    out[..end].copy_from_slice(name);
    out[end..end + 2].copy_from_slice(&qtype.to_be_bytes());
    out[end + 2..end + 4].copy_from_slice(&qclass.to_be_bytes());
    end + 4
}

#[derive(Default)]
struct Cache {
    map: HashMap<Box<[u8]>, Entry>,
}

impl Cache {
    /// Копирует ответ в `out` и проставляет идентификатор клиента.
    fn get(&self, key: &[u8], id: u16, out: &mut [u8], now: Instant) -> Option<usize> {
        let entry = self.map.get(key)?;
        if entry.expires <= now || entry.packet.len() > out.len() {
            return None;
        }
        out[..entry.packet.len()].copy_from_slice(&entry.packet);
        set_packet_id(out, id);
        Some(entry.packet.len())
    }

    fn put(&mut self, key: &[u8], packet: &[u8], now: Instant) {
        // Обрезанные ответы бессмысленны в кэше: клиент всё равно переспросит
        // по TCP. Ошибки сервера не кэшируем, чтобы сбой не залипал.
        if truncated(packet) {
            return;
        }
        let ttl = match rcode(packet) {
            0 => CACHE_TTL,
            3 => NEG_TTL,
            _ => return,
        };

        if self.map.len() >= CACHE_CAP {
            // Сначала выбрасываем протухшее, и лишь если это не помогло —
            // чистим целиком. Учёт порядка обращений стоил бы дороже, чем
            // редкий поход наверх за тем, что мы только что забыли.
            self.map.retain(|_, e| e.expires > now);
            if self.map.len() >= CACHE_CAP {
                self.map.clear();
            }
        }

        self.map.insert(
            key.into(),
            Entry {
                packet: packet.into(),
                expires: now + ttl,
            },
        );
    }

    fn sweep(&mut self, now: Instant) {
        self.map.retain(|_, e| e.expires > now);
    }
}

/* ── учёт попыток ──────────────────────────────────────────── */

struct Hit {
    count: u64,
    /// Сколько уже ушло в журнал: пишем приростом, а не полным счётчиком.
    reported: u64,
    last: Instant,
}

#[derive(Serialize)]
pub struct BlockedDomain {
    pub name: String,
    pub count: u64,
    pub ago_secs: u64,
}

#[derive(Default)]
struct Attempts {
    domains: HashMap<Box<str>, Hit>,
}

impl Attempts {
    fn record(&mut self, name: &str, now: Instant) {
        if let Some(hit) = self.domains.get_mut(name) {
            hit.count += 1;
            hit.last = now;
            return;
        }
        if self.domains.len() >= STATS_CAP {
            // Вытесняем то, к чему дольше всего не обращались: панели интересно
            // происходящее сейчас, а не первые попытки полугодовой давности.
            if let Some(oldest) = self
                .domains
                .iter()
                .min_by_key(|(_, h)| h.last)
                .map(|(k, _)| k.clone())
            {
                self.domains.remove(&oldest);
            }
        }
        self.domains.insert(
            name.into(),
            Hit {
                count: 1,
                reported: 0,
                last: now,
            },
        );
    }
}

/* ── лента последних запросов ──────────────────────────────── */

/// Что резолвер сделал с запросом.
#[derive(Clone, Copy, Debug, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Verdict {
    /// Отсечён по списку, наверх не пошёл.
    Blocked,
    /// Ответ взят из своего кэша.
    Cached,
    /// Ушёл вышестоящему серверу.
    Upstream,
}

/// Одна запись ленты.
///
/// Имя лежит прямо в структуре массивом, а не отдельной строкой в куче. Кольцо
/// заводится один раз при запуске, и запись в него не трогает распределитель
/// памяти — иначе на каждый запрос в системе приходилось бы выделение и
/// освобождение, а горячий путь резолвера построен ровно на том, чтобы их не
/// было. Ценой идут постоянные 56 КБ, что на фоне кэша ответов незаметно.
struct Seen {
    name: [u8; MAX_NAME],
    len: u8,
    qtype: u16,
    verdict: Verdict,
    at: Instant,
}

#[derive(Serialize)]
pub struct RecentQuery {
    pub name: String,
    pub kind: String,
    pub verdict: Verdict,
    pub ago_secs: u64,
}

/// Кольцо: доходит до предела и начинает затирать самое старое.
///
/// Живёт только в памяти и умирает вместе со службой. На диск не попадает
/// намеренно — это перечень всего, что человек открывал, и хранить такое
/// дольше, чем нужно для ответа на вопрос «что сейчас произошло», незачем.
#[derive(Default)]
struct Recent {
    items: Vec<Seen>,
    /// Куда писать следующую запись; она же — самая старая, когда кольцо полно.
    next: usize,
}

impl Recent {
    fn record(&mut self, name: &str, qtype: u16, verdict: Verdict, at: Instant) {
        if self.items.len() < RECENT_CAP {
            self.items.push(Seen {
                name: [0; MAX_NAME],
                len: 0,
                qtype: 0,
                verdict: Verdict::Upstream,
                at,
            });
        }
        // Пишем на месте, а не собираем запись целиком и не двигаем её: имя
        // копируется ровно на свою длину, а не на все 255 байт массива.
        let slot = &mut self.items[self.next];
        let len = name.len().min(MAX_NAME);
        slot.name[..len].copy_from_slice(&name.as_bytes()[..len]);
        slot.len = len as u8;
        slot.qtype = qtype;
        slot.verdict = verdict;
        slot.at = at;

        self.next = (self.next + 1) % RECENT_CAP;
    }

    /// Отдаёт ленту от свежего к старому — в том порядке, в каком её читают.
    fn snapshot(&self, now: Instant) -> Vec<RecentQuery> {
        let n = self.items.len();
        (0..n)
            .map(|k| {
                let slot = &self.items[(self.next + n - 1 - k) % n];
                RecentQuery {
                    name: String::from_utf8_lossy(&slot.name[..slot.len as usize]).into_owned(),
                    kind: type_name(slot.qtype),
                    verdict: slot.verdict,
                    ago_secs: now.saturating_duration_since(slot.at).as_secs(),
                }
            })
            .collect()
    }
}

/// Название типа записи для панели. Считается только на запрос из панели, на
/// горячем пути хранится сырое число.
fn type_name(qtype: u16) -> String {
    match qtype {
        TYPE_A => "A".to_string(),
        TYPE_AAAA => "AAAA".to_string(),
        5 => "CNAME".to_string(),
        12 => "PTR".to_string(),
        15 => "MX".to_string(),
        16 => "TXT".to_string(),
        33 => "SRV".to_string(),
        65 => "HTTPS".to_string(),
        other => format!("#{other}"),
    }
}

/// Счётчики для панели. Обновляются на горячем пути, поэтому атомарные:
/// блокировка ради статистики — худшая из возможных причин ждать.
#[derive(Default)]
struct Counters {
    queries: AtomicU64,
    blocked: AtomicU64,
    cache_hits: AtomicU64,
    timeouts: AtomicU64,
    /// Подряд не отвеченных запросов — по ним переключаемся на запасной сервер.
    misses_in_row: AtomicU64,
}

#[derive(Serialize)]
pub struct DnsStatus {
    pub listening: bool,
    pub upstream: String,
    pub queries: u64,
    pub blocked: u64,
    pub cache_hits: u64,
    pub timeouts: u64,
    pub error: Option<String>,
    pub attempts: Vec<BlockedDomain>,
    pub recent: Vec<RecentQuery>,
}

/* ── общее состояние ───────────────────────────────────────── */

/// То, что видит остальная программа: список запрещённого и показания счётчиков.
pub struct Dns {
    blocklist: watch::Sender<Arc<Blocklist>>,
    counters: Arc<Counters>,
    attempts: Arc<Mutex<Attempts>>,
    recent: Arc<Mutex<Recent>>,
    /// Строка с адресом вышестоящего сервера либо причина, почему резолвер не
    /// работает. Панель должна показывать это раньше, чем родитель заметит, что
    /// блокировка молчит.
    state: Mutex<(bool, String, Option<String>)>,
}

impl Dns {
    pub fn new() -> Self {
        let (blocklist, _) = watch::channel(Arc::new(Blocklist::default()));
        Self {
            blocklist,
            counters: Arc::default(),
            attempts: Arc::default(),
            recent: Arc::default(),
            state: Mutex::new((false, String::new(), None)),
        }
    }

    /// Подменяет список запрещённого. Работающие задачи подхватят его на
    /// следующем же запросе — перезапускать резолвер не нужно.
    ///
    /// Именно `send_replace`, а не `send`: последний при отсутствии получателей
    /// не только возвращает ошибку, но и **оставляет прежнее значение**. А
    /// получателей нет ровно до того, как резолвер подпишется, — то есть при
    /// загрузке правил и при первом включении. С `send` список молча оставался
    /// бы пустым: адаптеры уведены на резолвер, панель показывает блокировку,
    /// а не блокируется ничего.
    pub fn set_domains<I: IntoIterator<Item = String>>(&self, domains: I) {
        self.blocklist.send_replace(Arc::new(Blocklist::new(domains)));
    }

    /// Отмечает, что резолвер погашен. Без этого панель показывала бы «слушаю»
    /// после выключения блокировки — то есть врала бы о состоянии системы.
    pub fn mark_stopped(&self) {
        if let Ok(mut s) = self.state.lock() {
            *s = (false, String::new(), None);
        }
        // Лента уходит вместе с резолвером: выключенная блокировка не должна
        // оставлять после себя перечень того, что человек открывал.
        if let Ok(mut recent) = self.recent.lock() {
            *recent = Recent::default();
        }
    }

    pub fn status(&self) -> DnsStatus {
        let now = Instant::now();
        let mut attempts: Vec<BlockedDomain> = self
            .attempts
            .lock()
            .map(|a| {
                a.domains
                    .iter()
                    .map(|(name, hit)| BlockedDomain {
                        name: name.to_string(),
                        count: hit.count,
                        ago_secs: now.saturating_duration_since(hit.last).as_secs(),
                    })
                    .collect()
            })
            .unwrap_or_default();
        attempts.sort_by_key(|d| d.ago_secs);

        let (listening, upstream, error) = self
            .state
            .lock()
            .map(|s| s.clone())
            .unwrap_or((false, String::new(), None));

        let recent = self
            .recent
            .lock()
            .map(|r| r.snapshot(now))
            .unwrap_or_default();

        DnsStatus {
            listening,
            upstream,
            queries: self.counters.queries.load(Ordering::Relaxed),
            blocked: self.counters.blocked.load(Ordering::Relaxed),
            cache_hits: self.counters.cache_hits.load(Ordering::Relaxed),
            timeouts: self.counters.timeouts.load(Ordering::Relaxed),
            error,
            attempts,
            recent,
        }
    }
}

/* ── резолвер ──────────────────────────────────────────────── */

struct Pending {
    client: SocketAddr,
    /// Сокет, на который пришёл запрос: ответить нужно с того же адреса, иначе
    /// клиент сочтёт ответ чужим. Клон `Arc` дешевле разбора адреса на ответе.
    sock: Arc<UdpSocket>,
    /// Идентификатор, с которым пришёл клиент: наружу уходит наш, обратно ему
    /// должен вернуться его собственный.
    client_id: u16,
    /// Ключ кэша, посчитанный на запросе, — на ответе имя разбирать не придётся.
    key: Option<Box<[u8]>>,
    deadline: Instant,
}

struct Resolver {
    upstream_sock: Arc<UdpSocket>,
    upstreams: Vec<SocketAddr>,
    active: AtomicUsize,
    inflight: Mutex<HashMap<u16, Pending>>,
    next_id: AtomicU16,
    cache: Mutex<Cache>,
    blocklist: watch::Receiver<Arc<Blocklist>>,
    counters: Arc<Counters>,
    attempts: Arc<Mutex<Attempts>>,
    recent: Arc<Mutex<Recent>>,
}

impl Resolver {
    fn upstream(&self) -> SocketAddr {
        self.upstreams[self.active.load(Ordering::Relaxed) % self.upstreams.len()]
    }

    /// Заносит запрос в ленту последних. Кольцо заведено заранее, поэтому
    /// вызов не выделяет памяти и держит мьютекс считаные наносекунды.
    fn note(&self, name: &str, qtype: u16, verdict: Verdict, now: Instant) {
        if let Ok(mut recent) = self.recent.lock() {
            recent.record(name, qtype, verdict, now);
        }
    }

    /// Решает судьбу запроса. Возвращает готовый ответ, если отвечаем сами.
    ///
    /// Здесь не должно появиться ни одного `await`: пока задача разбирает
    /// пакет, она не читает сокет, и очередь ядра копится.
    fn decide(&self, pkt: &mut [u8], len: usize, now: Instant) -> Decision {
        let mut name = [0u8; MAX_NAME];
        let Some(q) = parse_question(&pkt[..len], &mut name) else {
            // Разобрать не смогли — пусть отвечает вышестоящий сервер.
            return Decision::Forward(None);
        };
        // Имя собрано из меток, приведённых к нижнему регистру ASCII, так что
        // проверка кодировки здесь не нужна.
        let name = match std::str::from_utf8(&name[..q.len]) {
            Ok(name) => name,
            Err(_) => return Decision::Forward(None),
        };

        if self.blocklist.borrow().matches(name) {
            if let Ok(mut attempts) = self.attempts.lock() {
                attempts.record(name, now);
            }
            self.counters.blocked.fetch_add(1, Ordering::Relaxed);
            self.note(name, q.qtype, Verdict::Blocked, now);
            return match write_blocked(pkt, &q) {
                Some(n) => Decision::Answer(n),
                None => Decision::Drop,
            };
        }

        let mut key_buf = [0u8; MAX_NAME + 4];
        let key_len = cache_key(name.as_bytes(), q.qtype, q.qclass, &mut key_buf);
        let key = &key_buf[..key_len];

        // Вопрос спросившего сохраняем до того, как поверх ляжет ответ из кэша.
        // В кэше лежит вопрос того, кто спрашивал первым, а регистр имени у них
        // может отличаться: ключ приведён к нижнему, сам пакет — нет. Клиент
        // сверяет возвращённый вопрос со своим, и чужой регистр он вправе счесть
        // ответом не ему. Длины совпадают всегда — имена различаются только
        // регистром, — поэтому подмена ровно на месте.
        let mut asked = [0u8; MAX_NAME + 6];
        let asked_len = q.end - HEADER;
        if asked_len > asked.len() {
            return Decision::Forward(Some(key.into()));
        }
        asked[..asked_len].copy_from_slice(&pkt[HEADER..q.end]);

        if let Ok(cache) = self.cache.lock() {
            if let Some(n) = cache.get(key, packet_id(pkt), pkt, now) {
                pkt[HEADER..q.end].copy_from_slice(&asked[..asked_len]);
                self.counters.cache_hits.fetch_add(1, Ordering::Relaxed);
                drop(cache); // мьютекс кэша не держим на время записи в ленту
                self.note(name, q.qtype, Verdict::Cached, now);
                return Decision::Answer(n);
            }
        }
        self.note(name, q.qtype, Verdict::Upstream, now);
        Decision::Forward(Some(key.into()))
    }
}

enum Decision {
    /// Отвечаем сами; в буфере лежит готовый ответ указанной длины.
    Answer(usize),
    /// Спрашиваем наверх; внутри — ключ кэша, если его удалось посчитать.
    Forward(Option<Box<[u8]>>),
    Drop,
}

/// Принимает запросы от системы. Одна задача на сокет: разбор занимает доли
/// микросекунды, а параллельные приёмники только толкались бы на нём.
async fn client_loop(
    res: Arc<Resolver>,
    sock: Arc<UdpSocket>,
    mut shutdown: watch::Receiver<bool>,
) {
    // В куче, а не на стеке: массив такого размера внутри `async fn` раздул бы
    // само будущее задачи. Выделение одно на весь срок жизни цикла, на запросах
    // памяти не трогаем вовсе.
    let mut buf = vec![0u8; BUF];
    loop {
        let (len, client) = tokio::select! {
            r = sock.recv_from(&mut buf) => match r {
                Ok(v) => v,
                // Windows возвращает ошибку на сокет, если предыдущая отправка
                // упёрлась в закрытый порт. Это не повод ронять резолвер.
                Err(_) => continue,
            },
            _ = shutdown.changed() => return,
        };
        if len < HEADER {
            continue;
        }
        res.counters.queries.fetch_add(1, Ordering::Relaxed);

        let now = Instant::now();
        match res.decide(&mut buf, len, now) {
            Decision::Answer(n) => {
                let _ = sock.send_to(&buf[..n], client).await;
            }
            Decision::Forward(key) => {
                let client_id = packet_id(&buf);
                let id = res.next_id.fetch_add(1, Ordering::Relaxed);

                let accepted = match res.inflight.lock() {
                    Ok(mut table) => {
                        // Переполнение таблицы означает, что наверху затор.
                        // Молча отбрасываем: клиент переспросит сам, а расти
                        // без предела нельзя.
                        if table.len() >= MAX_INFLIGHT {
                            false
                        } else {
                            table.insert(
                                id,
                                Pending {
                                    client,
                                    sock: Arc::clone(&sock),
                                    client_id,
                                    key,
                                    deadline: now + UPSTREAM_TIMEOUT,
                                },
                            );
                            true
                        }
                    }
                    Err(_) => false,
                };

                if accepted {
                    set_packet_id(&mut buf, id);
                    let _ = res.upstream_sock.send_to(&buf[..len], res.upstream()).await;
                }
            }
            Decision::Drop => {}
        }
    }
}

/// Разбирает ответы вышестоящего сервера и раздаёт их тем, кто спрашивал.
async fn upstream_loop(res: Arc<Resolver>, mut shutdown: watch::Receiver<bool>) {
    let mut buf = vec![0u8; BUF];
    loop {
        let (len, from) = tokio::select! {
            r = res.upstream_sock.recv_from(&mut buf) => match r {
                Ok(v) => v,
                Err(_) => continue,
            },
            _ = shutdown.changed() => return,
        };
        // Сокет открыт наружу, поэтому ответ от постороннего адреса —
        // либо чужой пакет, либо попытка подсунуть свой адрес.
        if len < HEADER || !res.upstreams.contains(&from) {
            continue;
        }

        let id = packet_id(&buf);
        let Some(pending) = res.inflight.lock().ok().and_then(|mut t| t.remove(&id)) else {
            continue; // опоздал: запрос уже сочли потерянным
        };
        res.counters.misses_in_row.store(0, Ordering::Relaxed);

        if let Some(key) = &pending.key {
            if let Ok(mut cache) = res.cache.lock() {
                cache.put(key, &buf[..len], Instant::now());
            }
        }

        set_packet_id(&mut buf, pending.client_id);
        let _ = pending.sock.send_to(&buf[..len], pending.client).await;
    }
}

/// Чистит просроченное и переводит стрелку на запасной сервер, когда основной
/// замолчал. Без этого падение роутера означало бы, что имена не разрешаются
/// вовсе, — а виноватым выглядел бы родительский контроль.
async fn janitor(res: Arc<Resolver>, state: Shared, mut shutdown: watch::Receiver<bool>) {
    let mut last_report = Instant::now();
    loop {
        tokio::select! {
            _ = tokio::time::sleep(SWEEP_INTERVAL) => {}
            _ = shutdown.changed() => return,
        }
        let now = Instant::now();

        let expired = match res.inflight.lock() {
            Ok(mut table) => {
                let before = table.len();
                table.retain(|_, p| p.deadline > now);
                before - table.len()
            }
            Err(_) => 0,
        };

        if expired > 0 {
            res.counters
                .timeouts
                .fetch_add(expired as u64, Ordering::Relaxed);
            let row = res
                .counters
                .misses_in_row
                .fetch_add(expired as u64, Ordering::Relaxed)
                + expired as u64;
            if row >= FAILOVER_AFTER && res.upstreams.len() > 1 {
                res.counters.misses_in_row.store(0, Ordering::Relaxed);
                let next = res.active.fetch_add(1, Ordering::Relaxed) + 1;
                let addr = res.upstreams[next % res.upstreams.len()];
                if let Ok(mut s) = state.dns.state.lock() {
                    s.1 = addr.to_string();
                }
                state
                    .log(format!("DNS: сервер не отвечает, перехожу на {addr}"))
                    .await;
            }
        }

        if let Ok(mut cache) = res.cache.lock() {
            cache.sweep(now);
        }

        // Журнал попыток пишем пачкой. Discord переспрашивает своё имя по
        // нескольку раз в секунду, и построчная запись превратила бы и журнал в
        // мусор, и запись в файл — в то самое узкое место.
        if last_report.elapsed() >= REPORT_INTERVAL {
            last_report = now;
            let summary = res.attempts.lock().ok().map(|mut a| {
                let mut fresh: Vec<(String, u64)> = a
                    .domains
                    .iter_mut()
                    .filter(|(_, h)| h.count > h.reported)
                    .map(|(name, h)| {
                        let delta = h.count - h.reported;
                        h.reported = h.count;
                        (name.to_string(), delta)
                    })
                    .collect();
                fresh.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
                fresh
            });

            if let Some(fresh) = summary.filter(|f| !f.is_empty()) {
                let listed: Vec<String> = fresh
                    .iter()
                    .take(5)
                    .map(|(name, n)| format!("{name} ×{n}"))
                    .collect();
                let tail = match fresh.len().saturating_sub(5) {
                    0 => String::new(),
                    more => format!(" и ещё {more}"),
                };
                state
                    .log(format!("DNS: отклонены запросы — {}{tail}", listed.join(", ")))
                    .await;
            }
        }
    }
}

/* ── DNS поверх TCP ────────────────────────────────────────── */

/// По TCP клиент приходит, когда ответ не влез в пакет. Таких запросов единицы,
/// поэтому здесь всё просто: задача на соединение и пересылка как есть. Не
/// поддержать TCP нельзя — иначе часть имён переставала бы разрешаться, и
/// причину искали бы неделю.
async fn tcp_loop(res: Arc<Resolver>, listener: TcpListener, mut shutdown: watch::Receiver<bool>) {
    let live = Arc::new(AtomicUsize::new(0));
    loop {
        let (mut stream, _) = tokio::select! {
            r = listener.accept() => match r {
                Ok(v) => v,
                Err(_) => continue,
            },
            _ = shutdown.changed() => return,
        };

        if live.load(Ordering::Relaxed) >= MAX_TCP {
            continue; // соединение закроется само вместе с `stream`
        }
        live.fetch_add(1, Ordering::Relaxed);

        let res = Arc::clone(&res);
        let live = Arc::clone(&live);
        tokio::spawn(async move {
            let _ = tokio::time::timeout(TCP_IDLE, serve_tcp(&res, &mut stream)).await;
            live.fetch_sub(1, Ordering::Relaxed);
        });
    }
}

async fn serve_tcp(res: &Resolver, stream: &mut TcpStream) {
    // На соединение, а не на задачу: по TCP приходят единицы клиентов, и
    // держать под каждого постоянный буфер незачем.
    let mut buf = vec![0u8; BUF];
    loop {
        // По TCP сообщению предшествует его длина в двух байтах.
        let mut len_bytes = [0u8; 2];
        if stream.read_exact(&mut len_bytes).await.is_err() {
            return;
        }
        let len = u16::from_be_bytes(len_bytes) as usize;
        if !(HEADER..=BUF).contains(&len) {
            return;
        }
        if stream.read_exact(&mut buf[..len]).await.is_err() {
            return;
        }
        res.counters.queries.fetch_add(1, Ordering::Relaxed);

        let now = Instant::now();
        let reply = match res.decide(&mut buf, len, now) {
            Decision::Answer(n) => Some(n),
            Decision::Drop => return,
            // Кэш по TCP не наполняем: ответы здесь длинные, а приходят такие
            // запросы единицами — экономить на них нечего.
            Decision::Forward(_) => forward_tcp(res, len, &mut buf).await,
        };

        let Some(n) = reply else { return };
        if stream.write_all(&(n as u16).to_be_bytes()).await.is_err()
            || stream.write_all(&buf[..n]).await.is_err()
        {
            return;
        }
    }
}

/// Спрашивает вышестоящий сервер по TCP. Запрос лежит в начале `buf`, туда же
/// кладётся ответ; возвращается его длина.
async fn forward_tcp(res: &Resolver, query_len: usize, buf: &mut [u8]) -> Option<usize> {
    let mut up = tokio::time::timeout(UPSTREAM_TIMEOUT, TcpStream::connect(res.upstream()))
        .await
        .ok()?
        .ok()?;

    let header = (query_len as u16).to_be_bytes();
    up.write_all(&header).await.ok()?;
    up.write_all(&buf[..query_len]).await.ok()?;

    let mut len_bytes = [0u8; 2];
    tokio::time::timeout(UPSTREAM_TIMEOUT, up.read_exact(&mut len_bytes))
        .await
        .ok()?
        .ok()?;
    let len = u16::from_be_bytes(len_bytes) as usize;
    if !(HEADER..=BUF).contains(&len) {
        return None;
    }
    tokio::time::timeout(UPSTREAM_TIMEOUT, up.read_exact(&mut buf[..len]))
        .await
        .ok()?
        .ok()?;
    Some(len)
}

/* ── запуск ────────────────────────────────────────────────── */

/// Сообщает, что резолвер не поднялся. Панель должна показать причину: молчащая
/// блокировка выглядит одинаково с выключенной.
async fn fail(state: &Shared, msg: String) {
    if let Ok(mut s) = state.dns.state.lock() {
        *s = (false, String::new(), Some(msg.clone()));
    }
    state.log(format!("DNS: {msg}")).await;
}

/// Поднимает резолвер и возвращает признак «порт занят, запросы идут».
///
/// Возвращаемое значение — не формальность: адаптеры переключаются на нас
/// только после успешного запуска, иначе неудачный старт оставил бы машину
/// вообще без разрешения имён.
pub async fn run(state: Shared, upstreams: Vec<SocketAddr>, shutdown: watch::Receiver<bool>) -> bool {
    if upstreams.is_empty() {
        fail(&state, "не нашёл ни одного вышестоящего сервера".to_string()).await;
        return false;
    }

    let client_v4 = match UdpSocket::bind(BIND_ADDR).await {
        Ok(s) => Arc::new(s),
        // Порт 53 занимают «Общий доступ к подключению» и Docker. Без него
        // блокировка сайтов не работает, но правила по процессам — работают.
        Err(e) => {
            fail(
                &state,
                format!("{BIND_ADDR} занят ({e}), блокировка сайтов выключена"),
            )
            .await;
            return false;
        }
    };
    // Шестая версия необязательна: на машине без IPv6 сокет не поднимется, и
    // это нормально — тогда и запросов по нему не будет.
    let client_v6 = UdpSocket::bind(BIND_ADDR_V6).await.ok().map(Arc::new);
    let tcp_v4 = TcpListener::bind(BIND_ADDR).await.ok();
    let tcp_v6 = TcpListener::bind(BIND_ADDR_V6).await.ok();

    let upstream_sock = match UdpSocket::bind("0.0.0.0:0").await {
        Ok(s) => Arc::new(s),
        Err(e) => {
            fail(&state, format!("нет исходящего сокета: {e}")).await;
            return false;
        }
    };

    let resolver = Arc::new(Resolver {
        upstream_sock,
        active: AtomicUsize::new(0),
        inflight: Mutex::new(HashMap::with_capacity(256)),
        // Начальный идентификатор случаен: предсказуемая нумерация облегчила бы
        // подделку ответа тому, кто видит наш исходящий порт.
        next_id: AtomicU16::new(rand::random()),
        cache: Mutex::new(Cache::default()),
        blocklist: state.dns.blocklist.subscribe(),
        counters: Arc::clone(&state.dns.counters),
        attempts: Arc::clone(&state.dns.attempts),
        recent: Arc::clone(&state.dns.recent),
        upstreams,
    });

    if let Ok(mut s) = state.dns.state.lock() {
        *s = (true, resolver.upstream().to_string(), None);
    }
    state
        .log(format!(
            "DNS: слушаю {BIND_ADDR}, спрашиваю {}",
            resolver.upstream()
        ))
        .await;

    for listener in [tcp_v4, tcp_v6].into_iter().flatten() {
        tokio::spawn(tcp_loop(Arc::clone(&resolver), listener, shutdown.clone()));
    }
    if let Some(sock) = client_v6 {
        tokio::spawn(client_loop(Arc::clone(&resolver), sock, shutdown.clone()));
    }
    tokio::spawn(upstream_loop(Arc::clone(&resolver), shutdown.clone()));
    tokio::spawn(janitor(
        Arc::clone(&resolver),
        Arc::clone(&state),
        shutdown.clone(),
    ));
    // Приём запросов уходит в фоновую задачу: вызвавшему нужно знать, что порт
    // занят, и идти переключать адаптеры, а не ждать здесь до остановки службы.
    tokio::spawn(client_loop(resolver, client_v4, shutdown));
    true
}

#[cfg(test)]
#[path = "dns_tests.rs"]
mod tests;