netget-tool 0.3.5

//! DC (Direct Connect) server implementation
pub mod actions;

use crate::llm::action_helper::call_llm;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::Event;
use crate::server::connection::ConnectionId;
use crate::server::DcProtocol;
use crate::state::app_state::AppState;
use actions::DC_COMMAND_RECEIVED_EVENT;
use anyhow::Result;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc;

/// The `Pk=` value inside the `$Lock` challenge, and the hub name used when only
/// `hub_topic` was given — so a topic-only configuration still names the hub it belongs to
/// rather than presenting the topic as the name.
const DEFAULT_HUB_NAME: &str = "NetGetHub";

/// Largest single NMDC command accepted from a client, in bytes.
///
/// NMDC frames by `|`, and the accumulator grew until one arrived — with no cap and no timeout,
/// an **unauthenticated** peer that connects and never sends a `|` grows it as fast as its link
/// allows until the process dies. That is the whole attack: `nc host 411 < /dev/zero`. Real
/// commands are tens of bytes; `$MyINFO` with a long description is the largest and stays well
/// under a kilobyte, so 64 KiB is generous and still finite.
pub const MAX_COMMAND_LEN: usize = 64 * 1024;

/// How long to wait for the first byte from a client that has only connected.
///
/// NMDC is hub-speaks-first: the hub sends `$Lock` and the client answers with `$Key` and
/// `$ValidateNick` from inside its own connect path, with no human in the loop. A minute is far
/// longer than that takes. The `$Lock` itself is written before the read loop begins, so the
/// model's time over it — and a `manual` rule parking it for a human — is outside this bound by
/// construction.
const FIRST_BYTE_READ_TIMEOUT: Duration = Duration::from_secs(60);

/// How long to wait for further bytes once the client has sent a complete command.
///
/// Five minutes. An NMDC client in a hub is never silent for long: a bare `|` is the protocol's
/// keepalive and DC++ sends one about once a minute when it has nothing else to say, so five
/// minutes is five missed keepalives — a link that is gone, not a user who is quiet. Making it
/// longer would not help a real client and would only widen the window an unauthenticated peer
/// can hold.
///
/// The same bound covers a peer that has begun a command and stalled before its `|`, which is
/// the other way this loop can be held: [`MAX_COMMAND_LEN`] bounds how much such a peer can
/// make the hub buffer, and this bounds how long it can sit there having sent less.
///
/// The LLM round-trip and a `manual` rule parking a command for a human
/// (`src/state/intercepts.rs`, 300s by default) both happen after a complete command has been
/// read, so neither can be timed out from under itself.
const IDLE_BETWEEN_COMMANDS_TIMEOUT: Duration = Duration::from_secs(300);

/// Concurrent connections this hub admits.
///
/// [`crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS`]. Each connection may buffer up to
/// [`MAX_COMMAND_LEN`], so this is the multiplier that turns that per-connection bound into a
/// total one.
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;

/// What a peer over [`MAX_CONNECTIONS`] is told before the socket closes.
///
/// `$HubIsFull` is NMDC's own way for a hub to turn a client away, and it is what a client
/// expects in place of the `$Lock` it was waiting for. A client that does not recognise it sees
/// a well-formed, `|`-terminated command followed by a clean close rather than an unexplained
/// disconnection. Fixed text, so there is no placeholder an internal error could reach.
const CONNECTION_CAP_REFUSAL: &[u8] = b"$HubIsFull|";

/// Strip the NMDC framing characters from a value that goes inside a command.
///
/// Every command this hub writes is `format!("$Something {}|", value)`, so a `|` anywhere in
/// `value` ends the command early and everything after it is parsed by the client as a *second*
/// command — `$Kick`, `$ForceMove`, anything. `\r` and `\n` go for the same reason. The values
/// come from the model, and a hub that echoes a user's chat line back is putting peer-controlled
/// text through that format string.
///
/// `hub_name_command` has always done this for its two startup parameters; the ten action
/// executors did not. `send_dc_raw` is deliberately exempt — sending an arbitrary frame is its
/// entire declared purpose.
pub(crate) fn nmdc_field(value: &str) -> String {
    value
        .chars()
        .filter(|c| *c != '|' && *c != '\r' && *c != '\n')
        .collect()
}

