mirror of
https://github.com/tokio-rs/axum.git
synced 2024-12-29 15:49:16 +01:00
8013165908
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(...)`.
95 lines
2.5 KiB
Rust
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;
|
|
}
|
|
}
|