predawn 0.9.0

Predawn is a Rust web framework like Spring Boot.
Documentation
use std::{collections::BTreeMap, convert::Infallible, future::Future};

use headers::{Connection, HeaderMapExt, SecWebsocketAccept, Upgrade};
use http::{header::SEC_WEBSOCKET_PROTOCOL, StatusCode};
use hyper_util::rt::TokioIo;
use predawn_core::{
    api_response::ApiResponse,
    body::ResponseBody,
    into_response::IntoResponse,
    openapi::{self, Schema},
    response::{MultiResponse, Response, SingleResponse},
};
use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream};

use super::{OnFailedUpgrade, WebSocket, WebSocketRequest};

#[derive(Debug)]
pub struct WebSocketResponse(Response);

impl WebSocketResponse {
    pub(crate) fn new<F, C, Fut>(request: WebSocketRequest<F>, callback: C) -> WebSocketResponse
    where
        F: OnFailedUpgrade,
        C: FnOnce(WebSocket) -> Fut + Send + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let WebSocketRequest {
            config,
            protocol,
            sec_websocket_key,
            on_upgrade,
            on_failed_upgrade,
            sec_websocket_protocol: _,
        } = request;

        {
            let protocol = protocol.clone();

            tokio::spawn(async move {
                let upgraded = match on_upgrade.await {
                    Ok(upgraded) => upgraded,
                    Err(err) => {
                        on_failed_upgrade.call(err);
                        return;
                    }
                };
                let upgraded = TokioIo::new(upgraded);

                let socket =
                    WebSocketStream::from_raw_socket(upgraded, Role::Server, Some(config)).await;
                let socket = WebSocket {
                    inner: socket,
                    protocol,
                };

                callback(socket).await;
            });
        }

        let mut response = http::Response::builder()
            .status(StatusCode::SWITCHING_PROTOCOLS)
            .body(ResponseBody::empty())
            .unwrap();

        let headers = response.headers_mut();

        headers.typed_insert(Connection::upgrade());
        headers.typed_insert(Upgrade::websocket());
        headers.typed_insert(SecWebsocketAccept::from(sec_websocket_key));

        if let Some(protocol) = protocol {
            headers.insert(SEC_WEBSOCKET_PROTOCOL, protocol);
        }

        WebSocketResponse(response)
    }
}

impl IntoResponse for WebSocketResponse {
    type Error = Infallible;

    fn into_response(self) -> Result<Response, Self::Error> {
        Ok(self.0)
    }
}

impl SingleResponse for WebSocketResponse {
    const STATUS_CODE: u16 = 101;

    fn response(_: &mut BTreeMap<String, Schema>, _: &mut Vec<String>) -> openapi::Response {
        openapi::Response {
            description: "A WebSocket response".to_string(),
            ..Default::default()
        }
    }
}

impl ApiResponse for WebSocketResponse {
    fn responses(
        schemas: &mut BTreeMap<String, Schema>,
        schemas_in_progress: &mut Vec<String>,
    ) -> Option<BTreeMap<StatusCode, openapi::Response>> {
        Some(<Self as MultiResponse>::responses(
            schemas,
            schemas_in_progress,
        ))
    }
}