/// Build the `$HubName` command from the `hub_name` / `hub_topic` startup parameters, or
/// `None` when neither was supplied.
///
/// `$HubName <name> - <topic>` is how NMDC carries a topic: there is no separate topic
/// command in the base protocol, and DC++ splits the name at the first ` - `.
///
/// Both values are sanitised. NMDC frames commands by `|`, so an unescaped pipe in either
/// would end the command early and leave the remainder to be parsed as a second one; `\r`
/// and `\n` are stripped for the same reason. Everything else is left alone — this is the
/// operator's text, not ours to rewrite.
fn hub_name_command(hub_name: Option<&str>, hub_topic: Option<&str>) -> Option<String> {
    fn clean(value: &str) -> String {
        value
            .chars()
            .filter(|c| *c != '|' && *c != '\r' && *c != '\n')
            .collect::<String>()
            .trim()
            .to_string()
    }

    let name = hub_name.map(clean).filter(|s| !s.is_empty());
    let topic = hub_topic.map(clean).filter(|s| !s.is_empty());

    match (name, topic) {
        (None, None) => None,
        (Some(name), None) => Some(format!("$HubName {name}|")),
        (name, Some(topic)) => {
            let name = name.unwrap_or_else(|| DEFAULT_HUB_NAME.to_string());
            Some(format!("$HubName {name} - {topic}|"))
        }
    }
}

/// DC server that forwards commands to LLM
pub struct DcServer;

