websocket添加account

This commit is contained in:
李运家 2025-03-09 20:21:01 +08:00
parent 06774af7fe
commit 7a6ac2f09d
2 changed files with 7 additions and 2 deletions

View File

@ -4,16 +4,18 @@ use axum::extract::ws::WebSocketUpgrade;
use axum::extract::ConnectInfo; use axum::extract::ConnectInfo;
use axum::response::IntoResponse; use axum::response::IntoResponse;
use axum_extra::{headers, TypedHeader}; use axum_extra::{headers, TypedHeader};
use library::context::Context;
use macros::ws; use macros::ws;
use crate::service::websocket_service; use crate::service::websocket_service;
#[ws("/ws")] #[ws("/ws")]
pub async fn websocket_handler( pub async fn websocket_handler(
context: Context,
ws: WebSocketUpgrade, ws: WebSocketUpgrade,
user_agent: Option<TypedHeader<headers::UserAgent>>, user_agent: Option<TypedHeader<headers::UserAgent>>,
ConnectInfo(addr): ConnectInfo<SocketAddr> ConnectInfo(addr): ConnectInfo<SocketAddr>
) -> impl IntoResponse { ) -> impl IntoResponse {
tracing::info!("`{:?}` at {:?} connected.", user_agent, addr); tracing::info!("`{:?}` at {:?} connected.", user_agent, addr);
ws.on_upgrade(move |socket| websocket_service::handle_socket(socket, addr)) ws.on_upgrade(move |socket| websocket_service::handle_socket(socket, addr, context))
} }

View File

@ -4,9 +4,12 @@ use std::ops::ControlFlow;
use axum::extract::ws::{Message, WebSocket}; use axum::extract::ws::{Message, WebSocket};
use futures::stream::{SplitSink, SplitStream}; use futures::stream::{SplitSink, SplitStream};
use futures::{SinkExt, StreamExt}; use futures::{SinkExt, StreamExt};
use library::context::Context;
/// Actual websocket statemachine (one will be spawned per connection) /// Actual websocket statemachine (one will be spawned per connection)
pub async fn handle_socket(socket: WebSocket, _who: SocketAddr) { pub async fn handle_socket(socket: WebSocket, _who: SocketAddr, context: Context) {
let account = context.get_account().unwrap();
tracing::info!("`{:?}` at {:?} connected, user is {:?}", account, _who, account.username);
let (sender, receiver) = socket.split(); let (sender, receiver) = socket.split();
tokio::spawn(write(sender)); tokio::spawn(write(sender));