reach_websocket/server/

remote.rs

// SPDX-FileCopyrightText: 2025 Michael Goldenberg <m@mgoldenberg.net>
// SPDX-FileCopyrightText: 2025 eaon <eaon@posteo.net>
// SPDX-License-Identifier: EUPL-1.2

use std::str;

use axum::{
    body::Body,
    extract::FromRequestParts,
    extract::ws::rejection::{
        ConnectionNotUpgradable, InvalidConnectionHeader, InvalidProtocolPseudoheader,
        InvalidUpgradeHeader, InvalidWebSocketVersionHeader, MethodNotConnect, MethodNotGet,
        WebSocketKeyHeaderMissing, WebSocketUpgradeRejection,
    },
    http::{HeaderValue, Method, Version, request::Parts},
    response::Response,
};
use hyper::{
    StatusCode, header,
    upgrade::{OnUpgrade, Upgraded},
};
use hyper_util::rt::TokioIo;
use rand::seq::SliceRandom;

use reach_aliases::*;

use super::*;

impl<Incoming, Outgoing> IntoWebSocketContext<RemoteServerContext<Incoming, Outgoing>>
    for RemoteServerOptions<Incoming, Outgoing>
where
    Incoming: CommunicableType + Send,
    Outgoing: CommunicableType + From<GenericWebSocketError> + Send,
{
    type Channel = RawItemSender<Incoming, Outgoing>;

    fn into_web_socket_context(
        self,
        channel: Self::Channel,
    ) -> RemoteServerContext<Incoming, Outgoing> {
        ServerContext {
            web_socket: self.web_socket,
            incoming_sink: channel,
            outgoing_stream: Default::default(),
            extensions: self.extensions,
            hold: None,
        }
    }
}

pub struct WebSocketUpgrade<E, D, C>
where
    D: RequestDelegator<C>,
    C: CommunicableTypes,
    C::Resp: From<GenericWebSocketError>,
{
    server_context_extensions: E,
    global_context: Arc<D::GlobalContext>,
    on_upgrade: OnUpgrade,
    sec_websocket_signed_key: Option<HeaderValue>,
}

macro_rules! if_then_err {
    ($boolean:expr, $error:tt) => {
        if $boolean {
            Err($error::default())?;
        }
    };
}

macro_rules! header_ne {
    ($headers:expr, $key:expr, $value:expr) => {
        !$headers
            .get($key)
            .is_some_and(|hdr| hdr.as_bytes().eq_ignore_ascii_case($value.as_bytes()))
    };
}

#[inline]
fn validate_http_11(parts: &Parts) -> Result<(), WebSocketUpgradeRejection> {
    if_then_err!(parts.method != Method::GET, MethodNotGet);

    if_then_err!(
        !parts.headers.get(header::CONNECTION).is_some_and(|h| {
            str::from_utf8(h.as_bytes()).is_ok_and(|b| b.to_ascii_lowercase().contains("upgrade"))
        }),
        InvalidConnectionHeader
    );

    if_then_err!(
        header_ne!(parts.headers, header::UPGRADE, "websocket"),
        InvalidUpgradeHeader
    );

    Ok(())
}

#[inline]
fn validate_http_2(parts: &Parts) -> Result<(), WebSocketUpgradeRejection> {
    if_then_err!(parts.method != Method::CONNECT, MethodNotConnect);

    if_then_err!(
        parts
            .extensions
            .get::<hyper::ext::Protocol>()
            .is_none_or(|p| p.as_str() != "websocket"),
        InvalidProtocolPseudoheader
    );

    Ok(())
}

#[inline]
fn sign(key: &HeaderValue) -> Result<HeaderValue, WebSocketUpgradeRejection> {
    use sha1::{Digest, Sha1};

    let mut hasher = Sha1::new();
    hasher.update(key);
    hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    use base64::prelude::*;

    let x = BASE64_STANDARD.encode(hasher.finalize());

    HeaderValue::from_str(&x).map_err(|_| ConnectionNotUpgradable::default().into())
}

impl<E, D, C> FromRequestParts<(E, Arc<D::GlobalContext>)> for WebSocketUpgrade<E, D, C>
where
    E: Clone + Send + Sync,
    D: RequestDelegator<C>,
    D::GlobalContext: Send + Sync,
    C: CommunicableTypes,
    C::Resp: From<GenericWebSocketError>,
{
    type Rejection = WebSocketUpgradeRejection;

    async fn from_request_parts(
        parts: &mut Parts,
        (server_context_extensions, global_context): &(E, Arc<D::GlobalContext>),
    ) -> Result<Self, Self::Rejection> {
        let headers = &parts.headers;

        let sec_websocket_signed_key = match parts.version {
            Version::HTTP_2 => validate_http_2(parts).map(|_| None)?,
            Version::HTTP_11 => {
                validate_http_11(parts)?;

                let key = headers
                    .get(header::SEC_WEBSOCKET_KEY)
                    .ok_or(WebSocketKeyHeaderMissing::default())?;

                Some(sign(key)?)
            }
            _ => Err(ConnectionNotUpgradable::default())?,
        };

        if_then_err!(
            header_ne!(headers, header::SEC_WEBSOCKET_VERSION, "13"),
            InvalidWebSocketVersionHeader
        );

        Ok(Self {
            global_context: global_context.clone(),
            sec_websocket_signed_key,
            server_context_extensions: server_context_extensions.clone(),
            on_upgrade: parts
                .extensions
                .remove()
                .ok_or(ConnectionNotUpgradable::default())?,
        })
    }
}

