bevy_net_backend 0.1.0

//! The pending-request bookkeeping ([`InFlight`]) and the plugin's three systems. The ECS side is
//! the ONLY thing that answers requests: every request gets exactly one answer, whatever the
//! transport does.

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, Instant};

use bevy_app::AppExit;
use bevy_ecs::message::{MessageReader, MessageWriter};
use bevy_ecs::resource::Resource;
use bevy_ecs::system::{Commands, Res, ResMut};
use bevy_time::{Real, Time};
use http::header::HeaderName;
use http::Method;

use crate::client::{HttpClient, Queued, Route};
use crate::config::HttpConfig;
use crate::credentials::BackendCredentials;
use crate::request::{build_uri, OutgoingRequest, PreparedRequest, RequestId};
use crate::response::{BackendError, HttpResponse};
use crate::transport::{HttpTransportRes, HttpTransportResult};

/// How long after its own timeout a request is answered with a timeout by the plugin, in case the
/// transport never reports it (a stuck worker, a custom transport that forgets it).
pub const DEADLINE_GRACE: Duration = Duration::from_secs(5);

/// The kind of connection a pending request belongs to. `#[non_exhaustive]`: later versions may
/// add kinds.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum RequestKind {
    /// An HTTP request.
    Http,
    /// A request on a WebSocket connection (feature `ws`).
    WebSocket,
    /// A command on an SSH connection (feature `ssh`).
    Ssh,
    /// A file operation on an SSH connection (feature `sftp`).
    Sftp,
}

/// What a pending request is, for display (a "saving…" list, a debug overlay). Never holds a
/// query string, header value or body.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct RequestInfo {
    /// The kind of connection.
    pub kind: RequestKind,
    /// The HTTP method, for [`RequestKind::Http`].
    pub method: Option<Method>,
    /// What the request targets: for HTTP the URL path without the query (as the game wrote it),
    /// for WebSocket and SSH the connection's name (never an SSH command line: it may hold a
    /// secret).
    pub target: String,
}

struct Entry {
    route: Route,
    info: RequestInfo,
    deadline: Duration,
    allowed: Duration,
    generation: u64,
}

type Answer = (RequestId, Route, HttpTransportResult);

/// Which protocol's systems own a set of rows in [`InFlight`] (HTTP keeps its own map). Each
/// protocol replaces only its own rows.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum Protocol {
    #[cfg(feature = "ws")]
    WebSocket,
    #[cfg(feature = "ssh")]
    Ssh,
}

/// The cancel list every client shares: `HttpClient::cancel`, `WsClient::cancel` and
/// `SshClient::cancel` push into it. Each protocol's systems CLAIM only the ids they own (a request
/// they hold or have queued), in any order, and leave the rest. An id nobody claimed during two
/// `Send` phases in a row (an unknown or already answered id) is dropped by the core `Send`
/// system, which runs first; two phases, so an id pushed by a game system while a phase was
/// running still gets one full phase.
#[derive(Clone, Default)]
pub(crate) struct CancelList(Arc<Mutex<Vec<(RequestId, u8)>>>);

impl CancelList {
    fn lock(&self) -> std::sync::MutexGuard<'_, Vec<(RequestId, u8)>> {
        self.0.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// Ask for `id` to be cancelled (any protocol).
    pub(crate) fn push(&self, id: RequestId) {
        self.lock().push((id, 0));
    }

    /// Take the ids `owns` recognises, oldest first; leave the others.
    pub(crate) fn claim(&self, mut owns: impl FnMut(RequestId) -> bool) -> Vec<RequestId> {
        let mut list = self.lock();
        let mut claimed = Vec::new();
        list.retain(|(id, _)| {
            if owns(*id) {
                claimed.push(*id);
                false
            } else {
                true
            }
        });
        claimed
    }

    /// Start of a `Send` phase: drop ids that stayed unclaimed through two phases.
    pub(crate) fn age(&self) {
        self.lock().retain_mut(|(_, passes)| {
            *passes = passes.saturating_add(1);
            *passes < 3
        });
    }

    #[cfg(test)]
    pub(crate) fn len(&self) -> usize {
        self.lock().len()
    }
}

/// Every request waiting for its answer, HTTP, WebSocket and SSH alike (optional read-only
/// tracking, e.g. for a "saving…" spinner). An HTTP request appears here in `PostUpdate`
/// ([`BackendSystems::Send`](crate::BackendSystems::Send)) of the frame it was made in, a
/// WebSocket request or SSH command in the same place (also while it waits for its connection). A
/// request leaves it when it is answered; the answer message follows in `First` (of that frame, or
/// of the next one for answers decided in `PostUpdate`, such as a cancel).
#[derive(Resource)]
pub struct InFlight {
    entries: HashMap<RequestId, Entry>,
    /// Rows of the other protocols, each replaced only by its own systems.
    rows: HashMap<Protocol, HashMap<RequestId, RequestInfo>>,
    ready: Vec<Answer>,
    cancels: CancelList,
    epoch: Instant,
}

