/
main.rs
116 lines (107 loc) · 3.31 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
//! Example websocket server.
//!
//! Run with
//!
//! ```not_rust
//! cd examples && cargo run -p example-websockets
//! ```
use axum::{
extract::{
ws::{Message, WebSocket, WebSocketUpgrade},
TypedHeader,
},
http::StatusCode,
response::IntoResponse,
routing::{get, get_service},
Router,
};
use std::{net::SocketAddr, path::PathBuf};
use tower_http::{
services::ServeDir,
trace::{DefaultMakeSpan, TraceLayer},
};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
#[tokio::main]
async fn main() {
tracing_subscriber::registry()
.with(tracing_subscriber::EnvFilter::new(
std::env::var("RUST_LOG")
.unwrap_or_else(|_| "example_websockets=debug,tower_http=debug".into()),
))
.with(tracing_subscriber::fmt::layer())
.init();
let assets_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets");
// build our application with some routes
let app = Router::new()
.fallback(
get_service(ServeDir::new(assets_dir).append_index_html_on_directories(true))
.handle_error(|error: std::io::Error| async move {
(
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 {
match msg {
Message::Text(t) => {
println!("client sent str: {:?}", t);
}
Message::Binary(_) => {
println!("client sent binary data");
}
Message::Ping(_) => {
println!("socket ping");
}
Message::Pong(_) => {
println!("socket pong");
}
Message::Close(_) => {
println!("client disconnected");
return;
}
}
} 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;
}
}