引入shell模块

Cargo.toml

@@ -16,3 +16,5 @@ log = "0.4.25"
 fern = "0.7.1"
 colored = "3.0.0"
 chrono = "0.4.39"
+lazy_static = "1.5.0"
+once_cell = "1.20.2"

src/bin/main.rs

@@ -8,11 +8,15 @@ use async_trait::async_trait;
 use chrono::Local;
 use colored::{Color, Colorize};
 use fern::Dispatch;
+use lazy_static::lazy_static;
 use log::{Level, LevelFilter};
-use std::sync::Arc;
+use once_cell::sync::{Lazy, OnceCell};
+use std::sync::{Arc, OnceLock};
+use tokio::task::JoinHandle;
 use LingTransmit::server::event::ServerEvent;
 use LingTransmit::server::Client::Client;
 use LingTransmit::server::Server;
+use LingTransmit::shell::{register_command, start_shell, CommandActuators};
 use LingTransmit::ssl::ServerCert;
 
 #[tokio::main]
@@ -21,17 +25,30 @@ async fn main() {
     let pri = include_bytes!("../../ssl/test.pem");
     let passwd = include_str!("../../ssl/pass.txt");
     init_log();
-    let server_cert = ServerCert::init_buffer_password(
-        &cert.to_vec(),
-        &pri.to_vec(),
-        passwd,
-    )
+    let server_cert = ServerCert::init_buffer_password(&cert.to_vec(), &pri.to_vec(), passwd)
         .expect("解析证书失败");
-    let server = Server::new_tcp("0.0.0.0:11451", server_cert, Arc::new(Event {}))
+    let server = Arc::new(
+        Server::new_tcp("0.0.0.0:11451", server_cert, Arc::new(Event {}))
             .await
-        .expect("启动服务端失败");
-    server.start_accept().await;
+            .expect("启动服务端失败"),
+    );
+    let task = start_server(server.clone());
+    register_command(
+        "exit".to_string(),
+        Box::new(ExitCommand::new(server.clone())),
+    )
+    .await
+    .expect("注册命令失败");
+    start_shell().await;
+    task.await.expect("关闭服务器失败");
 }
