use crate::common::Frame;
use crate::server::{ConnectionHandler, ConnectionManagerTrait, ServerEventHandler};
use async_trait::async_trait;
use std::sync::Arc;
use tracing::info;
pub struct ServerMessageWrapper {
pub(crate) event_handler: Arc<dyn ServerEventHandler>,
connection_manager: Option<Arc<crate::server::connection::ConnectionManager>>,
device_manager: Option<Arc<crate::server::device::DeviceManager>>,
parser: crate::common::MessageParser,
}
impl ServerMessageWrapper {
pub fn new(
event_handler: Arc<dyn ServerEventHandler>,
connection_manager: Option<Arc<crate::server::connection::ConnectionManager>>,
device_manager: Option<Arc<crate::server::device::DeviceManager>>,
parser: crate::common::MessageParser,
) -> Self {
Self {
event_handler,
connection_manager,
device_manager,
parser,
}
}
async fn handle_and_send_response<F>(
&self,
handler_future: F,
message_id: String,
connection_id: &str,
log_context: &str,
) -> crate::common::error::Result<()>
where
F: std::future::Future<Output = crate::common::error::Result<Option<Frame>>>,
{
let response_frame = match handler_future.await {
Ok(Some(response)) => {
tracing::trace!(
"[ServerMessageWrapper] {}: 自定义响应: connection_id={}, message_id={}",
log_context,
connection_id,
message_id
);
response
}
Ok(None) => {
tracing::trace!(
"[ServerMessageWrapper] {}: 自动 ACK: connection_id={}, message_id={}",
log_context,
connection_id,
message_id
);
use crate::common::protocol::Reliability;
use crate::common::protocol::builder::{ack_message, frame_with_payload_command};
frame_with_payload_command(ack_message(message_id, None), Reliability::AtLeastOnce)
}
Err(e) => {
tracing::error!(
"[ServerMessageWrapper] {}: 处理失败,发送错误 ACK, connection_id={}, message_id={}, error={}",
log_context,
connection_id,
message_id,
e
);
use crate::common::protocol::Reliability;
use crate::common::protocol::builder::{ack_message, frame_with_payload_command};
use std::collections::HashMap;
let mut metadata = HashMap::new();
metadata.insert("error".to_string(), b"true".to_vec());
metadata.insert("error_message".to_string(), e.to_string().into_bytes());
frame_with_payload_command(
ack_message(message_id, Some(metadata)),
Reliability::AtLeastOnce,
)
}
};
if let Some(manager) = &self.connection_manager {
self.send_response_frame_async(response_frame, connection_id, log_context, manager)
.await;
}
Ok(())
}
async fn send_response_frame_async(
&self,
frame_to_send: Frame,
connection_id: &str,
log_context: &str,
manager: &Arc<crate::server::connection::ConnectionManager>,
) {
let manager_trait = Arc::clone(manager) as Arc<dyn ConnectionManagerTrait>;
let conn_id: Arc<str> = Arc::from(connection_id);
let message_id = frame_to_send.message_id.clone();
let log_ctx: Arc<str> = Arc::from(log_context);
tokio::spawn(async move {
match manager_trait
.send_frame_to(conn_id.as_ref(), &frame_to_send, None)
.await
{
Ok(()) => tracing::debug!(
"[ServerMessageWrapper] {}: 已发送, connection_id={}, message_id={}",
log_ctx,
conn_id,
message_id
),
Err(e) => {
tracing::warn!(
"[ServerMessageWrapper] {}: 发送失败, connection_id={}, message_id={}, error={}",
log_ctx,
conn_id,
message_id,
e
);
}
}
});
}
fn update_connection_active_async(&self, connection_id: &str) {
if let Some(manager) = &self.connection_manager {
let manager_trait = Arc::clone(manager) as Arc<dyn ConnectionManagerTrait>;
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
let _ = manager_trait.update_connection_active(&conn_id).await;
});
}
}
async fn handle_system_command(
&self,
frame: &Frame,
sys_type: i32,
connection_id: &str,
) -> crate::common::error::Result<()> {
use crate::common::protocol::flare::core::commands::system_command::Type as SysType;
match SysType::try_from(sys_type) {
Ok(SysType::Ping) => {
self.update_connection_active_async(connection_id);
match self.event_handler.handle_ping(frame, connection_id).await {
Ok(Some(custom_response)) => {
if let Some(manager) = &self.connection_manager {
self.send_response_frame_async(
custom_response,
connection_id,
"handle_ping",
manager,
)
.await;
}
}
_ => {
use crate::common::protocol::{
Reliability, frame_with_system_command, pong,
};
let pong_frame =
frame_with_system_command(pong(), Reliability::AtLeastOnce);
if let Some(manager) = &self.connection_manager {
self.send_response_frame_async(
pong_frame,
connection_id,
"handle_ping",
manager,
)
.await;
}
}
}
}
Ok(SysType::Pong) => {
let _ = self.event_handler.handle_pong(frame, connection_id).await;
self.update_connection_active_async(connection_id);
}
Ok(SysType::Event) => {
self.update_connection_active_async(connection_id);
let frame_clone = frame.clone();
let frame_message_id = frame.message_id.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
let wrapper = self.clone_for_async();
tokio::spawn(async move {
if let Err(e) = wrapper
.handle_and_send_response(
wrapper
.event_handler
.handle_system_event(&frame_clone, &conn_id),
frame_message_id,
&conn_id,
"handle_system_event",
)
.await
{
tracing::error!(
"[ServerMessageWrapper] handle_system_event: 处理失败, connection_id={}, error={}",
conn_id,
e
);
}
});
}
Ok(SysType::NegotiationReady) => {
self.update_connection_active_async(connection_id);
if let Some(manager) = &self.connection_manager {
let manager_clone = Arc::clone(manager);
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Err(e) = (*manager_clone).mark_negotiation_confirmed(&conn_id) {
tracing::error!(
"[ServerMessageWrapper] 标记协商确认失败: connection_id={}, error={}",
conn_id,
e
);
} else {
tracing::debug!(
"[ServerMessageWrapper] ✅ 协商已确认: connection_id={}",
conn_id
);
}
});
}
}
_ => {
tracing::debug!("[ServerMessageWrapper] 未处理的系统命令类型: {}", sys_type);
}
}
Ok(())
}
async fn handle_message_command(
&self,
_frame: &Frame,
command: &crate::common::protocol::PayloadCommand,
connection_id: &str,
) -> crate::common::error::Result<()> {
let message_id = command.message_id.clone();
use crate::common::protocol::flare::core::commands::payload_command::Type as PayloadType;
if let Ok(payload_type) = PayloadType::try_from(command.r#type) {
let handler_future = match payload_type {
PayloadType::Message => self.event_handler.handle_message(command, connection_id),
PayloadType::Event => self.event_handler.handle_event(command, connection_id),
PayloadType::Ack => self.event_handler.handle_ack(command, connection_id),
PayloadType::Data => self.event_handler.handle_data(command, connection_id),
PayloadType::Unspecified => {
tracing::warn!(
"[ServerMessageWrapper] handle_message_command: Unspecified payload type, connection_id={}",
connection_id
);
return Ok(());
}
};
self.handle_and_send_response(
handler_future,
message_id,
connection_id,
"handle_message_command",
)
.await?;
return Ok(());
}
tracing::error!(
"[ServerMessageWrapper] handle_message_command: 无法识别载荷类型, connection_id={}, message_id={}",
connection_id,
message_id
);
Err(crate::common::error::FlareError::general_error(
"Unknown message type",
))
}
async fn handle_notification_command(
&self,
frame: &Frame,
command: &crate::common::protocol::NotificationCommand,
connection_id: &str,
) -> crate::common::error::Result<()> {
self.handle_and_send_response(
self.event_handler
.handle_notification_command(command, connection_id),
frame.message_id.clone(),
connection_id,
"handle_notification_command",
)
.await
}
async fn handle_custom_command(
&self,
frame: &Frame,
command: &crate::common::protocol::flare::core::commands::CustomCommand,
connection_id: &str,
) -> crate::common::error::Result<()> {
self.update_connection_active_async(connection_id);
let cmd_name = command.name.clone();
let log_ctx = format!("handle_custom_command[{}]", cmd_name);
self.handle_and_send_response(
self.event_handler
.handle_custom_command(command, connection_id),
frame.message_id.clone(),
connection_id,
&log_ctx,
)
.await
}
fn clone_for_async(&self) -> Self {
Self {
event_handler: self.event_handler.clone(),
connection_manager: self.connection_manager.clone(),
device_manager: self.device_manager.clone(),
parser: self.parser.clone(),
}
}
}
impl Clone for ServerMessageWrapper {
fn clone(&self) -> Self {
self.clone_for_async()
}
}
#[async_trait]
impl ConnectionHandler for ServerMessageWrapper {
async fn handle_frame(
&self,
frame: &Frame,
connection_id: &str,
) -> crate::common::error::Result<Option<Frame>> {
if let Some(cmd) = &frame.command {
match &cmd.r#type {
Some(crate::common::protocol::flare::core::commands::command::Type::System(
sys_cmd,
)) => {
let sys_type = sys_cmd.r#type;
use crate::common::protocol::flare::core::commands::system_command::Type as SysType;
if sys_type != SysType::Connect as i32 {
let wrapper = self.clone_for_async();
let frame_clone = frame.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Err(e) = wrapper
.handle_system_command(&frame_clone, sys_type, &conn_id)
.await
{
tracing::error!("[ServerMessageWrapper] 处理系统命令失败: {}", e);
}
});
}
}
Some(crate::common::protocol::flare::core::commands::command::Type::Payload(
msg_cmd,
)) => {
let wrapper = self.clone_for_async();
let msg_cmd_clone = msg_cmd.clone();
let frame_clone = frame.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Err(e) = wrapper
.handle_message_command(&frame_clone, &msg_cmd_clone, &conn_id)
.await
{
tracing::error!("[ServerMessageWrapper] 处理消息命令失败: {}", e);
}
});
}
Some(
crate::common::protocol::flare::core::commands::command::Type::Notification(
notif_cmd,
),
) => {
let wrapper = self.clone_for_async();
let notif_cmd_clone = notif_cmd.clone();
let frame_clone = frame.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Err(e) = wrapper
.handle_notification_command(&frame_clone, ¬if_cmd_clone, &conn_id)
.await
{
tracing::error!("[ServerMessageWrapper] 处理通知命令失败: {}", e);
}
});
}
Some(crate::common::protocol::flare::core::commands::command::Type::Custom(
custom_cmd,
)) => {
let wrapper = self.clone_for_async();
let custom_cmd_clone = custom_cmd.clone();
let frame_clone = frame.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Err(e) = wrapper
.handle_custom_command(&frame_clone, &custom_cmd_clone, &conn_id)
.await
{
tracing::error!("[ServerMessageWrapper] 处理自定义命令失败: {}", e);
}
});
}
None => {
tracing::debug!("[ServerMessageWrapper] 未处理的命令类型");
}
}
}
Ok(None)
}
async fn on_connect(&self, connection_id: &str) -> crate::common::error::Result<()> {
info!("[ServerMessageWrapper] ✅ 新连接: {}", connection_id);
self.event_handler.on_connect(connection_id).await
}
async fn on_disconnect(&self, connection_id: &str) -> crate::common::error::Result<()> {
info!("[ServerMessageWrapper] ❌ 连接断开: {}", connection_id);
let _ = self.event_handler.on_disconnect(connection_id, None).await;
if let (Some(device_mgr), Some(manager)) = (&self.device_manager, &self.connection_manager)
{
let manager_trait = Arc::clone(manager) as Arc<dyn ConnectionManagerTrait>;
let device_mgr_clone = device_mgr.clone();
let conn_id: Arc<str> = Arc::from(connection_id);
tokio::spawn(async move {
if let Some((_, conn_info)) = manager_trait.get_connection(&conn_id).await
&& let Some(user_id) = conn_info.user_id
{
if let Err(e) = device_mgr_clone.remove_device(&user_id, &conn_id).await {
tracing::debug!(
"[ServerMessageWrapper] Failed to remove device from DeviceManager: {}",
e
);
} else {
tracing::info!(
"[ServerMessageWrapper] Successfully removed device from DeviceManager: user_id={}, connection_id={}",
user_id,
conn_id
);
}
}
});
}
Ok(())
}
}
impl ServerMessageWrapper {
pub(crate) async fn on_error(
&self,
connection_id: &str,
error: &str,
) -> crate::common::error::Result<()> {
tracing::error!(
"[ServerMessageWrapper] ❌ 连接错误: connection_id={}, error={}",
connection_id,
error
);
if let Err(e) = self.event_handler.on_error(connection_id, error).await {
tracing::error!(
"[ServerMessageWrapper] 事件处理器处理错误失败: connection_id={}, error={}",
connection_id,
e
);
}
Ok(())
}
}