//! RFC 7395 WebSocket adapter to the shared XMPP stream engine.
use crate::wire::{Framer, escape};
use axum::{
    Router,
    extract::{
        State, WebSocketUpgrade,
        ws::{Message, WebSocket},
    },
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::get,
};
use futures_util::{SinkExt, StreamExt};
use std::{future::Future, pin::Pin, sync::Arc};
use tokio::io::{AsyncWriteExt, DuplexStream};
pub type ConnectionHandler =
    Arc<dyn Fn(DuplexStream) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
/// The effective HTTP connection context is supplied by the HTTP listener.
pub type ContextConnectionHandler = Arc<
    dyn Fn(
            DuplexStream,
            Option<crate::http_proxy::RequestContext>,
        ) -> Pin<Box<dyn Future<Output = ()> + Send>>
        + Send
        + Sync,
>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HttpRoute {
    Websocket,
    Bosh,
    Altconnect,
}

#[derive(Clone)]
struct ConfiguredState {
    handler: ContextConnectionHandler,
    bosh: BoshState,
    config: crate::http_config::HttpConfig,
    websocket: bool,
    bosh_enabled: bool,
    secure: bool,
}
/// A fallback dispatcher uses literal mount paths, avoiding route-pattern conflicts and panics.
pub fn configured_router(
    handler: ContextConnectionHandler,
    config: crate::http_config::HttpConfig,
    websocket: bool,
    bosh: bool,
    secure: bool,
) -> anyhow::Result<Router> {
    configured_router_with_fallback(handler, config, websocket, bosh, secure, Router::new())
}
/// Compose transport mounts with an existing HTTP app without merging fallbacks.
pub fn configured_router_with_fallback(
    handler: ContextConnectionHandler,
    config: crate::http_config::HttpConfig,
    websocket: bool,
    bosh: bool,
    secure: bool,
    fallback: Router,
) -> anyhow::Result<Router> {
    anyhow::ensure!(
        !(websocket && bosh && config.websocket_path == config.bosh_path),
        "websocket and bosh HTTP paths conflict"
    );
    let legacy: ConnectionHandler = {
        let handler = handler.clone();
        Arc::new(move |stream| handler(stream, None))
    };
    let mut bosh_state = new_bosh_state(legacy);
    bosh_state.context_handler = Some(handler.clone());
    let state = ConfiguredState {
        handler,
        bosh: bosh_state,
        config,
        websocket,
        bosh_enabled: bosh,
        secure,
    };
    Ok(fallback
        .layer(axum::middleware::from_fn_with_state(
            state,
            configured_middleware,
        ))
        .layer(axum::extract::DefaultBodyLimit::max(256 * 1024)))
}
async fn configured_middleware(
    State(state): State<ConfiguredState>,
    request: axum::extract::Request,
    next: axum::middleware::Next,
) -> Response {
    let path = request.uri().path();
    let matches =
        |mount: &str| path == mount || (mount != "/" && path.strip_suffix('/') == Some(mount));
    if (state.websocket && matches(&state.config.websocket_path))
        || (state.bosh_enabled && matches(&state.config.bosh_path))
    {
        let route = if state.websocket && matches(&state.config.websocket_path) {
            HttpRoute::Websocket
        } else {
            HttpRoute::Bosh
        };
        let mut response = configured_request(State(state), request).await;
        response.extensions_mut().insert(route);
        response
    } else {
        next.run(request).await
    }
}
async fn configured_request(
    State(state): State<ConfiguredState>,
    request: axum::extract::Request,
) -> Response {
    use axum::extract::{FromRequest, FromRequestParts};
    let host = request
        .headers()
        .get(axum::http::header::HOST)
        .and_then(|v| v.to_str().ok());
    let Some(host) = host else {
        return StatusCode::BAD_REQUEST.into_response();
    };
    if !state.config.accepts_host(host) {
        return StatusCode::NOT_FOUND.into_response();
    }
    let path = request.uri().path();
    let matches =
        |mount: &str| path == mount || (mount != "/" && path.strip_suffix('/') == Some(mount));
    let websocket = state.websocket && matches(&state.config.websocket_path);
    let bosh = state.bosh_enabled && matches(&state.config.bosh_path);
    if !websocket && !bosh {
        return StatusCode::NOT_FOUND.into_response();
    }
    let context = request
        .extensions()
        .get::<crate::http_proxy::RequestContext>()
        .cloned();
    let secure = context.as_ref().map(|c| c.secure).unwrap_or(state.secure);
    let considered = if websocket {
        state.config.consider_websocket_secure
    } else {
        state.config.consider_bosh_secure
    };
    if !secure && !considered {
        return StatusCode::FORBIDDEN.into_response();
    }
    let origins = if websocket {
        &state.config.websocket_origins
    } else {
        &state.config.bosh_origins
    };
    let origin = request.headers().get(axum::http::header::ORIGIN).cloned();
    if let Some(origin) = &origin
        && !origin
            .to_str()
            .ok()
            .is_some_and(|o| origins.iter().any(|a| a == o || a == "*"))
    {
        return StatusCode::FORBIDDEN.into_response();
    }
    let method = request.method().clone();
    let mut response = if method == axum::http::Method::OPTIONS {
        preflight(if websocket {
            "GET, OPTIONS"
        } else {
            "POST, OPTIONS"
        })
    } else if websocket && method == axum::http::Method::GET {
        let (mut parts, _) = request.into_parts();
        match WebSocketUpgrade::from_request_parts(&mut parts, &()).await {
            Ok(ws) => {
                let handler: ConnectionHandler =
                    Arc::new(move |stream| (state.handler)(stream, context));
                upgrade(State(handler), parts.headers, ws).await
            }
            Err(e) => e.into_response(),
        }
    } else if bosh && method == axum::http::Method::POST {
        match Bytes::from_request(request, &()).await {
            Ok(bytes) => bosh_with_context(state.bosh, bytes, context).await,
            Err(e) => e.into_response(),
        }
    } else {
        (
            StatusCode::METHOD_NOT_ALLOWED,
            [(
                axum::http::header::ALLOW,
                if websocket {
                    "GET, OPTIONS"
                } else {
                    "POST, OPTIONS"
                },
            )],
        )
            .into_response()
    };
    response
        .headers_mut()
        .remove(axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN);
    if let Some(origin) = origin {
        response
            .headers_mut()
            .insert(axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN, origin);
        response.headers_mut().append(
            axum::http::header::VARY,
            axum::http::HeaderValue::from_static("Origin"),
        );
    }
    response
}
pub fn websocket_router(handler: ConnectionHandler) -> Router {
    let route = get(upgrade).options(|| async { preflight("GET, OPTIONS") });
    Router::new()
        .route("/xmpp-websocket", route.clone())
        .route("/xmpp-websocket/", route)
        .with_state(handler)
}
fn preflight(methods: &'static str) -> Response {
    (
        [
            ("access-control-allow-origin", "*"),
            ("access-control-allow-methods", methods),
            ("access-control-allow-headers", "Content-Type"),
        ],
        "",
    )
        .into_response()
}
async fn upgrade(
    State(handler): State<ConnectionHandler>,
    headers: axum::http::HeaderMap,
    ws: WebSocketUpgrade,
) -> Response {
    if !headers
        .get("sec-websocket-protocol")
        .and_then(|v| v.to_str().ok())
        .is_some_and(|p| p.split(',').any(|p| p.trim() == "xmpp"))
    {
        return StatusCode::BAD_REQUEST.into_response();
    }
    ws.protocols(["xmpp"])
        .max_message_size(256 * 1024)
        .max_frame_size(256 * 1024)
        .on_upgrade(move |socket| bridge(socket, handler))
}
fn from_websocket(input: &str) -> anyhow::Result<String> {
    let d = roxmltree::Document::parse(input)?;
    let n = d.root_element();
    if n.tag_name().namespace() == Some("urn:ietf:params:xml:ns:xmpp-framing") {
        match n.tag_name().name() {
            "open" => {
                let mut out="<stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams'".to_string();
                for key in ["to", "from", "id", "version"] {
                    if let Some(v) = n.attribute(key) {
                        out.push_str(&format!(" {key}='{}'", escape(v)));
                    }
                }
                out.push('>');
                Ok(out)
            }
            "close" => Ok("</stream:stream>".into()),
            _ => anyhow::bail!("unknown framing element"),
        }
    } else {
        anyhow::ensure!(
            matches!(
                n.tag_name().namespace(),
                Some(
                    "jabber:client"
                        | "urn:ietf:params:xml:ns:xmpp-sasl"
                        | "urn:xmpp:sasl:2"
                        | "urn:xmpp:sm:3"
                        | "urn:xmpp:csi:0"
                )
            ),
            "unexpected websocket namespace"
        );
        Ok(input.into())
    }
}
fn to_websocket(input: &str) -> anyhow::Result<Option<String>> {
    if input.starts_with("<?xml") {
        return Ok(None);
    }
    if input == "</stream:stream>" {
        return Ok(Some(
            "<close xmlns='urn:ietf:params:xml:ns:xmpp-framing'/>".into(),
        ));
    }
    if input.starts_with("<stream:stream ") {
        let xml = format!("{input}</stream:stream>");
        let d = roxmltree::Document::parse(&xml)?;
        let n = d.root_element();
        let mut out = "<open xmlns='urn:ietf:params:xml:ns:xmpp-framing'".to_string();
        for key in ["to", "from", "id", "version"] {
            if let Some(v) = n.attribute(key) {
                out.push_str(&format!(" {key}='{}'", escape(v)));
            }
        }
        out.push_str("/>");
        return Ok(Some(out));
    }
    let wrapped = format!(
        "<root xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams'>{input}</root>"
    );
    let d = roxmltree::Document::parse(&wrapped)?;
    let n = d
        .root_element()
        .first_element_child()
        .ok_or_else(|| anyhow::anyhow!("no XML"))?;
    // Every message is an independent XML document. Serialize the parsed
    // subtree so inherited namespaces also survive on children and attributes.
    Ok(Some(crate::protocol::xml(n)))
}
async fn bridge(socket: WebSocket, handler: ConnectionHandler) {
    let (internal, transport) = tokio::io::duplex(256 * 1024);
    let mut worker = tokio::spawn(handler(internal));
    let (mut read, mut write) = tokio::io::split(transport);
    let (mut tx, mut rx) = socket.split();
    let incoming = async {
        while let Some(Ok(message)) = rx.next().await {
            match message {
                Message::Text(text) => {
                    let text = from_websocket(&text)?;
                    write.write_all(text.as_bytes()).await?;
                }
                Message::Close(_) => break,
                Message::Binary(_) => anyhow::bail!("binary XMPP frames forbidden"),
                _ => {}
            }
        }
        anyhow::Ok(())
    };
    let outgoing = async {
        let mut framer = Framer::default();
        loop {
            let frame = framer.next(&mut read).await?;
            if let Some(text) = to_websocket(&frame)? {
                tx.send(Message::Text(text.into())).await?;
            }
            if frame == "</stream:stream>" {
                break;
            }
        }
        anyhow::Ok(())
    };
    tokio::select! {_=incoming=>{},_=outgoing=>{}}
    let _ = tokio::time::timeout(Duration::from_secs(2), tx.close()).await;
    // Deliver buffered XMPP close frames before retiring the engine. EOF alone
    // remains an ungraceful disconnect, allowing stream management to hibernate.
    let _ = write.shutdown().await;
    if tokio::time::timeout(Duration::from_secs(5), &mut worker)
        .await
        .is_err()
    {
        worker.abort();
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    fn configured_state(secure: bool) -> ConfiguredState {
        let config = serde_json::from_value(serde_json::json!({"global":{"http_paths":{"websocket":"/custom/ws","bosh":"/custom/bind"},"http_cors_override":{"websocket":{"origins":["https://allowed.test"]},"bosh":{"origins":["https://allowed.test"]}}},"hosts":{"a.test":{"options":{}}}})).unwrap();
        let handler: ContextConnectionHandler =
            Arc::new(|_, _| Box::pin(async { panic!("no stream expected") }));
        let legacy: ConnectionHandler =
            Arc::new(|_| Box::pin(async { panic!("no stream expected") }));
        ConfiguredState {
            handler,
            bosh: new_bosh_state(legacy),
            config: crate::http_config::HttpConfig::for_host(&config, "a.test").unwrap(),
            websocket: true,
            bosh_enabled: true,
            secure,
        }
    }
    async fn configured_status(
        state: ConfiguredState,
        method: &str,
        path: &str,
        host: &str,
        origin: Option<&str>,
        context: Option<crate::http_proxy::RequestContext>,
    ) -> Response {
        let mut request = axum::http::Request::builder()
            .method(method)
            .uri(path)
            .header("host", host);
        if let Some(origin) = origin {
            request = request.header("origin", origin);
        }
        let mut request = request.body(axum::body::Body::empty()).unwrap();
        if let Some(context) = context {
            request.extensions_mut().insert(context);
        }
        configured_request(State(state), request).await
    }
    #[tokio::test]
    async fn configured_mount_host_origin_and_secure_dispatch() {
        assert_eq!(
            configured_status(
                configured_state(true),
                "OPTIONS",
                "/custom/ws/",
                "a.test",
                Some("https://allowed.test"),
                None
            )
            .await
            .status(),
            StatusCode::OK
        );
        assert_eq!(
            configured_status(
                configured_state(true),
                "OPTIONS",
                "/xmpp-websocket",
                "a.test",
                None,
                None
            )
            .await
            .status(),
            StatusCode::NOT_FOUND
        );
        assert_eq!(
            configured_status(
                configured_state(true),
                "OPTIONS",
                "/custom/bind",
                "unknown.test",
                None,
                None
            )
            .await
            .status(),
            StatusCode::NOT_FOUND
        );
        assert_eq!(
            configured_status(
                configured_state(true),
                "OPTIONS",
                "/custom/bind",
                "a.test",
                Some("https://bad.test"),
                None
            )
            .await
            .status(),
            StatusCode::FORBIDDEN
        );
        assert_eq!(
            configured_status(
                configured_state(false),
                "OPTIONS",
                "/custom/bind",
                "a.test",
                None,
                None
            )
            .await
            .status(),
            StatusCode::FORBIDDEN
        );
        let context = crate::http_proxy::RequestContext {
            socket_peer: "127.0.0.1:1234".parse().unwrap(),
            peer_ip: "198.51.100.1".parse().unwrap(),
            secure: true,
        };
        let response = configured_status(
            configured_state(false),
            "OPTIONS",
            "/custom/bind",
            "a.test",
            Some("https://allowed.test"),
            Some(context),
        )
        .await;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response.headers()[axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN],
            "https://allowed.test"
        );
        assert_eq!(
            configured_status(
                configured_state(true),
                "GET",
                "/custom/bind",
                "a.test",
                None,
                None
            )
            .await
            .status(),
            StatusCode::METHOD_NOT_ALLOWED
        );
    }
    #[test]
    fn independent_framing() {
        assert!(from_websocket("<open xmlns='urn:ietf:params:xml:ns:xmpp-framing' to='example.test' version='1.0'/>").unwrap().starts_with("<stream:stream"));
        assert!(from_websocket("<message/>").is_err());
        assert!(from_websocket("<inactive xmlns='urn:xmpp:csi:0'/>").is_ok());
        assert!(from_websocket("<active xmlns='urn:xmpp:csi:0'/>").is_ok());
        assert!(to_websocket("<stream:features><mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl'/></stream:features>").unwrap().unwrap().contains("xmlns:stream="));
        let xml = to_websocket("<message to='a'><body>x</body></message>")
            .unwrap()
            .unwrap();
        assert_eq!(
            roxmltree::Document::parse(&xml)
                .unwrap()
                .root_element()
                .tag_name()
                .namespace(),
            Some("jabber:client")
        );
    }
    #[test]
    fn websocket_preserves_inherited_namespaces_and_spaced_declarations() {
        for input in [
            "<message xmlns = 'jabber:client'><body>hello</body></message>",
            "<stream:features><message><body>hello</body></message></stream:features>",
            "<message><body stream:test='value'>hello</body></message>",
        ] {
            let xml = to_websocket(input).unwrap().unwrap();
            let doc = roxmltree::Document::parse(&xml).unwrap();
            let body = doc.descendants().find(|n| n.has_tag_name("body")).unwrap();
            assert_eq!(body.tag_name().namespace(), Some("jabber:client"));
            if input.contains("stream:test") {
                assert_eq!(
                    body.attribute(("http://etherx.jabber.org/streams", "test")),
                    Some("value")
                );
            }
        }
    }
    #[tokio::test]
    async fn bosh_processing_capacity_is_bounded() {
        let requests = Arc::new(tokio::sync::Semaphore::new(256));
        let _occupied = requests.clone().acquire_many_owned(256).await.unwrap();
        let state = BoshState {
            handler: Arc::new(|_| Box::pin(async { panic!("must not create a session") })),
            sessions: Arc::new(Mutex::new(BTreeMap::new())),
            requests,
            context_handler: None,
        };
        let response = bosh(State(state), Bytes::from_static(b"<body/>")).await;
        let bytes = axum::body::to_bytes(response.into_body(), 1024)
            .await
            .unwrap();
        assert!(
            std::str::from_utf8(&bytes)
                .unwrap()
                .contains("resource-constraint")
        );
    }
}