impl DcServer {
    /// Spawn DC server with integrated LLM actions
    pub async fn spawn_with_llm_actions(
        listen_addr: SocketAddr,
        llm_client: OllamaClient,
        app_state: Arc<AppState>,
        status_tx: mpsc::UnboundedSender<String>,
        server_id: crate::state::ServerId,
        startup_params: Option<crate::protocol::StartupParams>,
    ) -> Result<SocketAddr> {
        // `hub_name` and `hub_topic` are declared as what this hub calls itself, and both
        // were read by nothing: the only identity a client ever saw was the hardcoded
        // `Pk=NetGetHub` inside the `$Lock` challenge, so setting either changed nothing an
        // NMDC client could observe.
        let hub_announcement = match &startup_params {
            Some(params) => hub_name_command(
                params.get_optional_string("hub_name")?.as_deref(),
                params.get_optional_string("hub_topic")?.as_deref(),
            ),
            None => None,
        };

        let listener =
            crate::server::socket_helpers::create_reusable_tcp_listener(listen_addr).await?;
        let local_addr = listener.local_addr()?;
        Log::new(Some(&status_tx)).info(format!(
            "DC server (action-based) listening on {}",
            local_addr
        ));

        let protocol = Arc::new(DcProtocol::new());

        let task_registrar = app_state.clone();
        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
        let accept_handle = tokio::spawn(async move {
            loop {
                match crate::server::accept_bounded::accept_bounded(
                    &listener,
                    &limiter,
                    CONNECTION_CAP_REFUSAL,
                    "DC",
                    Some(&status_tx),
                )
                .await
                {
                    Ok((stream, remote_addr, permit)) => {
                        let connection_id =
                            ConnectionId::new(app_state.get_next_unified_id().await);
                        let local_addr_conn = stream.local_addr().unwrap_or(local_addr);
                        let llm_clone = llm_client.clone();
                        let state_clone = app_state.clone();
                        let status_clone = status_tx.clone();
                        let protocol_clone = protocol.clone();
                        let hub_announcement = hub_announcement.clone();

                        // Tracked, not detached: stop_server must abort this task too.
                        let task_owner = app_state.clone();
                        task_owner
                            .spawn_server_task(server_id, async move {
                                // Held for the life of the connection: releasing it here would
                                // cap the accept rate rather than the number of live
                                // connections.
                                let _permit = permit;
                                let (read_half, write_half) = tokio::io::split(stream);
                                let write_half_arc = Arc::new(tokio::sync::Mutex::new(write_half));

                                // Add connection to ServerInstance
                                use crate::state::server::{
                                    ConnectionState as ServerConnectionState, ConnectionStatus,
                                    ProtocolConnectionInfo,
                                };
                                let now = crate::utils::clock::Instant::now();
                                let conn_state = ServerConnectionState {
                                    id: connection_id,
                                    remote_addr,
                                    local_addr: local_addr_conn,
                                    bytes_sent: 0,
                                    bytes_received: 0,
                                    packets_sent: 0,
                                    packets_received: 0,
                                    last_activity: now,
                                    status: ConnectionStatus::Active,
                                    status_changed_at: now,
                                    protocol_info: ProtocolConnectionInfo::empty(),
                                };
                                state_clone
                                    .add_connection_to_server(server_id, conn_state)
                                    .await;
                                let _ = status_clone.send("__UPDATE_UI__".to_string());

                                // Register the client in the protocol's own state map. This was
                                // never called from anywhere, so `DcProtocol::clients` stayed
                                // permanently empty and every accessor that goes through
                                // `get_mut` - set_nickname, set_client_info, set_operator - was a
                                // silent no-op, while get_nickname always returned None.
                                protocol_clone.add_connection(connection_id).await;

                                // Peer messaging: the dashboard's "message this peer" /
                                // "disconnect this peer" inject actions into THIS connection
                                // through the same executor the LLM path uses. Registered
                                // before the $Lock goes out so the operator can reach the
                                // connection from its first moment, including while a manual
                                // `*` rule parks the first command's answer.
                                let peer_rx = crate::server::peer_support::register_peer_channel(
                                    &state_clone,
                                    server_id,
                                    connection_id.as_u32(),
                                )
                                .await;
                                crate::server::peer_support::spawn_peer_command_task(
                                    peer_rx,
                                    protocol_clone.clone(),
                                    state_clone.clone(),
                                    server_id,
                                    connection_id.as_u32(),
                                    write_half_arc.clone(),
                                    status_clone.clone(),
                                );

                                Self::run_connection(
                                    read_half,
                                    &write_half_arc,
                                    connection_id,
                                    server_id,
                                    &llm_clone,
                                    &state_clone,
                                    &status_clone,
                                    &protocol_clone,
                                    hub_announcement.as_deref(),
                                )
                                .await;

                                // Every exit path - EOF, read error, close_connection, a failed
                                // $Lock write - lands here.
                                state_clone
                                    .remove_peer_handle(server_id, connection_id.as_u32())
                                    .await;
                                protocol_clone.remove_connection(&connection_id).await;
                                state_clone
                                    .remove_connection_from_server(server_id, connection_id)
                                    .await;
                                let _ = status_clone.send("__UPDATE_UI__".to_string());
                            })
                            .await;
                    }
                    Err(e) => {
                        Log::new(Some(&status_tx))
                            .error(format!("Failed to accept DC connection: {}", e));
                    }
                }
            }
        });

        // Register the accept loop so stop_server can abort it and release the port.
        task_registrar
            .register_server_task(server_id, accept_handle)
            .await;

        Ok(local_addr)
    }

