cp_tor/

runtime.rs

//! Tor runtime: manages Arti bootstrap, circuit pool, and onion service.
//!
//! Provides the `TorRuntime` struct that wraps arti-client for the Canon
//! live search protocol (CP-013 §13-17). Manages a pool of warm Tor
//! circuits to peers, an onion service for incoming queries, and the
//! full lifecycle from bootstrap to teardown.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use arti_client::{DataStream, StreamPrefs, TorClient, TorClientConfig};
use safelog::DisplayRedacted;
use tokio::sync::Mutex;
use tor_cell::relaycell::msg::Connected;
use tor_hsservice::config::OnionServiceConfigBuilder;
use tor_hsservice::HsNickname;
use tor_rtcompat::PreferredRuntime;
use tracing::{debug, info, warn};

use cp_graph::GraphStore;

use crate::error::{Result, TorError};
use crate::rate_limit::RateLimiter;
use crate::server::ServerConfig;
use crate::types::{
    PeerCapabilities, PeerRegistration, PeerScore, CANON_PORT, CIRCUIT_POOL_SIZE,
    CIRCUIT_ROTATION_SECS, KEEPALIVE_INTERVAL_SECS,
};
use crate::wire;

fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64
}

/// Configuration for the Tor runtime.
pub struct TorConfig {
    /// Whether to launch a Tor onion service for incoming queries.
    pub enable_onion_service: bool,
    /// Number of warm circuits to maintain in the pool.
    pub circuit_pool_size: usize,
    /// Keepalive interval in seconds.
    pub keepalive_interval_secs: u64,
    /// Circuit rotation interval in seconds.
    pub circuit_rotation_secs: u64,
    /// Node identity secret key (Ed25519 seed).
    pub identity_secret: [u8; 32],
    /// Node's Ed25519 public key.
    pub identity_public: [u8; 32],
    /// BLAKE3 hash of the embedding model.
    pub model_hash: [u8; 32],
    /// Topics this node covers.
    pub topics: Vec<String>,
}

impl Default for TorConfig {
    fn default() -> Self {
        Self {
            enable_onion_service: true,
            circuit_pool_size: CIRCUIT_POOL_SIZE,
            keepalive_interval_secs: KEEPALIVE_INTERVAL_SECS,
            circuit_rotation_secs: CIRCUIT_ROTATION_SECS,
            identity_secret: [0u8; 32],
            identity_public: [0u8; 32],
            model_hash: [0u8; 32],
            topics: Vec::new(),
        }
    }
}

/// State of the Tor runtime.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuntimeState {
    /// Not yet bootstrapped.
    Idle,
    /// Bootstrap in progress.
    Bootstrapping,
    /// Connected and ready for search.
    Ready,
    /// Shutting down.
    ShuttingDown,
}

/// A warm circuit entry holding an actual `DataStream` to a peer.
#[allow(dead_code)]
struct PooledCircuit {
    peer: PeerRegistration,
    stream: DataStream,
    created_at: i64,
    last_used: i64,
    last_keepalive: i64,
}

impl PooledCircuit {
    fn needs_keepalive(&self, now_ms: i64) -> bool {
        let elapsed = (now_ms - self.last_keepalive) / 1000;
        elapsed >= KEEPALIVE_INTERVAL_SECS as i64
    }

    fn needs_rotation(&self, now_ms: i64) -> bool {
        let elapsed = (now_ms - self.created_at) / 1000;
        elapsed >= CIRCUIT_ROTATION_SECS as i64
    }
}

/// Manages Tor connectivity for the Canon node.
///
/// Handles bootstrap, circuit pool management with warm circuits to
/// peers, and an onion service for incoming search queries. The circuit
/// pool targets `CIRCUIT_POOL_SIZE` warm connections with keepalive probes
/// every `KEEPALIVE_INTERVAL_SECS` and rotation every `CIRCUIT_ROTATION_SECS`.
pub struct TorRuntime {
    client: TorClient<PreferredRuntime>,
    config: TorConfig,
    state: Arc<Mutex<RuntimeState>>,
    /// The node's onion address (set after onion service launches).
    onion_address: Arc<Mutex<Option<String>>>,
    /// Handle to the running onion service — must be kept alive or the service shuts down.
    #[allow(dead_code)]
    onion_service: Mutex<Option<Arc<tor_hsservice::RunningOnionService>>>,
}

impl TorRuntime {
    /// Bootstrap the Tor client and connect to the network.
    ///
    /// This blocks until the Tor client is bootstrapped and connected,
    /// which typically takes under 1 minute on first launch. Returns
    /// when the client is ready for circuit establishment.
    pub async fn bootstrap(config: TorConfig) -> Result<Self> {
        info!("Starting Tor bootstrap...");
        let start = Instant::now();

        let tor_config = TorClientConfig::default();
        let client = TorClient::create_bootstrapped(tor_config)
            .await
            .map_err(|e| TorError::Connection(format!("Tor bootstrap failed: {e}")))?;

        let elapsed = start.elapsed();
        info!("Tor bootstrapped in {:.1}s", elapsed.as_secs_f64());

        Ok(Self {
            client,
            config,
            state: Arc::new(Mutex::new(RuntimeState::Ready)),
            onion_address: Arc::new(Mutex::new(None)),
            onion_service: Mutex::new(None),
        })
    }

