Skip to content

Axum integration #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Axum integration #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
75b3953
initial implemenation of hyper/axum integration
clarkohw Sep 19, 2024
6bb97f0
quick PR feedback
clarkohw Sep 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
[workspace]
resolver = "2"
members = [
"ratchet_rs",
"ratchet_axum",
"ratchet_core",
"ratchet_deflate",
"ratchet_ext",
"ratchet_fixture",
"ratchet_rs",
"ratchet_rs/autobahn/client",
"ratchet_rs/autobahn/server",
"ratchet_rs/autobahn/split_client",
Expand Down Expand Up @@ -43,3 +44,10 @@ flate2 = { version = "1.0", default-features = false }
anyhow = "1.0"
serde_json = "1.0"
tracing-subscriber = "0.3.18"
hyper = "1.4.1"
axum = "0.7.5"
axum-core = "0.4.3"
hyper-util = "0.1.0"
pin-project = "1.1.5"
async-trait = "0.1.79"
sha1 = "0.10.4"
30 changes: 30 additions & 0 deletions ratchet_axum/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
[package]
name = "ratchet_axum"
description = "Axum Integration for Ratchet"
readme = "README.md"
repository = "https://github.com/swimos/ratchet/"
version = "0.1.0"
edition.workspace = true
authors.workspace = true
license.workspace = true
categories.workspace = true


[dependencies]
ratchet_rs = { version = "1.0.3", path = "../ratchet_rs" }
hyper = { workspace = true }
axum-core = { workspace = true }
hyper-util = { workspace = true , features = ["tokio"]}
pin-project = { workspace = true }
async-trait = { workspace = true }
base64 = { workspace = true }
http = { workspace = true }
sha1 = { workspace = true }

[dev-dependencies]
axum = { workspace = true }
tokio = { workspace = true, features = ["full"] }
bytes = { workspace = true }

[[example]]
name = "axum"
40 changes: 40 additions & 0 deletions ratchet_axum/examples/axum.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
use axum::{response::IntoResponse, routing::get, Router};
use bytes::BytesMut;
use ratchet_axum::{UpgradeFut, WebSocketUpgrade};
use ratchet_rs::{Message, NegotiatedExtension, NoExt, PayloadType, Role, WebSocketConfig};

#[tokio::main]
async fn main() {
let app = Router::new().route("/", get(ws_handler));

let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app).await.unwrap();
}

async fn handle_client(fut: UpgradeFut) {
let io = fut.await.unwrap();
let mut websocket = ratchet_rs::WebSocket::from_upgraded(
WebSocketConfig::default(),
io,
NegotiatedExtension::from(NoExt),
BytesMut::new(),
Role::Server,
);
let mut buf = BytesMut::new();

loop {
match websocket.read(&mut buf).await.unwrap() {
Message::Text => {
websocket.write(&mut buf, PayloadType::Text).await.unwrap();
buf.clear();
}
_ => break,
}
}
}

async fn ws_handler(incoming_upgrade: WebSocketUpgrade) -> impl IntoResponse {
let (response, fut) = incoming_upgrade.upgrade().unwrap();
tokio::task::spawn(async move { handle_client(fut).await });
response
}
160 changes: 160 additions & 0 deletions ratchet_axum/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
//todo missing docs

#![deny(
// missing_docs,
missing_copy_implementations,
missing_debug_implementations,
trivial_numeric_casts,
unstable_features,
unused_must_use,
unused_mut,
unused_imports,
unused_import_braces
)]

use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use axum_core::body::Body;
use base64;
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use http::HeaderMap;
use hyper::Response;
use hyper_util::rt::TokioIo;
use pin_project::pin_project;
use sha1::Digest;
use sha1::Sha1;

const HEADER_CONNECTION: &str = "upgrade";
const HEADER_UPGRADE: &str = "websocket";
const WEBSOCKET_VERSION: &[u8] = b"13";

type Error = hyper::Error;

#[derive(Debug)]
pub struct WebSocketUpgrade {
key: String,
headers: HeaderMap,
on_upgrade: hyper::upgrade::OnUpgrade,
pub permessage_deflate: bool,
}

impl WebSocketUpgrade {
pub fn upgrade(self) -> Result<(Response<Body>, UpgradeFut), Error> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be more ergonomic if this was inverted a bit. It would also be safer as it ensures that any contracts made during the upgrade are upheld as we create the WebSocket instance. Otherwise, we could negotiate an extension and then a user could create a WebSocket without it.

Something like this:

pub fn upgrade<E,F, Fut>(self, f:F) -> Response<Body>
    where
        F: FnOnce(UpgradedServer<TokioIo<hyper::upgrade::Upgraded>, E>) -> Fut,
        Fut: Future<Output=()>,
        E: Extension,
    {
          // await the upgrade future and spawn the user's handler after creating the WebSocket. 
    }

Then you could use it as follows:

async fn ws_handler<E>(incoming_upgrade: IncomingUpgrade, state: State<E>) -> impl IntoResponse {
    incoming_upgrade.upgrade(|mut upgraded| async {
        let UpgradedServer {
            request,
            websocket,
            subprotocol,
        } = upgraded;

        let mut buf = BytesMut::new();

        loop {
            match websocket.read(&mut buf).await.unwrap() {
                Message::Text => {
                    websocket.write(&mut buf, PayloadType::Text).await.unwrap();
                    buf.clear();
                }
                _ => break,
            }
        }
    })
}

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@SirCipher I am trying to implement the upgrade like this but having trouble because ExtensionProvider is not Send. I tried it with the below setup, and was able to get it to compile/run, but running into this error:

called `Result::unwrap()` on an `Err` value: Error { inner: Inner { kind: IO, source: Some(Kind(UnexpectedEof)) } }

wondering if you have any advice here

