//! Tor as a transport (`--transports tor`, in builds with the `tor` feature), for
//! reachability: every node runs an onion service (arti), reachable from anywhere
//! without a forwarded port or a relay. Not for anonymity: a node's other addresses are
//! in the member list next to its onion address.
//!
//! A node's transport identity is its node key (64 hex digits); where it is, its onion
//! service: `<onion>.onion:<port>/<key>` (`locate`, `locations`). The onion service's
//! own key lives in arti's state (`--tor-state-dir`), so its address stays the same
//! across restarts.
//!
//! Peers talk over one long-lived connection per pair (opening a stream over Tor costs
//! a round trip through the rendezvous circuit): a Tor stream to the other's onion
//! service, either way. Tor authenticates the onion service to the one who dials it,
//! but nothing about the dialer to the service: the challenge handshake
//! (`transport::challenge`) at the start of every connection proves both keys, bound to
//! the onion address dialed. Then messages go as on libp2p: their length (4 bytes, big
//! endian), then the message as `wire` encodes it. A connection that breaks is dialed
//! again on the next send. Both ends may dial at once; each then sends on the newest
//! connection and reads from both until the older one closes.
//!
//! Onion circuits are slow (round trips of a second or more, seconds to set one up),
//! so the timeouts are generous and messages smaller than on libp2p; the paths layer
//! (`transport::paths`) prefers faster transports while they work.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use anyhow::{Context, Result, anyhow, bail};
use arti_client::config::TorClientConfigBuilder;
use arti_client::{DataStream, TorClient};
use async_trait::async_trait;
use futures::StreamExt;
use safelog::DisplayRedacted as _;
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::sync::{mpsc, watch};
use tor_cell::relaycell::msg::{Connected, End, EndReason};
use tor_hsservice::config::OnionServiceConfigBuilder;
use tor_hsservice::{RunningOnionService, StreamRequest, handle_rend_requests};
use tor_proto::stream::IncomingStreamRequest;
use tor_rtcompat::PreferredRuntime;
use tracing::{debug, info, warn};

use crate::identity::NodeKey;
use crate::locks::MutexExt;
use crate::transport::challenge::{self, Role};
use crate::transport::{EncodedMessage, Hex, MessageHandler, ParsedMessage, Transport, wire};

/// The onion service's nickname in arti's state (one per node).
const NICKNAME: &str = "continuum";
/// The largest message read. Sync messages above 16 chunks travel as blobs, and a chunk
/// is at most 4 MiB.
const MAX_MESSAGE: usize = 8 << 20;
/// How long reaching an onion service may take (building circuits, the rendezvous).
const CONNECT_TIMEOUT: Duration = Duration::from_secs(60);
/// How long the challenge handshake may take, a few round trips.
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
/// How long a message may take, before what its size adds (`stream_timeout`).
const STREAM_TIMEOUT: Duration = Duration::from_secs(30);
/// The slowest a message is expected to go over Tor, in bytes per second.
const MIN_RATE: u64 = 64 << 10;
/// Messages read and not handled yet, at most.
const READS: usize = 32;
/// Peers with a known location, at most.
const MAX_PEERS: usize = 4096;

/// How long a message of `bytes` may take.
fn stream_timeout(bytes: usize) -> Duration {
    STREAM_TIMEOUT + Duration::from_millis(bytes as u64 * 1000 / MIN_RATE)
}

/// What a Tor transport is made with.
pub(crate) struct TorConfig {
    /// The node key: it proves who we are in the challenge handshake.
    pub(crate) key: Arc<NodeKey>,
    /// Arti's state (the onion service's key, among others) and cache, in `state` and
    /// `cache` below.
    pub(crate) state_dir: PathBuf,
    /// The onion service's port.
    pub(crate) port: u16,
}

/// Whether `id` is a Tor transport identity: a public key, 64 lowercase hex digits.
fn is_key(id: &str) -> bool {
    id.len() == 64
        && id
            .bytes()
            .all(|byte| matches!(byte, b'0'..=b'9' | b'a'..=b'f'))
}

/// A location, `<onion>.onion:<port>`, and the key there, if `target` names one.
fn parse_location(target: &str) -> Option<(String, Option<String>)> {
    let (location, key) = match target.split_once('/') {
        Some((location, key)) => (location, Some(key)),
        None => (target, None),
    };
    let (host, port) = location.rsplit_once(':')?;
    let name = host.strip_suffix(".onion")?;
    if name.is_empty() || !name.bytes().all(|byte| byte.is_ascii_alphanumeric()) {
        return None;
    }
    port.parse::<u16>().ok()?;
    if key.is_some_and(|key| !is_key(key)) {
        return None;
    }
    Some((location.to_ascii_lowercase(), key.map(str::to_owned)))
}