impl Default for InFlight {
    fn default() -> Self {
        Self { entries: HashMap::new(), rows: HashMap::new(), ready: Vec::new(), cancels: CancelList::default(), epoch: Instant::now() }
    }
}

impl fmt::Debug for InFlight {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("InFlight").field("pending", &self.len()).finish_non_exhaustive()
    }
}

impl InFlight {
    /// Whether `id` is still waiting for its answer.
    pub fn contains(&self, id: RequestId) -> bool {
        self.entries.contains_key(&id) || self.rows.values().any(|rows| rows.contains_key(&id))
    }

    /// How many requests are waiting.
    pub fn len(&self) -> usize {
        self.rows.values().fold(self.entries.len(), |n, rows| n.saturating_add(rows.len()))
    }

    /// Whether nothing is waiting.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty() && self.rows.values().all(HashMap::is_empty)
    }

    /// The waiting ids, oldest first.
    pub fn ids(&self) -> Vec<RequestId> {
        let mut ids: Vec<RequestId> = self.entries.keys().chain(self.rows.values().flat_map(HashMap::keys)).copied().collect();
        ids.sort_unstable();
        ids
    }

    /// What a waiting request is (kind, method, target without query / connection name).
    pub fn describe(&self, id: RequestId) -> Option<&RequestInfo> {
        self.entries.get(&id).map(|e| &e.info).or_else(|| self.rows.values().find_map(|rows| rows.get(&id)))
    }

    fn now(&self, time: Option<&Time<Real>>) -> Duration {
        time.map_or_else(|| self.epoch.elapsed(), Time::elapsed)
    }

    pub(crate) fn cancel_list(&self) -> CancelList {
        self.cancels.clone()
    }

    /// Take the cancel ids `owns` recognises (see [`CancelList`]); the others stay for the other
    /// protocols.
    #[cfg_attr(not(any(feature = "ws", feature = "ssh")), allow(dead_code))]
    pub(crate) fn claim_cancels(&self, owns: impl FnMut(RequestId) -> bool) -> Vec<RequestId> {
        self.cancels.claim(owns)
    }

    /// Replace the rows of `protocol` (its systems call this after every change); the rows of the
    /// other protocols are untouched.
    #[cfg_attr(not(any(feature = "ws", feature = "ssh")), allow(dead_code))]
    pub(crate) fn set_rows(&mut self, protocol: Protocol, rows: impl IntoIterator<Item = (RequestId, RequestInfo)>) {
        let map = self.rows.entry(protocol).or_default();
        map.clear();
        map.extend(rows);
    }
}

/// Defaults, credentials and the URL applied to a game's request.
pub(crate) fn prepare(request: OutgoingRequest, config: &HttpConfig, credentials: Option<&BackendCredentials>) -> Result<PreparedRequest, BackendError> {
    let mut request = request;
    if let Some(error) = request.take_error() {
        return Err(error);
    }
    config.validate().map_err(|e| BackendError::InvalidRequest(e.to_string()))?;
    check_method(request.method(), request.body().is_some())?;
    let own: HashSet<HeaderName> = request.headers().keys().cloned().collect();
    for (name, value) in config.headers() {
        if !own.contains(name) {
            request.headers_mut().append(name.clone(), value.clone());
        }
    }
    if request.uses_credentials() {
        if let Some(credentials) = credentials {
            credentials.apply(&mut request);
        }
    }
    if let Some(error) = request.take_error() {
        return Err(error);
    }
    let uri = build_uri(config.base_url(), request.path(), request.query(), config.insecure_http_allowed())?;
    let (method, headers, body, timeout, purpose) = request.into_parts();
    Ok(PreparedRequest { method, uri, headers, body, timeout: timeout.unwrap_or(config.timeout()), max_body_bytes: config.max_body_bytes(), purpose })
}

/// The methods sent: the standard ones except `CONNECT` (a tunnel, not an API call). `HEAD` has no
/// body. Anything else would fail inside the HTTP client after the request counted as "sent".
fn check_method(method: &Method, has_body: bool) -> Result<(), BackendError> {
    let standard = [Method::GET, Method::POST, Method::PUT, Method::PATCH, Method::DELETE, Method::HEAD, Method::OPTIONS, Method::TRACE];
    if !standard.contains(method) {
        return Err(BackendError::InvalidRequest(format!("method `{method}` is not supported (standard methods only, no CONNECT)")));
    }
    if has_body && *method == Method::HEAD {
        return Err(BackendError::InvalidRequest("a HEAD request cannot have a body".into()));
    }
    Ok(())
}