    /// Send the $Lock challenge, then read pipe-delimited commands until the peer goes
    /// away or an action closes the connection. Returns on every exit so the caller runs
    /// the one cleanup path.
    #[allow(clippy::too_many_arguments)]
    async fn run_connection<R, W>(
        read_half: R,
        write_half_arc: &Arc<tokio::sync::Mutex<W>>,
        connection_id: ConnectionId,
        server_id: crate::state::ServerId,
        llm_clone: &OllamaClient,
        state_clone: &Arc<AppState>,
        status_clone: &mpsc::UnboundedSender<String>,
        protocol_clone: &Arc<DcProtocol>,
        hub_announcement: Option<&str>,
    ) where
        R: tokio::io::AsyncRead + Unpin,
        W: tokio::io::AsyncWrite + Unpin,
    {
        // Send initial $Lock challenge, followed by the hub's own name/topic when the
        // operator configured one. NMDC hubs announce `$HubName` right after `$Lock`, and
        // it is the first thing a client displays.
        let lock_command =
            format!("$Lock EXTENDEDPROTOCOLABCABCABCABCABCABC Pk={DEFAULT_HUB_NAME}|");
        {
            let mut writer = write_half_arc.lock().await;
            if let Err(e) = writer.write_all(lock_command.as_bytes()).await {
                Log::new(Some(status_clone)).error(format!("Failed to send initial Lock: {}", e));
                return;
            }
            Log::new(Some(status_clone))
                .debug(format!("DC sent {} bytes (Lock)", lock_command.len()));

            // Update stats
            state_clone
                .update_connection_stats(
                    server_id,
                    connection_id,
                    Some(0),
                    Some(lock_command.len() as u64),
                    Some(0),
                    Some(1),
                )
                .await;
        }

        // `$HubName` from the `hub_name` / `hub_topic` startup parameters. Sent only when
        // one of them was set, so a hub configured with neither puts exactly the bytes on
        // the wire it always did.
        if let Some(announcement) = hub_announcement {
            let mut writer = write_half_arc.lock().await;
            if let Err(e) = writer.write_all(announcement.as_bytes()).await {
                Log::new(Some(status_clone))
                    .error(format!("Failed to send initial HubName: {}", e));
                return;
            }
            Log::new(Some(status_clone)).debug(format!("DC sent HubName: {}", announcement));
            state_clone
                .update_connection_stats(
                    server_id,
                    connection_id,
                    Some(0),
                    Some(announcement.len() as u64),
                    Some(0),
                    Some(1),
                )
                .await;
        }

        // Buffered: this loop reads a byte at a time looking for `|`, and unbuffered that was
        // one syscall per byte of every command.
        let mut read_half = tokio::io::BufReader::new(read_half);
        let mut buffer = Vec::new();

        // Labelled so `close_connection` can leave the *connection* loop. It used to `break`
        // out of the inner `for protocol_result in ...` only, so the model asking to hang up
        // logged "DC closing connection N", skipped the remaining actions in that one batch,
        // and then went straight back to reading commands. The action's own description says
        // "Disconnect this client without sending anything further"; it did neither.
        let mut seen_a_command = false;
        'connection: loop {
            let mut byte = [0u8; 1];
            // The deadline wraps this read and nothing else. Everything that can legitimately
            // take minutes — the LLM round-trip, and a `manual` rule parking the command for a
            // human to answer — happens below, after a whole `|`-terminated command has been
            // read.
            let read_timeout = if seen_a_command {
                IDLE_BETWEEN_COMMANDS_TIMEOUT
            } else {
                FIRST_BYTE_READ_TIMEOUT
            };
            let read =
                match tokio::time::timeout(read_timeout, read_half.read_exact(&mut byte)).await {
                    Ok(read) => read,
                    Err(_) => {
                        Log::new(Some(status_clone)).info(format!(
                            "DC connection {} sent nothing for {}s; closing idle connection",
                            connection_id,
                            read_timeout.as_secs()
                        ));
                        break;
                    }
                };
            match read {
                Ok(_) => {
                    buffer.push(byte[0]);

                    // A client that never sends the `|` terminator must not be able to grow
                    // this buffer without limit; see MAX_COMMAND_LEN.
                    if buffer.len() > MAX_COMMAND_LEN {
                        Log::new(Some(status_clone)).warn(format!(
                            "DC connection {} sent {} bytes with no '|' terminator (limit {}); \
                             closing: decision=oversize_command",
                            connection_id,
                            buffer.len(),
                            MAX_COMMAND_LEN
                        ));
                        break;
                    }

                    // Check for pipe delimiter
                    if byte[0] == b'|' {
                        // We have a complete command
                        seen_a_command = true;
                        let command_bytes = buffer.clone();
                        buffer.clear();

                        // Convert to string
                        let command_str = match String::from_utf8(command_bytes.clone()) {
                            Ok(s) => s,
                            Err(e) => {
                                Log::new(Some(status_clone))
                                    .warn(format!("DC received non-UTF8 data: {}", e));
                                continue;
                            }
                        };

                        // Remove trailing pipe
                        let command = command_str.trim_end_matches('|');

                        // Byte-count summary and full command are FileOnly.
                        let preview = crate::utils::truncate_for_log(command, 100);
                        let log = Log::new(Some(status_clone));
                        log.debug(format!(
                            "DC received {} bytes on connection {}: {}",
                            command_bytes.len(),
                            connection_id,
                            preview
                        ));
                        log.trace(format!("DC command: {:?}", command));

                        // Update receive stats
                        state_clone
                            .update_connection_stats(
                                server_id,
                                connection_id,
                                Some(command_bytes.len() as u64),
                                Some(0),
                                Some(1),
                                Some(0),
                            )
                            .await;

                        // Record the nickname the client just claimed.
                        //
                        // Nothing ever called `set_nickname`, so
                        // `client_nickname` below was null on every event
                        // this hub has ever raised - the field is
                        // documented as carrying the nickname "if set",
                        // and it never was. NMDC announces it in exactly
                        // two places: `$ValidateNick <nick>` at login, and
                        // the second field of `$MyINFO $ALL <nick> ...`.
                        if let Some(rest) = command.strip_prefix("$ValidateNick ") {
                            if let Some(nick) = rest.split_whitespace().next() {
                                protocol_clone
                                    .set_nickname(connection_id, nick.to_string())
                                    .await;
                            }
                        } else if let Some(rest) = command.strip_prefix("$MyINFO ") {
                            let mut fields = rest.split_whitespace();
                            if fields.next() == Some("$ALL") {
                                if let Some(nick) = fields.next() {
                                    protocol_clone
                                        .set_nickname(connection_id, nick.to_string())
                                        .await;
                                }
                            }
                        }

                        // Parse command type
                        let command_type = if command.starts_with('$') {
                            command
                                .split_whitespace()
                                .next()
                                .unwrap_or("$Unknown")
                                .trim_start_matches('$')
                        } else if command.starts_with('<') {
                            "Chat"
                        } else {
                            "Unknown"
                        };

                        // Get client nickname if available
                        let client_nickname = protocol_clone.get_nickname(&connection_id).await;

                        let event = Event::new(
                            &DC_COMMAND_RECEIVED_EVENT,
                            serde_json::json!({
                                "command": command,
                                "command_type": command_type,
                                "client_nickname": client_nickname,
                            }),
                        );

                        Log::new(Some(status_clone))
                            .debug(format!("DC calling LLM for connection {}", connection_id));

                        let result = call_llm(
                            llm_clone,
                            state_clone,
                            server_id,
                            Some(connection_id),
                            &event,
                            protocol_clone.as_ref(),
                        )
                        .await;

                        match result {
                            Ok(execution_result) => {
                                let log = Log::new(Some(status_clone));
                                for message in &execution_result.messages {
                                    log.info(format!("{}", message));
                                }

                                log.debug(format!(
                                    "DC got {} protocol results",
                                    execution_result.protocol_results.len()
                                ));

                                // Answering with nothing is a real answer in NMDC - plenty
                                // of client commands ($Version, $MyINFO) want no reply - so
                                // it is tagged, not treated as a failure. The tag is what
                                // separates it in the log from the backend erroring below.
                                if execution_result.protocol_results.is_empty() {
                                    log.debug(format!(
                                        "DC connection {} decision=model_no_actions",
                                        connection_id
                                    ));
                                }

                                for protocol_result in execution_result.protocol_results {
                                    match protocol_result {
                                        ActionResult::Output(data) => {
                                            let mut write = write_half_arc.lock().await;
                                            if let Err(e) = write.write_all(&data).await {
                                                // The socket is gone; reading on is pointless
                                                // and every further LLM call would be spent on
                                                // answers nobody can receive.
                                                log.error(format!(
                                                    "Failed to write DC response: {}",
                                                    e
                                                ));
                                                break 'connection;
                                            }
                                            let _ = write.flush().await;
                                            drop(write);
                                            state_clone
                                                .update_connection_stats(
                                                    server_id,
                                                    connection_id,
                                                    None,
                                                    Some(data.len() as u64),
                                                    None,
                                                    Some(1),
                                                )
                                                .await;

                                            log.debug(format!(
                                                "DC sent {} bytes on connection {}",
                                                data.len(),
                                                connection_id
                                            ));
                                            log.trace(format!(
                                                "DC sent data: {:?}",
                                                String::from_utf8_lossy(&data)
                                            ));
                                        }
                                        ActionResult::CloseConnection => {
                                            log.info(format!(
                                                "DC closing connection {}: decision=model_close",
                                                connection_id
                                            ));
                                            break 'connection;
                                        }
                                        _ => {}
                                    }
                                }
                            }
                            Err(e) => {
                                let log = Log::new(Some(status_clone));
                                let failure = crate::utils::WireFailure::classify(&e);
                                let class = if failure.is_overloaded() {
                                    "overloaded"
                                } else {
                                    "unavailable"
                                };
                                // The whole error - backend URL, model name, anyhow chain -
                                // goes to the log and the status stream, where an operator
                                // looks. Nothing derived from it reaches the peer.
                                log.warn(format!(
                                    "DC command failed on connection {}: decision=fail_closed_llm_error class={} error={}",
                                    connection_id, class, e
                                ));

                                // Silence would leave a half-finished NMDC handshake
                                // hanging until the client's own timeout. The hub answers
                                // with a category instead, in two distinct NMDC shapes so
                                // the client can tell a transient overload (chat notice,
                                // session continues, retry is sensible) from a hard
                                // failure ($Error, the protocol's own error command).
                                let frame = match failure {
                                    crate::utils::WireFailure::Overloaded => {
                                        format!("<Hub> {}|", failure.prefixed_text())
                                    }
                                    crate::utils::WireFailure::Unavailable => {
                                        format!("$Error {}|", failure.prefixed_text())
                                    }
                                };
                                let bytes = frame.into_bytes();
                                let mut write = write_half_arc.lock().await;
                                if let Err(werr) = write.write_all(&bytes).await {
                                    log.error(format!(
                                        "Failed to write DC failure notice on connection {}: {}",
                                        connection_id, werr
                                    ));
                                } else {
                                    let _ = write.flush().await;
                                    drop(write);
                                    state_clone
                                        .update_connection_stats(
                                            server_id,
                                            connection_id,
                                            None,
                                            Some(bytes.len() as u64),
                                            None,
                                            Some(1),
                                        )
                                        .await;
                                }
                            }
                        }
                    }
                }
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::UnexpectedEof {
                        Log::new(Some(status_clone))
                            .debug(format!("DC connection {} closed by client", connection_id));
                    } else {
                        Log::new(Some(status_clone)).error(format!(
                            "DC read error on connection {}: {}",
                            connection_id, e
                        ));
                    }
                    break;
                }
            }
        }
    }
}