/// One connection to a peer: where to write. A reader task reads the other half.
struct Connection {
    id: u64,
    writer: tokio::sync::Mutex<WriteHalf<DataStream>>,
}

struct Shared {
    key: Arc<NodeKey>,
    local_id: String,
    /// Where each peer's onion service is, by key.
    book: Mutex<HashMap<String, String>>,
    /// The connection to write to each peer on, by key.
    connections: Mutex<HashMap<String, Arc<Connection>>>,
    /// One dial at a time per location.
    dialing: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
    next_connection: AtomicU64,
    /// Messages read, with the sender's key, for the consumer.
    inbound: mpsc::Sender<(String, Vec<u8>)>,
}

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex.lock_or_recover()
}

pub(crate) struct TorTransport {
    shared: Arc<Shared>,
    client: Arc<TorClient<PreferredRuntime>>,
    /// Our onion service, `<onion>.onion:<port>`.
    onion: String,
    /// Kept: dropping it stops the onion service.
    _service: Arc<RunningOnionService>,
    incoming: tokio::sync::Mutex<Option<mpsc::Receiver<(String, Vec<u8>)>>>,
    stopped: watch::Receiver<bool>,
}

impl TorTransport {
    /// Bootstraps arti (this takes a while: it fetches what it needs to know of the Tor
    /// network) and launches our onion service.
    pub(crate) async fn start(config: TorConfig) -> Result<Self> {
        let tor_config = TorClientConfigBuilder::from_directories(
            config.state_dir.join("state"),
            config.state_dir.join("cache"),
        )
        .build()
        .context("Configuring Tor")?;
        info!(state_dir = %config.state_dir.display(), "Bootstrapping Tor");
        let client = TorClient::create_bootstrapped(tor_config)
            .await
            .context("Bootstrapping Tor")?;
        let service_config = OnionServiceConfigBuilder::default()
            .nickname(
                NICKNAME
                    .parse()
                    .map_err(|error| anyhow!("Onion service nickname: {error}"))?,
            )
            .build()
            .context("Configuring the onion service")?;
        let (service, requests) = client
            .launch_onion_service(service_config)
            .context("Launching the onion service")?
            .ok_or_else(|| anyhow!("The onion service is disabled in arti's configuration"))?;
        let name = service
            .onion_address()
            .ok_or_else(|| anyhow!("The onion service has no address"))?;
        // Arti shows onion addresses redacted unless asked: we publish ours.
        let onion = format!("{}:{}", name.display_unredacted(), config.port);

        let local_id = Hex(&config.key.public()).to_string();
        let (inbound, incoming) = mpsc::channel(READS);
        let shared = Arc::new(Shared {
            key: config.key,
            local_id,
            book: Mutex::default(),
            connections: Mutex::default(),
            dialing: Mutex::default(),
            next_connection: AtomicU64::new(0),
            inbound,
        });
        let (stopping, stopped) = watch::channel(false);
        {
            let shared = Arc::clone(&shared);
            let onion = onion.clone();
            let port = config.port;
            tokio::spawn(async move {
                let requests = handle_rend_requests(requests);
                futures::pin_mut!(requests);
                while let Some(request) = requests.next().await {
                    let shared = Arc::clone(&shared);
                    let onion = onion.clone();
                    tokio::spawn(async move { shared.accept(request, &onion, port).await });
                }
                warn!("The onion service stopped");
                let _ = stopping.send(true);
            });
        }
        info!(address = %onion, "Tor transport started (the onion service may take a minute to be reachable)");
        Ok(Self {
            shared,
            client,
            onion,
            _service: service,
            incoming: tokio::sync::Mutex::new(Some(incoming)),
            stopped,
        })
    }

    /// A connection to the peer with `key`, made if there is none.
    async fn connection(&self, key: &str) -> Result<Arc<Connection>> {
        if let Some(connection) = lock(&self.shared.connections).get(key) {
            return Ok(Arc::clone(connection));
        }
        let Some(location) = lock(&self.shared.book).get(key).cloned() else {
            bail!("no onion address known for {key}");
        };
        self.dial(&location, Some(key)).await?;
        lock(&self.shared.connections)
            .get(key)
            .cloned()
            .ok_or_else(|| anyhow!("the connection to {key} closed at once"))
    }