/// `PostUpdate` ([`BackendSystems::Send`](crate::BackendSystems::Send)): hand every queued request
/// to the transport, or put its answer aside when it cannot be sent; apply cancels.
///
/// In a frame with an `AppExit` message nothing is submitted: the queue is left to
/// [`shutdown_on_exit`], which answers it with `Shutdown`, and a request answered `Shutdown` must
/// never reach the server.
#[allow(clippy::too_many_arguments)]
pub(crate) fn send_requests(
    client: Res<HttpClient>,
    mut inflight: ResMut<InFlight>,
    mut transport: Option<ResMut<HttpTransportRes>>,
    config: Res<HttpConfig>,
    credentials: Option<Res<BackendCredentials>>,
    time: Option<Res<Time<Real>>>,
    mut exit: MessageReader<AppExit>,
) {
    // This system runs first in the phase: drop cancel ids nobody claimed for two phases.
    inflight.cancels.age();
    if exit.read().count() > 0 {
        return;
    }
    let queued = client.drain();
    // Only HTTP's own ids (a request queued this frame or waiting for its answer); the rest stays
    // in the shared list for the other protocols.
    let queued_ids: HashSet<RequestId> = queued.iter().map(|Queued::Send { id, .. }| *id).collect();
    let cancels = {
        let entries = &inflight.entries;
        inflight.cancels.claim(|id| queued_ids.contains(&id) || entries.contains_key(&id))
    };
    if queued.is_empty() && cancels.is_empty() {
        return;
    }
    let now = inflight.now(time.as_deref());
    // A cancel for a request queued in the same frame: answer it without ever submitting it.
    let cancelled: HashSet<RequestId> = cancels.iter().copied().collect();
    let mut claimed: HashSet<RequestId> = HashSet::new();
    for item in queued {
        match item {
            Queued::Send { id, route, .. } if cancelled.contains(&id) => {
                claimed.insert(id);
                inflight.ready.push((id, route, Err(BackendError::Cancelled)));
            }
            Queued::Send { id, request, route } => {
                let info = RequestInfo { kind: RequestKind::Http, method: Some(request.method().clone()), target: request.path().to_string() };
                match prepare(*request, &config, credentials.as_deref()) {
                    Err(error) => inflight.ready.push((id, route, Err(error))),
                    Ok(prepared) => match transport.as_mut() {
                        None => inflight.ready.push((id, route, Err(BackendError::NoTransport))),
                        Some(transport) => {
                            let allowed = prepared.timeout.saturating_add(DEADLINE_GRACE);
                            if let Some(method) = &info.method {
                                tracing::debug!(">>> NET-BACKEND: {id} {method} {}", info.target);
                            }
                            inflight
                                .entries
                                .insert(id, Entry { route, info, deadline: now.saturating_add(allowed), allowed, generation: transport.generation() });
                            transport.get_mut().submit(id, prepared);
                        }
                    },
                }
            }
        }
    }
    for id in cancels {
        if claimed.contains(&id) {
            continue;
        }
        if let Some(entry) = inflight.entries.remove(&id) {
            if let Some(transport) = transport.as_mut() {
                transport.get_mut().cancel(id);
            }
            inflight.ready.push((id, entry.route, Err(BackendError::Cancelled)));
        }
    }
}

/// `First` ([`BackendSystems::Receive`](crate::BackendSystems::Receive)): collect what the
/// transport reports, answer requests that can no longer be answered by it (deadline passed,
/// transport gone), and write every answer as a message for this frame.
pub(crate) fn receive_answers(
    mut inflight: ResMut<InFlight>,
    mut transport: Option<ResMut<HttpTransportRes>>,
    config: Res<HttpConfig>,
    time: Option<Res<Time<Real>>>,
    mut raw: MessageWriter<HttpResponse>,
    mut commands: Commands,
) {
    let mut answers = std::mem::take(&mut inflight.ready);
    let generation = transport.as_ref().map(|t| t.generation());
    if let Some(transport) = transport.as_mut() {
        for (id, result) in transport.get_mut().poll() {
            match inflight.entries.get(&id) {
                Some(entry) if Some(entry.generation) == generation => {
                    if let Some(entry) = inflight.entries.remove(&id) {
                        answers.push((id, entry.route, result));
                    }
                }
                _ => tracing::debug!(">>> NET-BACKEND: {id}: a late result was discarded (already answered)"),
            }
        }
    }
    if !inflight.entries.is_empty() {
        let now = inflight.now(time.as_deref());
        let gone: Vec<(RequestId, Entry)> = inflight.entries.extract_if(|_, entry| Some(entry.generation) != generation).collect();
        answers.extend(gone.into_iter().map(|(id, entry)| (id, entry.route, Err(BackendError::NoTransport))));
        let late: Vec<(RequestId, Entry)> = inflight.entries.extract_if(|_, entry| entry.deadline <= now).collect();
        for (id, entry) in late {
            if let Some(transport) = transport.as_mut() {
                transport.get_mut().cancel(id);
            }
            answers.push((id, entry.route, Err(BackendError::Timeout(format!("no answer from the transport within {:?}", entry.allowed)))));
        }
    }
    deliver(answers, config.max_body_bytes(), &mut raw, &mut commands);
}