+
+fn start_server(server: Arc<Server>) -> JoinHandle<()> {
+    tokio::spawn(async move {
+        server.start_accept().await;
+    })
+}
+
 fn get_time() -> String {
     let now = Local::now();
     now.format("%Y-%m-%d %H:%M:%S").to_string()
@@ -57,6 +74,24 @@ fn init_log() {
         .apply();
 }
 
+struct ExitCommand {
+    server: Arc<Server>,
+}
+
+impl ExitCommand {
+    pub fn new(server: Arc<Server>) -> Self {
+        ExitCommand { server }
+    }
+}
+
+#[async_trait]
+impl CommandActuators for ExitCommand {
+    async fn execute(&self, command: String) {
+        self.server.close().await;
+        self.exit_shell();
+    }
+}
+
 struct Event {}
 
 #[async_trait]

src/lib.rs

@@ -2,3 +2,4 @@ pub mod server;
 pub mod close_sender;
 pub mod packet;
 pub mod ssl;
+pub mod shell;

src/server/accept.rs

@@ -30,7 +30,7 @@ impl OwnedWriteHalfAbstraction for tcp::OwnedWriteHalf {}
 impl OwnedWriteHalfAbstraction for unix::OwnedWriteHalf {}
 
 #[async_trait]
-pub trait AcceptSocket {
+pub trait AcceptSocket: Send + Sync {
     async fn accept(
         &self,
     ) -> io::Result<(

src/server/mod.rs

@@ -29,7 +29,7 @@ pub type ClientList = Arc<Mutex<HashMap<ClientID, Arc<Client::Client>>>>;
 
 /// 服务器抽象
 pub struct Server {
-    listener: Box<dyn AcceptSocket>,
+    listener: Arc<dyn AcceptSocket>,
     close_sender: CloseSender,
     client_list: ClientList,
     next_id: AtomicU64,
@@ -38,7 +38,7 @@ pub struct Server {
 }
 
 impl Server {
-    fn new(listener: Box<dyn AcceptSocket>, cert: ServerCert, event: Arc<dyn ServerEvent>) -> Self {
+    fn new(listener: Arc<dyn AcceptSocket>, cert: ServerCert, event: Arc<dyn ServerEvent>) -> Self {
         Server {
             listener,
             close_sender: CloseSender::new(),
@@ -55,7 +55,7 @@ impl Server {
         event: Arc<dyn ServerEvent>,
     ) -> io::Result<Self> {
         let listener = TcpListener::bind(addr).await?;
-        Ok(Server::new(Box::new(listener), cert, event))
+        Ok(Server::new(Arc::new(listener), cert, event))
     }
 
     pub async fn new_unix<P>(
@@ -67,7 +67,7 @@ impl Server {
         P: AsRef<Path>,
     {
         let unix = UnixListener::bind(path)?;
-        Ok(Server::new(Box::new(unix), cert, event))
+        Ok(Server::new(Arc::new(unix), cert, event))
     }
 
     /// 广播关闭消息

src/shell/mod.rs

@@ -0,0 +1,81 @@
+// 版权所有 (c) ling 保留所有权利。
+// 除非另行说明，否则仅允许在LingTransmit中使用此文件中的代码。
+//
+// 由 ling 创建于 2025/1/19.
+#![allow(non_snake_case)]
+
+use colored::Colorize;
+use lazy_static::lazy_static;
+use once_cell::sync::Lazy;
+use std::collections::HashMap;
+use std::io;
+use std::io::Write;
+use std::sync::atomic::{AtomicBool, Ordering};
+use tokio::sync::Mutex;
+
+/// 命令驱动器
+#[async_trait::async_trait]
+pub trait CommandActuators: Send + Sync {
+    async fn execute(&self, command: String);
+
+    /// 调用后，shell将停止读取输入
+    fn exit_shell(&self) {
+        SHELL_FLAG.store(false, Ordering::Release);
+    }
+}
+
+lazy_static! {
+    static ref COMMAND_MAP: Lazy<Mutex<HashMap<String, Box<dyn CommandActuators>>>> =
+        Lazy::new(|| { Mutex::new(HashMap::new()) });
+}
+static SHELL_FLAG: AtomicBool = AtomicBool::new(true);
+
+/// 启动shell
+pub async fn start_shell() {
+    loop {
+        print_shell();
+        let mut input = String::new();
+        io::stdin().read_line(&mut input).expect("读取输入出错");
+        let input = input.trim();
+        run_command(input.to_string()).await;
+        if !SHELL_FLAG.load(Ordering::Acquire) {
+            break;
+        }
+    }
+}
+
+/// 注册命令
+pub async fn register_command(
+    name: String,
+    actuators: Box<dyn CommandActuators>,
+) -> Result<(), String> {
+    let mut map = COMMAND_MAP.lock().await;
+    if let Some(_) = map.get(&name) {
+        return Err(format!("命令 {} 已经存在", name));
+    }
+    map.insert(name, actuators);
+    Ok(())
+}
+
+/// 取消注册命令
+pub async fn unregister_command(name: &str) {
+    let mut map = COMMAND_MAP.lock().await;
+    map.remove(name);
+}
+
+async fn run_command(command: String) {
+    if let Some(first_word) = command.split_whitespace().next() {
+        let map = COMMAND_MAP.lock().await;
+        if let Some(actuators) = map.get(first_word) {
+            actuators.execute(command).await;
+        } else {
+            println!("{}", format!("找不到命令：{}", first_word).red());
+        }
+    }
+}
+
+/// 打印提示符
+fn print_shell() {
+    print!("[root@ling] # ");
+    io::stdout().flush().unwrap();
+}