    /// Dials the onion service at `location` and runs the handshake: returns the key
    /// there, which must be `expected` if given.
    async fn dial(&self, location: &str, expected: Option<&str>) -> Result<String> {
        let dialing = Arc::clone(
            lock(&self.shared.dialing)
                .entry(location.to_owned())
                .or_default(),
        );
        let _one_at_a_time = dialing.lock().await;
        // Someone else may have connected meanwhile.
        if let Some(expected) = expected
            && lock(&self.shared.connections).contains_key(expected)
        {
            return Ok(expected.to_owned());
        }
        let result = async {
            let (host, port) = location
                .rsplit_once(':')
                .and_then(|(host, port)| Some((host, port.parse::<u16>().ok()?)))
                .ok_or_else(|| anyhow!("{location} is no onion address with a port"))?;
            let mut stream =
                tokio::time::timeout(CONNECT_TIMEOUT, self.client.connect((host, port)))
                    .await
                    .map_err(|_| {
                        anyhow!("reaching {location} took longer than {CONNECT_TIMEOUT:?}")
                    })?
                    .with_context(|| format!("Reaching {location}"))?;
            let key = challenge::run(
                &mut stream,
                &self.shared.key,
                Role::Dialer,
                location,
                HANDSHAKE_TIMEOUT,
            )
            .await
            .with_context(|| format!("Handshake with {location}"))?;
            let key = Hex(&key).to_string();
            if let Some(expected) = expected
                && key != expected
            {
                bail!("{location} is another node now ({key}, not {expected})");
            }
            self.shared.remember(&key, location);
            self.shared.register(key.clone(), stream);
            Ok(key)
        }
        .await;
        lock(&self.shared.dialing).remove(location);
        result
    }
}

impl Shared {
    /// `key` is at `location`.
    fn remember(&self, key: &str, location: &str) {
        if key == self.local_id {
            return;
        }
        let mut book = lock(&self.book);
        if book.len() < MAX_PEERS || book.contains_key(key) {
            book.insert(key.to_owned(), location.to_owned());
        }
    }

    /// A stream to our onion service: accepted if for our port, then the handshake.
    async fn accept(self: Arc<Self>, request: StreamRequest, onion: &str, port: u16) {
        let ours = matches!(request.request(), IncomingStreamRequest::Begin(begin) if begin.port() == port);
        if !ours {
            let _ = request
                .reject(End::new_with_reason(EndReason::CONNECTREFUSED))
                .await;
            return;
        }
        let mut stream = match request.accept(Connected::new_empty()).await {
            Ok(stream) => stream,
            Err(error) => {
                debug!(error = %error, "Couldn't accept a stream to our onion service");
                return;
            }
        };
        match challenge::run(
            &mut stream,
            &self.key,
            Role::Listener,
            onion,
            HANDSHAKE_TIMEOUT,
        )
        .await
        {
            Ok(key) => self.register(Hex(&key).to_string(), stream),
            Err(error) => {
                debug!(error = %format!("{error:#}"), "Dropping a connection to our onion service: the handshake failed");
            }
        }
    }

    /// A connection to `key`, authenticated: the one to write to it on from now, and
    /// read from until it closes.
    fn register(self: &Arc<Self>, key: String, stream: DataStream) {
        let (reader, writer) = tokio::io::split(stream);
        let id = self.next_connection.fetch_add(1, Ordering::Relaxed);
        let connection = Arc::new(Connection {
            id,
            writer: tokio::sync::Mutex::new(writer),
        });
        lock(&self.connections).insert(key.clone(), connection);
        debug!(peer = %key, connection = id, "Tor connection up");
        let shared = Arc::clone(self);
        tokio::spawn(async move {
            let error = shared.read(&key, reader).await;
            debug!(peer = %key, connection = id, error = %format!("{error:#}"), "Tor connection closed");
            shared.forget(&key, id);
        });
    }

    /// Reads messages from `key` until the connection fails; returns why.
    async fn read(&self, key: &str, mut reader: ReadHalf<DataStream>) -> anyhow::Error {
        loop {
            let mut length = [0u8; 4];
            if let Err(error) = reader.read_exact(&mut length).await {
                return error.into();
            }
            let length = u32::from_be_bytes(length) as usize;
            if length > MAX_MESSAGE {
                return anyhow!("a message of {length} bytes; at most {MAX_MESSAGE} are read");
            }
            let mut body = vec![0u8; length];
            match tokio::time::timeout(stream_timeout(length), reader.read_exact(&mut body)).await {
                Ok(Ok(_)) => {}
                Ok(Err(error)) => return error.into(),
                Err(_) => {
                    return anyhow!(
                        "{length} bytes took longer than {:?}",
                        stream_timeout(length)
                    );
                }
            }
            if self.inbound.send((key.to_owned(), body)).await.is_err() {
                return anyhow!("the transport stopped");
            }
        }
    }

    /// The connection `id` to `key` is gone: not to be written to any more.
    fn forget(&self, key: &str, id: u64) {
        let mut connections = lock(&self.connections);
        if connections
            .get(key)
            .is_some_and(|connection| connection.id == id)
        {
            connections.remove(key);
        }
    }
}

#[async_trait]
impl Transport for TorTransport {
    fn local_id(&self) -> &str {
        &self.shared.local_id
    }