/// `Last` on `AppExit` ([`BackendSystems::Exit`](crate::BackendSystems::Exit)): answer
/// everything still open (results that already arrived as they are, the rest with
/// [`BackendError::Shutdown`]) and shut the transport down without waiting for it.
pub(crate) fn shutdown_on_exit(
    client: Res<HttpClient>,
    mut inflight: ResMut<InFlight>,
    mut transport: Option<ResMut<HttpTransportRes>>,
    config: Res<HttpConfig>,
    mut raw: MessageWriter<HttpResponse>,
    mut commands: Commands,
) {
    let mut answers = std::mem::take(&mut inflight.ready);
    let queued = client.drain();
    let queued_ids: HashSet<RequestId> = queued.iter().map(|Queued::Send { id, .. }| *id).collect();
    let cancelled: HashSet<RequestId> = {
        let entries = &inflight.entries;
        inflight.cancels.claim(|id| queued_ids.contains(&id) || entries.contains_key(&id)).into_iter().collect()
    };
    for Queued::Send { id, route, .. } in queued {
        let error = if cancelled.contains(&id) { BackendError::Cancelled } else { BackendError::Shutdown };
        answers.push((id, route, Err(error)));
    }
    for id in &cancelled {
        if let Some(entry) = inflight.entries.remove(id) {
            answers.push((*id, entry.route, Err(BackendError::Cancelled)));
        }
    }
    let generation = transport.as_ref().map(|t| t.generation());
    if let Some(transport) = transport.as_mut() {
        for (id, result) in transport.get_mut().poll() {
            if inflight.entries.get(&id).is_some_and(|e| Some(e.generation) == generation) {
                if let Some(entry) = inflight.entries.remove(&id) {
                    answers.push((id, entry.route, result));
                }
            }
        }
    }
    let open = inflight.entries.len();
    answers.extend(inflight.entries.drain().map(|(id, entry)| (id, entry.route, Err(BackendError::Shutdown))));
    if let Some(transport) = transport.as_mut() {
        transport.get_mut().shutdown();
    }
    if open > 0 {
        tracing::info!(">>> NET-BACKEND: app exit: {open} open request(s) answered with Shutdown");
    }
    deliver(answers, config.max_body_bytes(), &mut raw, &mut commands);
}

/// Turn transport results into final answers (status and body-limit rules) and write them, in
/// request order.
fn deliver(mut answers: Vec<Answer>, limit: u64, raw: &mut MessageWriter<HttpResponse>, commands: &mut Commands) {
    answers.sort_by_key(|(id, _, _)| *id);
    for (id, route, result) in answers {
        let result = result.and_then(|response| {
            if u64::try_from(response.body.len()).unwrap_or(u64::MAX) > limit {
                Err(BackendError::BodyTooLarge { limit })
            } else if response.status.is_success() {
                Ok(response)
            } else {
                Err(BackendError::Status(Box::new(response)))
            }
        });
        match &result {
            Ok(response) => tracing::debug!(">>> NET-BACKEND: {id} -> {}", response.status),
            Err(
                error @ (BackendError::InvalidRequest(_)
                | BackendError::InsecureHttp { .. }
                | BackendError::Encode(_)
                | BackendError::RequestTooLarge { .. }
                | BackendError::NoTransport),
            ) => {
                tracing::warn!(">>> NET-BACKEND: {id} not sent: {error}")
            }
            Err(error) => tracing::debug!(">>> NET-BACKEND: {id} -> {error}"),
        }
        match route {
            Route::Raw => {
                raw.write(HttpResponse { id, result });
            }
            #[cfg(feature = "json")]
            Route::Json(json) => json.deliver(id, result, commands),
        }
    }
    #[cfg(not(feature = "json"))]
    let _ = commands;
}