#[derive(Debug, Clone)]
pub struct RemoteServerContextExtensions {
    pub permissible_communication_lengths: Vec<usize>,
    pub csprng: ReachRng,
}

impl RemoteServerContextExtensions {
    pub fn new(permissible_communication_lengths: Vec<usize>, csprng: ReachRng) -> Self {
        Self {
            permissible_communication_lengths,
            csprng,
        }
    }
}

pub type RemoteServerOptions<Incoming, Outgoing> =
    ServerOptions<TokioIo<Upgraded>, Incoming, Outgoing, RemoteServerContextExtensions>;

pub type RemoteServerContext<Incoming, Outgoing> =
    ServerContext<TokioIo<Upgraded>, Incoming, Outgoing, (), RemoteServerContextExtensions>;

impl<Incoming, Outgoing> WebSocketContext for RemoteServerContext<Incoming, Outgoing>
where
    Incoming: CommunicableType + Send,
    Outgoing: CommunicableType + From<GenericWebSocketError> + Send,
{
    type Options = RemoteServerOptions<Incoming, Outgoing>;
    type Incoming = Incoming;
    type Outgoing = Outgoing;
    type Responder = ();

    async fn handle_incoming(
        &mut self,
        incoming: Communication<Self::Incoming>,
    ) -> Result<(), WebSocketError> {
        let (item, receiver) = WebSocketItem::channel(incoming);
        self.outgoing_stream.push(receiver);
        if let Err(e) = self.incoming_sink.send(item).await {
            if e.is_disconnected() {
                return Err(e.into());
            } else {
                // TODO: log error here
            }
        }

        Ok(())
    }

    async fn handle_outgoing(
        &mut self,
        outgoing: WebSocketItem<Communication<Self::Outgoing>, Self::Responder>,
    ) -> Result<(), WebSocketError> {
        // TODO: take maximum deviation from the base into account.
        // We don't want to blow up tiny messages to humongous sizes
        let length = loop {
            match self
                .extensions
                .permissible_communication_lengths
                .choose(&mut self.extensions.csprng)
            {
                Some(length) => {
                    if length > &outgoing.item.inner.len() {
                        break *length;
                    }
                }
                None => break 0,
            }
        };
        self.web_socket.send(outgoing.item.pad_to(length)).await?;

        Ok(())
    }
}

const UPGRADE_HEADER: HeaderValue = HeaderValue::from_static("upgrade");
const WEBSOCKET_HEADER: HeaderValue = HeaderValue::from_static("websocket");

pub fn on_upgrade<const B: usize, E, D, C, W>(
    upgrade: WebSocketUpgrade<E, D, C>,
    session_context_init: impl FnOnce() -> Arc<RwLock<D::SessionContext>> + Clone + Send + 'static,
) -> Response
where
    D: RequestDelegator<C> + 'static,
    D::GlobalContext: Send + Sync + 'static,
    D::SessionContext: Send + Sync + 'static,
    C: CommunicableTypes + 'static,
    C::Req: Send + Sync,
    C::Resp: From<GenericWebSocketError> + Send + Sync,
    E: Clone + Send + 'static,
    W: WebSocketContext<Incoming = C::Req, Outgoing = C::Resp> + Send + Unpin + 'static,
    W::Responder: Send + Sync,
    W::Options: From<(WebSocketChannel<TokioIo<Upgraded>, C::Req, C::Resp>, E)>,
    W::Options: IntoWebSocketContext<W, Channel = RawItemSender<W::Incoming, W::Outgoing>>,
{
    tokio::spawn(async move {
        if let Ok(upgraded) = upgrade.on_upgrade.await {
            let upgraded = TokioIo::new(upgraded);

            let web_socket_stream =
                WebSocketStream::from_raw_socket(upgraded, Role::Server, None).await;

            let _ = handle_connection::<B, _, E, D, C, W>(
                web_socket_stream,
                upgrade.server_context_extensions.clone(),
                upgrade.global_context,
                session_context_init,
            )
            .await;
        }
    });

    let mut response = Response::new(Body::empty());

    if let Some(signed_key) = upgrade.sec_websocket_signed_key {
        let (mut parts, body) = response.into_parts();
        parts.status = StatusCode::SWITCHING_PROTOCOLS;
        parts.headers.extend([
            (header::CONNECTION, UPGRADE_HEADER),
            (header::UPGRADE, WEBSOCKET_HEADER),
            (header::SEC_WEBSOCKET_ACCEPT, signed_key),
        ]);
        response = Response::from_parts(parts, body);
    }

    response
}