    /// Get the current runtime state.
    pub async fn state(&self) -> RuntimeState {
        *self.state.lock().await
    }

    /// Get the node's onion address (if the onion service is running).
    pub async fn onion_address(&self) -> Option<String> {
        self.onion_address.lock().await.clone()
    }

    /// Open an anonymous Tor stream to a peer's onion address.
    ///
    /// Returns a `DataStream` that implements `AsyncRead` + `AsyncWrite`.
    pub async fn connect_to_peer(&self, onion_address: &str) -> Result<DataStream> {
        let addr = format!("{onion_address}.onion:{CANON_PORT}");
        let prefs = StreamPrefs::new();
        let stream = self
            .client
            .connect_with_prefs(&addr, &prefs)
            .await
            .map_err(|e| TorError::Connection(format!("Failed to connect to {addr}: {e}")))?;

        debug!(
            "Connected to peer {}",
            &onion_address[..8.min(onion_address.len())]
        );
        Ok(stream)
    }

    /// Build a `PeerRegistration` for this node to publish on Arweave.
    ///
    /// The `onion_address` is populated from the running onion service.
    /// Call `start_onion_service` before this to ensure the address is set.
    pub async fn build_registration(&self) -> PeerRegistration {
        let node_id = crate::keys::node_id_from_public_key(&self.config.identity_public);

        let onion_address = self.onion_address.lock().await.clone().unwrap_or_default();

        let mut reg = PeerRegistration {
            onion_address,
            node_id,
            public_key: self.config.identity_public,
            capabilities: PeerCapabilities::default(),
            topics: self.config.topics.clone(),
            embedding_model: self.config.model_hash,
            timestamp: now_ms(),
            signature: [0u8; 64],
        };

        let signing_bytes = reg.signing_bytes();
        let signing_key = ed25519_dalek::SigningKey::from_bytes(&self.config.identity_secret);
        reg.signature = ed25519_dalek::Signer::sign(&signing_key, &signing_bytes).to_bytes();

        reg
    }

    /// Launch the Tor onion service for incoming search queries.
    ///
    /// Uses the arti `TorClient::launch_onion_service` API to create a real
    /// Tor hidden service. The service address is derived by arti from an
    /// auto-generated or stored keypair. Returns the stream of rendezvous
    /// requests that should be fed to `run_accept_loop`.
    pub async fn start_onion_service(
        &self,
    ) -> Result<futures::stream::BoxStream<'static, tor_hsservice::RendRequest>> {
        use futures::StreamExt;

        if !self.config.enable_onion_service {
            info!("Onion service disabled by config");
            return Err(TorError::Connection("Onion service disabled".into()));
        }

        info!("Launching onion service on port {}", CANON_PORT);

        let nickname: HsNickname = "canon-search"
            .parse()
            .map_err(|e| TorError::Connection(format!("Invalid service nickname: {e}")))?;

        let svc_config = OnionServiceConfigBuilder::default()
            .nickname(nickname)
            .build()
            .map_err(|e| TorError::Connection(format!("Onion service config error: {e}")))?;

        let (onion_service, rend_requests) = self
            .client
            .launch_onion_service(svc_config)
            .map_err(|e| TorError::Connection(format!("Failed to launch onion service: {e}")))?
            .ok_or_else(|| TorError::Connection("Onion service disabled by Tor config".into()))?;

        // Store the handle so the onion service stays alive
        *self.onion_service.lock().await = Some(onion_service.clone());

        // Wait briefly for the service to publish its descriptor
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;

        if let Some(hs_id) = onion_service.onion_address() {
            let address = hs_id.display_unredacted().to_string();
            // Strip the ".onion" suffix for internal storage
            let address = address.trim_end_matches(".onion").to_string();
            *self.onion_address.lock().await = Some(address.clone());
            info!("Onion service ready at {}.onion:{}", address, CANON_PORT);
        } else {
            warn!("Onion service launched but address not yet available");
        }

