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 1 commit
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
8 changes: 8 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ members = [
"ratchet_deflate",
"ratchet_ext",
"ratchet_fixture",
"ratchet_axum",
"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"
eyre = "0.6.12"
hyper-util = "0.1.0"
pin-project = "1.1.5"
async-trait = "0.1.79"
24 changes: 24 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,24 @@
[package]
name = "ratchet_axum"
description = "Axum Integration for Ratchet"
readme = "README.md"
repository = "https://github.com/swimos/ratchet/"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
categories.workspace = true


[dependencies]
ratchet_core = { version = "1.0.3", path = "../ratchet_core" }
ratchet_ext = { version = "1.0.3", path = "../ratchet_ext" }
hyper = { workspace = true }
axum-core = { workspace = true }
eyre = { workspace = true}
hyper-util = { workspace = true , features = ["tokio"]}
pin-project = { workspace = true }
async-trait = { workspace = true }
base64 = "0.22.1"
http = "1.1.0"
sha1 = "0.10.1"
162 changes: 162 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,162 @@
// Port of hyper_tunstenite for fastwebsockets.
Copy link
Contributor

Choose a reason for hiding this comment

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

I can understand the rationale behind this but I don't believe that this is required as it's a fairly simple integration and there isn't really another way of integrating it into Axum.

Copy link
Contributor

Choose a reason for hiding this comment

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

Some module-level documentation would be good for this. Coupled with an example of it

// https://github.com/de-vri-es/hyper-tungstenite-rs
//
// Copyright 2021, Maarten de Vries [email protected]
// BSD 2-Clause "Simplified" License
//
// Copyright 2023 Divy Srivastava <[email protected]>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//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 eyre::Report;
use http::HeaderMap;
use hyper::Response;
use hyper_util::rt::TokioIo;
use pin_project::pin_project;
use sha1::Digest;
use sha1::Sha1;

type Error = Report;

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

impl IncomingUpgrade {
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, "upgrade")
.header(hyper::header::UPGRADE, "websocket")
.header("Sec-WebSocket-Accept", self.key);
Copy link
Contributor

Choose a reason for hiding this comment

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

A user needs to be able to specify sec-websocket-protocol headers and the server needs to take the intersection of the two sets.


if self.permessage_deflate {
builder = builder.header("Sec-WebSocket-Extensions", "permessage-deflate");
Copy link
Contributor

Choose a reason for hiding this comment

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

This isn't quite correct. If the client has sent a permessage-deflate header then it must be negotiated by the extension otherwise the connection will be unreliable. If the server accepts the PMCE request, then the client may send a compressed frame that the server won't be able to read. While it is correct that the server can elect to ignore the configuration parameters, Ratchet does support Deflate PMCE configuration and it's worth integrating.

What it would be worth doing, is providing an https://docs.rs/ratchet_ext/latest/ratchet_ext/trait.ExtensionProvider.html which a user may specify to the Axum Router or the handler may capture it. Then you can use this provider to perform the extension negotiation using https://docs.rs/ratchet_ext/latest/ratchet_ext/trait.ExtensionProvider.html#tymethod.negotiate_server.

Also, so it's tidier, you can use http::http::header::SEC_WEBSOCKET_EXTENSIONS here.

}

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))
}
}

#[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 IncomingUpgrade
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("Sec-WebSocket-Key")
.ok_or(hyper::StatusCode::BAD_REQUEST)?;
if parts
.headers
.get("Sec-WebSocket-Version")
.map(|v| v.as_bytes())
!= Some(b"13".as_slice())
{
return Err(hyper::StatusCode::BAD_REQUEST);
}

let permessage_deflate = parts
.headers
.get("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[..])
}
10 changes: 9 additions & 1 deletion ratchet_rs/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,15 @@ default = []
deflate = ["ratchet_deflate"]
split = ["ratchet_core/split"]
fixture = ["ratchet_core/fixture"]
with_axum = ["ratchet_axum", "axum", "hyper"]

[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 }
ratchet_axum = { version = "1.0.3", path = "../ratchet_axum", optional = true }
tracing-subscriber = { workspace = true, features = ["env-filter"] }
log = { workspace = true }

Expand Down Expand Up @@ -49,4 +53,8 @@ name = "client"
required-features = ["split"]

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

[[example]]
name = "axum"
required-features = ["with_axum"]
40 changes: 40 additions & 0 deletions ratchet_rs/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::{IncomingUpgrade, UpgradeFut};
use ratchet_core::{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: IncomingUpgrade) -> impl IntoResponse {
let (response, fut) = incoming_upgrade.upgrade().unwrap();
tokio::task::spawn(async move { handle_client(fut).await });
response
}
Loading