axum/examples/websocket.rs
David Pedersen 8013165908
Move methods from ServiceExt to RoutingDsl (#160)
Previously, on `main`, this wouldn't compile:

```rust
let app = route("/", get(handler))
    .layer(
        ServiceBuilder::new()
            .timeout(Duration::from_secs(10))
            .into_inner(),
    )
    .handle_error(...)
    .route(...); // <-- doesn't work
```

That is because `handle_error` would be
`axum::service::ServiceExt::handle_error` which returns `HandleError<_,
_, _, HandleErrorFromService>` which does _not_ implement `RoutingDsl`.
So you couldn't call `route`. This was caused by
https://github.com/tokio-rs/axum/pull/120.

Basically `handle_error` when called on a `RoutingDsl`, the resulting
service should also implement `RoutingDsl`, but if called on another
random service it should _not_ implement `RoutingDsl`.

I don't think thats possible by having `handle_error` on `ServiceExt`
which is implemented for any service, since all axum routers are also
services by design.

This resolves the issue by removing `ServiceExt` and moving its methods
to `RoutingDsl`. Then we have more tight control over what has a
`handle_error` method.

`service::OnMethod` now also has a `handle_error` so you can still
handle errors from random services, by doing
`service::any(svc).handle_error(...)`.
2021-08-08 14:30:51 +02:00

95 lines
2.5 KiB
Rust

//! Example websocket server.
//!
//! Run with
//!
//! ```not_rust
//! cargo run --features=ws,headers --example websocket
//! ```
use axum::{
extract::{
ws::{Message, WebSocket, WebSocketUpgrade},
TypedHeader,
},
prelude::*,
response::IntoResponse,
routing::nest,
};
use http::StatusCode;
use std::net::SocketAddr;
use tower_http::{
services::ServeDir,
trace::{DefaultMakeSpan, TraceLayer},
};
#[tokio::main]
async fn main() {
// Set the RUST_LOG, if it hasn't been explicitly defined
if std::env::var("RUST_LOG").is_err() {
std::env::set_var("RUST_LOG", "websocket=debug,tower_http=debug")
}
tracing_subscriber::fmt::init();
// build our application with some routes
let app = nest(
"/",
axum::service::get(
ServeDir::new("examples/websocket").append_index_html_on_directories(true),
)
.handle_error(|error: std::io::Error| {
Ok::<_, std::convert::Infallible>((
StatusCode::INTERNAL_SERVER_ERROR,
format!("Unhandled internal error: {}", error),
))
}),
)
// routes are matched from bottom to top, so we have to put `nest` at the
// top since it matches all routes
.route("/ws", get(ws_handler))
// logging so we can see whats going on
.layer(
TraceLayer::new_for_http().make_span_with(DefaultMakeSpan::default().include_headers(true)),
);
// run it with hyper
let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
tracing::debug!("listening on {}", addr);
axum::Server::bind(&addr)
.serve(app.into_make_service())
.await
.unwrap();
}
async fn ws_handler(
ws: WebSocketUpgrade,
user_agent: Option<TypedHeader<headers::UserAgent>>,
) -> impl IntoResponse {
if let Some(TypedHeader(user_agent)) = user_agent {
println!("`{}` connected", user_agent.as_str());
}
ws.on_upgrade(handle_socket)
}
async fn handle_socket(mut socket: WebSocket) {
if let Some(msg) = socket.recv().await {
if let Ok(msg) = msg {
println!("Client says: {:?}", msg);
} else {
println!("client disconnected");
return;
}
}
loop {
if socket
.send(Message::Text(String::from("Hi!")))
.await
.is_err()
{
println!("client disconnected");
return;
}
tokio::time::sleep(std::time::Duration::from_secs(3)).await;
}
}