    fn is_peer_id(&self, id: &str) -> bool {
        is_key(id)
    }

    /// The longest a message may take.
    fn delivery_limit(&self) -> Duration {
        stream_timeout(MAX_MESSAGE)
    }

    /// Our onion service, with our key.
    fn local_addresses(&self) -> Vec<String> {
        vec![format!("{}/{}", self.onion, self.shared.local_id)]
    }

    /// A key, or a location with one (noted as where that key is).
    fn locate(&self, target: &str) -> Option<String> {
        if is_key(target) {
            return Some(target.to_owned());
        }
        let (location, key) = parse_location(target)?;
        let key = key?;
        self.shared.remember(&key, &location);
        Some(key)
    }

    fn locations(&self, target: &str) -> Vec<String> {
        lock(&self.shared.book)
            .get(target)
            .map(|location| format!("{location}/{target}"))
            .into_iter()
            .collect()
    }

    /// A location with a key is that key; one without is dialed, to find out.
    async fn resolve_seed(&self, seed: &str) -> Result<String> {
        if is_key(seed) {
            if lock(&self.shared.book).contains_key(seed) {
                return Ok(seed.to_owned());
            }
            bail!("no onion address known for {seed}; give its onion address");
        }
        let (location, key) = parse_location(seed)
            .ok_or_else(|| anyhow!("{seed} is no onion address with a port"))?;
        if location == self.onion {
            bail!("{seed} is this node");
        }
        match key {
            Some(key) => {
                self.shared.remember(&key, &location);
                Ok(key)
            }
            None => self.dial(&location, None).await,
        }
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        let body = wire::encode(message)?;
        Ok(EncodedMessage {
            peer: peer.to_owned(),
            address: peer.to_owned(),
            payload_len: body.len(),
            body: body.into(),
            kind: message.kind(),
            message: None,
        })
    }

    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        let key = &message.peer;
        if !is_key(key) {
            bail!("{key} is no Tor transport identity");
        }
        if message.body.len() > MAX_MESSAGE {
            bail!(
                "a message of {} bytes; at most {MAX_MESSAGE} go over Tor",
                message.body.len()
            );
        }
        let connection = self.connection(key).await?;
        let timeout = stream_timeout(message.body.len());
        let written = tokio::time::timeout(timeout, async {
            let mut writer = connection.writer.lock().await;
            writer
                .write_all(&(message.body.len() as u32).to_be_bytes())
                .await?;
            writer.write_all(&message.body).await?;
            writer.flush().await
        })
        .await;
        let result = match written {
            Ok(Ok(())) => Ok(()),
            Ok(Err(error)) => Err(anyhow::Error::from(error)),
            Err(_) => Err(anyhow!("sending to {key} took longer than {timeout:?}")),
        };
        if result.is_err() {
            // Dialed again on the next send.
            self.shared.forget(key, connection.id);
        }
        result.with_context(|| format!("Sending to {key} over Tor"))
    }

    /// Handles the messages read from every connection, one at a time, in the order
    /// they were read.
    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        let Some(mut incoming) = self.incoming.lock().await.take() else {
            bail!("Tor consumer already started");
        };
        while let Some((peer, body)) = incoming.recv().await {
            match wire::decode(&body) {
                Ok(message) => {
                    if let Err(error) = handler(peer.clone(), message, body.len()).await {
                        warn!(peer = %peer, error = %format!("{error:#}"), "Failed to process inbound message");
                    }
                }
                Err(error) => {
                    debug!(peer = %peer, error = %error, "Dropping a Tor message that doesn't decode");
                }
            }
        }
        Ok(())
    }

    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        let mut stopped = self.stopped.clone();
        let _ = stopped.wait_for(|stopped| *stopped).await;
        Err(std::io::Error::other("the onion service stopped"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const ONION: &str = "vww6ybal4bd7szmgncyruucpgfkqahzddi37ktceo3ah7ngmcopnpyyd.onion";

    #[test]
    fn locations_and_keys() {
        let key = "ab".repeat(32);
        assert!(is_key(&key));
        assert!(!is_key(&key.to_uppercase()));
        assert!(!is_key(&key[1..]));
        assert_eq!(
            parse_location(&format!("{ONION}:9662/{key}")),
            Some((format!("{ONION}:9662"), Some(key.clone())))
        );
        assert_eq!(
            parse_location(&format!("{ONION}:9662")),
            Some((format!("{ONION}:9662"), None))
        );
        // A port is needed, a key must be one, and only onion addresses.
        assert_eq!(parse_location(ONION), None);
        assert_eq!(parse_location(&format!("{ONION}:9662/xyz")), None);
        assert_eq!(parse_location("example.org:9662"), None);
        assert_eq!(parse_location(&format!("{ONION}:99999")), None);
    }
}