     pub fn upgrade_2<E, F, Fut>(self, f: F, provider: E) -> Response<Body>
    where
        F: FnOnce(UpgradedServer<TokioIo<hyper::upgrade::Upgraded>, E::Extension>) -> Fut
            + Send
            + 'static,
        Fut: Future<Output = ()> + Send,
        E: ExtensionProvider + Send + 'static,
        <E as ExtensionProvider>::Extension: Send,
    {
        tokio::spawn(async move {
            let upgraded = match self.on_upgrade.await {
                Ok(upgraded) => upgraded,
                Err(err) => {
                    return;
                }
            };
            let upgraded = TokioIo::new(upgraded);
            let upgrade_server = accept_with(
                upgraded,
                WebSocketConfig::default(),
                provider,
                SubprotocolRegistry::default(),
            )
            .await
            .unwrap()
            .upgrade_with(self.headers)
            .await
            .unwrap();

            f(upgrade_server).await;
        });
        let builder = Response::builder()
            .status(hyper::StatusCode::SWITCHING_PROTOCOLS)
            .header(hyper::header::CONNECTION, HEADER_CONNECTION)
            .header(hyper::header::UPGRADE, HEADER_UPGRADE)
            .header(hyper::header::SEC_WEBSOCKET_ACCEPT, self.key);

        let response = builder
            .body(Body::default())
            .expect("bug: failed to build response");
        response
    }

let mut builder = Response::builder()
.status(hyper::StatusCode::SWITCHING_PROTOCOLS)
.header(hyper::header::CONNECTION, HEADER_CONNECTION)
.header(hyper::header::UPGRADE, HEADER_UPGRADE)
.header(hyper::header::SEC_WEBSOCKET_ACCEPT, self.key);

if self.permessage_deflate {
builder = builder.header(
hyper::header::SEC_WEBSOCKET_EXTENSIONS,
"permessage-deflate",
);
}

let response = builder
.body(Body::default())
.expect("bug: failed to build response");

let stream = UpgradeFut {
inner: self.on_upgrade,
headers: self.headers,
};

Ok((response, stream))
}

// pub fn upgrade_2<E, F, Fut>(self, f: F) -> Response<Body>
// where
// F: FnOnce(UpgradedServer<TokioIo<hyper::upgrade::Upgraded>, E>) -> Fut,
// Fut: Future<Output = ()>,
// E: Extension,
// {
//
//
// }
}

#[async_trait::async_trait]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seeing as this is the only place that async_trait is used it's not really worth adding the dependency. You can just expand from_request_parts to return a BoxFuture

impl<S> axum_core::extract::FromRequestParts<S> for WebSocketUpgrade
where
S: Sync,
{
type Rejection = hyper::StatusCode;

async fn from_request_parts(
parts: &mut http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
let key = parts
.headers
.get(http::header::SEC_WEBSOCKET_KEY)
.ok_or(hyper::StatusCode::BAD_REQUEST)?;

if parts
.headers
.get(http::header::SEC_WEBSOCKET_VERSION)
.map(|v| v.as_bytes())
!= Some(WEBSOCKET_VERSION)
{
return Err(hyper::StatusCode::BAD_REQUEST);
}

let permessage_deflate = parts
.headers
.get(http::header::SEC_WEBSOCKET_EXTENSIONS)
.map(|val| {
val.to_str()
.unwrap_or_default()
.to_lowercase()
.contains("permessage-deflate")
})
.unwrap_or(false);

let on_upgrade = parts
.extensions
.remove::<hyper::upgrade::OnUpgrade>()
.ok_or(hyper::StatusCode::BAD_REQUEST)?;

Ok(Self {
on_upgrade,
key: sec_websocket_protocol(key.as_bytes()),
headers: parts.headers.clone(),
permessage_deflate,
})
}
}

/// A future that resolves to a websocket stream when the associated HTTP upgrade completes.
#[pin_project]
#[derive(Debug)]
pub struct UpgradeFut {
#[pin]
inner: hyper::upgrade::OnUpgrade,
pub headers: HeaderMap,
}

impl std::future::Future for UpgradeFut {
type Output = Result<TokioIo<hyper::upgrade::Upgraded>, Error>;

fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let this = self.project();
let upgraded = match this.inner.poll(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(x) => x,
};
Poll::Ready(upgraded.map(|u| TokioIo::new(u)).map_err(|e| e.into()))
}
}

fn sec_websocket_protocol(key: &[u8]) -> String {
let mut sha1 = Sha1::new();
sha1.update(key);
sha1.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); // magic string
let result = sha1.finalize();
STANDARD.encode(&result[..])
}
4 changes: 3 additions & 1 deletion ratchet_rs/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ split = ["ratchet_core/split"]
fixture = ["ratchet_core/fixture"]

[dependencies]
axum = { workspace = true, optional = true }
hyper = { workspace = true, features = ["http1", "server", "client"], optional = true }
ratchet_core = { version = "1.0.3", path = "../ratchet_core" }
ratchet_ext = { version = "1.0.3", path = "../ratchet_ext" }
ratchet_deflate = { version = "1.0.3", path = "../ratchet_deflate", optional = true }
Expand Down Expand Up @@ -49,4 +51,4 @@ name = "client"
required-features = ["split"]

[[example]]
name = "server"
name = "server"