        Ok(rend_requests.boxed())
    }

    /// Run the incoming connection accept loop for the onion service.
    ///
    /// Processes rendezvous requests from `start_onion_service`, accepts
    /// incoming streams, and spawns a handler for each connection. This
    /// runs until the runtime is shut down.
    pub async fn run_accept_loop(
        self: &Arc<Self>,
        rend_requests: futures::stream::BoxStream<'static, tor_hsservice::RendRequest>,
        graph: Arc<Mutex<GraphStore>>,
        server_config: Arc<ServerConfig>,
    ) {
        use futures::StreamExt;

        info!(
            "Starting onion service accept loop (max_concurrent={})",
            server_config.max_concurrent
        );

        // Enforce max_concurrent query limit via a semaphore
        let max_concurrent = server_config.max_concurrent.max(1) as usize;
        let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrent));

        let stream_requests = tor_hsservice::handle_rend_requests(rend_requests);
        futures::pin_mut!(stream_requests);

        // Shared rate limiter across all connections
        let rate_limiter = Arc::new(Mutex::new(RateLimiter::default()));

        while let Some(stream_request) = stream_requests.next().await {
            let state = *self.state.lock().await;
            if state == RuntimeState::ShuttingDown {
                break;
            }

            let graph = graph.clone();
            let config = server_config.clone();
            let sem = semaphore.clone();
            let rate_limiter = rate_limiter.clone();

            tokio::spawn(async move {
                // Acquire a permit before processing; drop it when done.
                let Ok(_permit) = sem.acquire().await else {
                    return; // Semaphore closed
                };

                match stream_request.accept(Connected::new_empty()).await {
                    Ok(mut data_stream) => {
                        if let Err(e) = crate::server::handle_connection_loop(
                            &mut data_stream,
                            &graph,
                            &rate_limiter,
                            &config,
                        )
                        .await
                        {
                            debug!("Connection handler error: {}", e);
                        }
                    }
                    Err(e) => {
                        debug!("Failed to accept stream: {}", e);
                    }
                }
            });
        }

        info!("Accept loop stopped");
    }

    /// Run the circuit pool maintenance loop.
    ///
    /// Manages warm circuits to peers: establishes `DataStream` connections
    /// to fill the pool, sends keepalive probes on held streams, and
    /// rotates circuits that have exceeded the rotation window. Runs
    /// until the runtime is shut down.
    pub async fn run_circuit_maintenance(
        self: &Arc<Self>,
        peers: Arc<Mutex<Vec<PeerRegistration>>>,
    ) {
        info!(
            "Starting circuit pool maintenance (size={}, keepalive={}s, rotation={}s)",
            self.config.circuit_pool_size,
            self.config.keepalive_interval_secs,
            self.config.circuit_rotation_secs
        );

        let mut pool: HashMap<[u8; 16], PooledCircuit> = HashMap::new();

        loop {
            let state = *self.state.lock().await;
            if state == RuntimeState::ShuttingDown {
                break;
            }

            let ts = now_ms();

            // 1. Rotate circuits that have exceeded the rotation window
            pool.retain(|node_id, circuit| {
                if circuit.needs_rotation(ts) {
                    debug!("Rotating circuit for peer {}", hex::encode(&node_id[..4]));
                    false
                } else {
                    true
                }
            });

            // 2. Send keepalive probes on held streams
            let needing_keepalive: Vec<[u8; 16]> = pool
                .iter()
                .filter(|(_, c)| c.needs_keepalive(ts))
                .map(|(id, _)| *id)
                .collect();

            for node_id in needing_keepalive {
                let Some(circuit) = pool.get_mut(&node_id) else {
                    continue;
                };

                if let Ok(()) = wire::write_keepalive(&mut circuit.stream).await {
                    circuit.last_keepalive = ts;
                    debug!("Keepalive OK for peer {}", hex::encode(&node_id[..4]));
                } else {
                    warn!(
                        "Keepalive failed for peer {}, removing",
                        hex::encode(&node_id[..4])
                    );
                    pool.remove(&node_id);
                }
            }

            // 3. Score and select new peers to fill empty pool slots
            let slots = self.config.circuit_pool_size.saturating_sub(pool.len());
            if slots > 0 {
                let peer_list = peers.lock().await;
                let mut scored: Vec<PeerScore> = peer_list
                    .iter()
                    .filter(|p| !p.is_expired(ts) && !pool.contains_key(&p.node_id))
                    .filter(|p| !p.onion_address.is_empty())
                    .map(|p| PeerScore::compute(p.clone(), &self.config.topics, None))
                    .collect();
                drop(peer_list);

                scored.sort_by(|a, b| {
                    b.composite
                        .partial_cmp(&a.composite)
                        .unwrap_or(std::cmp::Ordering::Equal)
                });

                for peer_score in scored.into_iter().take(slots) {
                    let peer = peer_score.peer;
                    match self.connect_to_peer(&peer.onion_address).await {
                        Ok(mut stream) => {
                            if wire::write_keepalive(&mut stream).await.is_ok() {
                                debug!(
                                    "Warm circuit established to peer {}",
                                    hex::encode(&peer.node_id[..4])
                                );
                                pool.insert(
                                    peer.node_id,
                                    PooledCircuit {
                                        peer,
                                        stream,
                                        created_at: ts,
                                        last_used: ts,
                                        last_keepalive: ts,
                                    },
                                );
                            }
                        }
                        Err(e) => {
                            debug!(
                                "Failed to connect to peer {}: {}",
                                hex::encode(&peer.node_id[..4]),
                                e
                            );
                        }
                    }
                }
            }

            debug!(
                "Circuit pool: {}/{} warm",
                pool.len(),
                self.config.circuit_pool_size
            );

            tokio::time::sleep(tokio::time::Duration::from_secs(
                self.config.keepalive_interval_secs,
            ))
            .await;
        }

        info!("Circuit maintenance loop stopped");
    }

    /// Initiate a graceful shutdown.
    pub async fn shutdown(&self) {
        info!("Shutting down Tor runtime");
        *self.state.lock().await = RuntimeState::ShuttingDown;
    }
}