use axum::{body::Bytes, routing::post};
use std::{
    collections::{BTreeMap, HashSet},
    time::{Duration, Instant},
};
use tokio::sync::{Mutex, Notify, mpsc};
const BOSH: &str = "http://jabber.org/protocol/httpbind";
struct Sequence {
    next: u64,
    pending: HashSet<u64>,
    responses: BTreeMap<u64, String>,
}
struct BoshSession {
    domain: String,
    wait: u64,
    sequence: Mutex<Sequence>,
    changed: Notify,
    poll_wakeup: Notify,
    writer: Mutex<tokio::io::WriteHalf<DuplexStream>>,
    output: Mutex<mpsc::Receiver<String>>,
    last: std::sync::Mutex<Instant>,
    tasks: Vec<tokio::task::JoinHandle<()>>,
    engine_done: Arc<Notify>,
}
impl Drop for BoshSession {
    fn drop(&mut self) {
        for task in &self.tasks {
            task.abort();
        }
    }
}
#[derive(Clone)]
struct BoshState {
    handler: ConnectionHandler,
    context_handler: Option<ContextConnectionHandler>,
    sessions: Arc<Mutex<BTreeMap<String, Arc<BoshSession>>>>,
    requests: Arc<tokio::sync::Semaphore>,
}
fn new_bosh_state(handler: ConnectionHandler) -> BoshState {
    let state = BoshState {
        handler,
        context_handler: None,
        sessions: Arc::new(Mutex::new(BTreeMap::new())),
        requests: Arc::new(tokio::sync::Semaphore::new(256)),
    };
    let weak = Arc::downgrade(&state.sessions);
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(30)).await;
            let Some(map) = weak.upgrade() else { break };
            map.lock()
                .await
                .retain(|_, s| s.last.lock().unwrap().elapsed() < Duration::from_secs(120));
        }
    });
    state
}
pub fn bosh_router(handler: ConnectionHandler) -> Router {
    let state = new_bosh_state(handler);
    let route = post(bosh).options(|| async { preflight("POST, OPTIONS") });
    Router::new()
        .route("/http-bind", route.clone())
        .route("/http-bind/", route)
        .layer(axum::extract::DefaultBodyLimit::max(256 * 1024))
        .with_state(state)
}
fn bosh_reply(xml: String) -> Response {
    (
        [
            (axum::http::header::CONTENT_TYPE, "text/xml; charset=utf-8"),
            (
                axum::http::HeaderName::from_static("access-control-allow-origin"),
                "*",
            ),
        ],
        xml,
    )
        .into_response()
}
fn terminate(condition: &str) -> Response {
    bosh_reply(format!(
        "<body xmlns='{BOSH}' type='terminate' condition='{condition}'/>"
    ))
}
async fn bosh(State(state): State<BoshState>, bytes: Bytes) -> Response {
    bosh_with_context(state, bytes, None).await
}
async fn bosh_with_context(
    state: BoshState,
    bytes: Bytes,
    context: Option<crate::http_proxy::RequestContext>,
) -> Response {
    // A disconnected HTTP client must not cancel stanza processing or leave a
    // pending RID without its cached response. A retry uses that same response.
    let Ok(permit) = state.requests.clone().try_acquire_owned() else {
        return terminate("resource-constraint");
    };
    match tokio::spawn(async move {
        let _permit = permit;
        bosh_request(state, bytes, context).await
    })
    .await
    {
        Ok(response) => response,
        Err(_) => terminate("remote-connection-failed"),
    }
}
async fn bosh_request(
    state: BoshState,
    bytes: Bytes,
    context: Option<crate::http_proxy::RequestContext>,
) -> Response {
    let Ok(text) = std::str::from_utf8(&bytes) else {
        return terminate("bad-request");
    };
    let Ok(doc) = roxmltree::Document::parse(text) else {
        return terminate("bad-request");
    };
    let body = doc.root_element();
    if body.tag_name().name() != "body" || body.tag_name().namespace() != Some(BOSH) {
        return terminate("bad-request");
    };
    let Some(rid) = body
        .attribute("rid")
        .and_then(|r| r.parse::<u64>().ok())
        .filter(|r| *r < 9_007_199_254_740_991)
    else {
        return terminate("bad-request");
    };
    let initial = body.attribute("sid").is_none();
    let sid = body
        .attribute("sid")
        .map(str::to_string)
        .unwrap_or_else(|| uuid::Uuid::new_v4().simple().to_string());
    let session = if initial {
        let Some(domain) = body.attribute("to").filter(|v| !v.is_empty()) else {
            return terminate("improper-addressing");
        };
        if body
            .attribute("hold")
            .and_then(|h| h.parse::<u64>().ok())
            .unwrap_or(1)
            > 1
        {
            return terminate("policy-violation");
        }
        let wait = body
            .attribute("wait")
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(30)
            .clamp(1, 60);
        let (internal, external) = tokio::io::duplex(256 * 1024);
        let engine_done = Arc::new(Notify::new());
        let done = engine_done.clone();
        let handler = if let Some(handler) = &state.context_handler {
            handler(internal, context)
        } else {
            (state.handler)(internal)
        };
        let engine = tokio::spawn(async move {
            handler.await;
            done.notify_one();
        });
        let (mut read, write) = tokio::io::split(external);
        let (tx, rx) = mpsc::channel(128);
        let pump = tokio::spawn(async move {
            let mut framer = Framer::default();
            while let Ok(frame) = framer.next(&mut read).await {
                if frame.starts_with("<?xml") || frame.starts_with("<stream:stream ") {
                    continue;
                }
                if frame == "</stream:stream>" {
                    break;
                }
                let Ok(Some(xml)) = to_websocket(&frame) else {
                    break;
                };
                if tx.send(xml).await.is_err() {
                    break;
                }
            }
        });
        let session = Arc::new(BoshSession {
            domain: domain.into(),
            wait,
            sequence: Mutex::new(Sequence {
                next: rid,
                pending: HashSet::new(),
                responses: BTreeMap::new(),
            }),
            changed: Notify::new(),
            poll_wakeup: Notify::new(),
            writer: Mutex::new(write),
            output: Mutex::new(rx),
            last: std::sync::Mutex::new(Instant::now()),
            tasks: vec![engine, pump],
            engine_done,
        });
        let mut map = state.sessions.lock().await;
        if map.len() >= 4096 {
            return terminate("resource-constraint");
        }
        map.insert(sid.clone(), session.clone());
        session
    } else {
        let map = state.sessions.lock().await;
        let Some(session) = map.get(&sid) else {
            return terminate("item-not-found");
        };
        session.clone()
    };
    *session.last.lock().unwrap() = Instant::now();
    match tokio::time::timeout(
        Duration::from_secs(90),
        process_bosh_request(&state, body, rid, initial, &sid, &session),
    )
    .await
    {
        Ok(response) => response,
        Err(_) => {
            state.sessions.lock().await.remove(&sid);
            for task in &session.tasks {
                task.abort();
            }
            terminate("remote-connection-failed")
        }
    }
}
async fn process_bosh_request(
    state: &BoshState,
    body: roxmltree::Node<'_, '_>,
    rid: u64,
    initial: bool,
    sid: &str,
    session: &BoshSession,
) -> Response {
    // A cached RID must never replay its stanzas. Two outstanding requests are allowed.
    let reorder_deadline = Instant::now() + Duration::from_secs(5);
    let mut seq = loop {
        let seq = session.sequence.lock().await;
        if rid <= seq.next {
            break seq;
        }
        if rid > seq.next.saturating_add(1) || Instant::now() >= reorder_deadline {
            return terminate("item-not-found");
        }
        drop(seq);
        tokio::time::sleep(Duration::from_millis(10)).await;
    };
    if let Some(cached) = seq.responses.get(&rid) {
        return bosh_reply(cached.clone());
    }
    if seq.pending.contains(&rid) {
        drop(seq);
        let deadline = Instant::now() + Duration::from_secs(session.wait + 5);
        loop {
            let notified = session.changed.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();
            if let Some(cached) = session.sequence.lock().await.responses.get(&rid).cloned() {
                return bosh_reply(cached);
            }
            if tokio::time::timeout_at(tokio::time::Instant::from_std(deadline), notified)
                .await
                .is_err()
            {
                return terminate("remote-connection-failed");
            }
        }
    }
    if rid != seq.next || seq.pending.len() >= 2 {
        return terminate("item-not-found");
    };
    seq.pending.insert(rid);
    if seq.pending.len() > 1 {
        session.poll_wakeup.notify_one();
    }
    seq.next += 1;
    let mut writer = session.writer.lock().await;
    drop(seq);
    let mut input = String::new();
    if initial || body.attribute(("urn:xmpp:xbosh", "restart")) == Some("true") {
        input.push_str(&format!("<stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' to='{}' version='1.0'>",escape(&session.domain)));
    }
    for node in body.children().filter(|n| n.is_element()) {
        input.push_str(&crate::protocol::xml(node));
    }
    let closing = body.attribute("type") == Some("terminate");
    if closing {
        input.push_str("</stream:stream>");
    }
    if writer.write_all(input.as_bytes()).await.is_err() {
        state.sessions.lock().await.remove(sid);
        return terminate("remote-connection-failed");
    }
    drop(writer);
    let mut payload = String::new();
    if closing {
        // Let the engine consume the stream close and run graceful cleanup before
        // removing the session: dropping it aborts the engine task. In particular,
        // a reconnect must not race with the old resource remaining registered.
        let mut rx = session.output.lock().await;
        let _ = tokio::time::timeout(Duration::from_secs(session.wait), async {
            while let Some(frame) = rx.recv().await {
                payload.push_str(&frame);
            }
            session.engine_done.notified().await;
        })
        .await;
    } else {
        let mut rx = session.output.lock().await;
        match tokio::select! { v=tokio::time::timeout(Duration::from_secs(session.wait), rx.recv())=>v, _=session.poll_wakeup.notified()=>Ok(Some(String::new())) }
        {
            Ok(Some(first)) => {
                payload.push_str(&first);
                while let Ok(next) = rx.try_recv() {
                    payload.push_str(&next);
                }
            }
            Ok(None) => {
                state.sessions.lock().await.remove(sid);
                return terminate("remote-stream-error");
            }
            Err(_) => {}
        }
    }
    let extra = if initial {
        format!(
            " sid='{sid}' wait='{}' hold='1' requests='2' inactivity='120' polling='1' ver='1.11' from='{}' xmlns:xmpp='urn:xmpp:xbosh' xmpp:version='1.0'",
            session.wait,
            escape(&session.domain)
        )
    } else if closing {
        " type='terminate'".into()
    } else {
        String::new()
    };
    let result = format!("<body xmlns='{BOSH}'{extra}>{payload}</body>");
    let mut seq = session.sequence.lock().await;
    seq.pending.remove(&rid);
    seq.responses.insert(rid, result.clone());
    while seq.responses.len() > 4 {
        seq.responses.pop_first();
    }
    drop(seq);
    session.changed.notify_waiters();
    if closing {
        state.sessions.lock().await.remove(sid);
    }
    bosh_reply(result)
}
/// Apply the exact configured origin allowlist; requests without Origin remain usable by native clients.
pub fn with_origins(router: Router, origins: Vec<String>) -> Router {
    let origins = Arc::new(origins);
    router.layer(axum::middleware::from_fn(
        move |req: axum::extract::Request, next: axum::middleware::Next| {
            let origins = origins.clone();
            async move {
                let origin = req.headers().get(axum::http::header::ORIGIN).cloned();
                if let Some(origin) = &origin
                    && !origin
                        .to_str()
                        .ok()
                        .is_some_and(|o| origins.iter().any(|a| a == o || a == "*"))
                {
                    return StatusCode::FORBIDDEN.into_response();
                }
                let mut response = next.run(req).await;
                response
                    .headers_mut()
                    .remove(axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN);
                if let Some(origin) = origin {
                    response
                        .headers_mut()
                        .insert(axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN, origin);
                    response.headers_mut().append(
                        axum::http::header::VARY,
                        axum::http::HeaderValue::from_static("Origin"),
                    );
                }
                response
            }
        },
    ))
}