改造消息架构

Author: pysrc

dsweb-core/src/server.rs

@@ -4,42 +4,59 @@ use crate::screen::Cap;
 use enigo::Enigo;
 use enigo::KeyboardControllable;
 use enigo::MouseControllable;
-use flate2::Compression;
 use flate2::write::DeflateEncoder;
-use websocket::OwnedMessage;
-use websocket::sync::Reader;
-use websocket::sync::Server;
-use websocket::sync::Writer;
+use flate2::Compression;
 use std::io::Write;
 use std::net::Ipv4Addr;
 use std::net::SocketAddrV4;
 use std::net::TcpStream;
-use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use websocket::sync::Reader;
+use websocket::sync::Server;
+use websocket::sync::Writer;
+use websocket::OwnedMessage;
 
 pub fn run(port: u16) {
     let host = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), port);
     let server = Server::bind(host).unwrap();
+    let isctrl = Arc::new(AtomicBool::new(false));
     for request in server.filter_map(Result::ok) {
-        if !request.protocols().contains(&"diffscreen".to_string()) {
+        if request.protocols().len() != 1 {
             request.reject().unwrap();
             continue;
         }
-        println!("New income !");
-        let client = request.use_protocol("diffscreen").accept().unwrap();
-        let (receiver, sender) = client.split().unwrap();
-        let running = Arc::new(AtomicBool::new(true));
-        let th1run = running.clone();
-        let _ = std::thread::spawn(move || {
-            screen_stream(sender, th1run);
-        });
-        let th = std::thread::spawn(move || {
-            event(receiver);
-        });
-        let _ = th.join();
-        running.store(false, Ordering::Relaxed);
-        println!("Break !");
+        let proto = request.protocols()[0].as_str();
+        match proto {
+            "diffscreen" => {
+                if isctrl.load(Ordering::Relaxed) {
+                    request.reject().unwrap();
+                    continue;
+                }
+                println!("New income !");
+                let client = request.use_protocol("diffscreen").accept().unwrap();
+                let (receiver, sender) = client.split().unwrap();
+                isctrl.store(true, Ordering::Relaxed);
+                let isctrlc = isctrl.clone();
+                let _ = std::thread::spawn(move || {
+                    screen_stream(sender, isctrlc);
+                });
+                let isctrlc = isctrl.clone();
+                let _ = std::thread::spawn(move || {
+                    event(receiver);
+                    isctrlc.store(false, Ordering::Relaxed);
+                    println!("Break !");
+                });
+            }
+            "paste" => {
+                // 处理粘贴
+            }
+            _ => {
+                request.reject().unwrap();
+                continue;
+            }
+        }
     }
 }
 
@@ -53,56 +70,52 @@ fn event(mut stream: Reader<TcpStream>) {
             Ok(message) => message,
             Err(e) => {
                 eprintln!("Msg err {}", e);
-                continue;
+                return;
             }
         };
         match message {
-            OwnedMessage::Binary(cmd) => {
-                match cmd[0] {
-                    dscom::KEY_UP => {
-                        if let Some(key) = key_mouse::key_to_enigo(cmd[1]) {
-                            enigo.key_up(key);
-                        }
-                    }
-                    dscom::KEY_DOWN => {
-                        if let Some(key) = key_mouse::key_to_enigo(cmd[1]) {
-                            enigo.key_down(key);
-                        }
+            OwnedMessage::Binary(cmd) => match cmd[0] {
+                dscom::KEY_UP => {
+                    if let Some(key) = key_mouse::key_to_enigo(cmd[1]) {
+                        enigo.key_up(key);
                     }
-                    dscom::MOUSE_KEY_UP => {
-                        if let Some(key) = key_mouse::mouse_to_engin(cmd[1]) {
-                            enigo.mouse_up(key);
-                        }
-                    }
-                    dscom::MOUSE_KEY_DOWN => {
-                        if let Some(key) = key_mouse::mouse_to_engin(cmd[1]) {
-                            enigo.mouse_down(key);
-                        }
-                    }
-                    dscom::MOUSE_WHEEL_UP => {
-                        enigo.mouse_scroll_y(-2);
-                    }
-                    dscom::MOUSE_WHEEL_DOWN => {
-                        enigo.mouse_scroll_y(2);
+                }
+                dscom::KEY_DOWN => {
+                    if let Some(key) = key_mouse::key_to_enigo(cmd[1]) {
+                        enigo.key_down(key);
                     }
-                    dscom::MOVE => {
-                        let x = ((cmd[1] as i32) << 8) | (cmd[2] as i32);
-                        let y = ((cmd[3] as i32) << 8) | (cmd[4] as i32);
-                        enigo.mouse_move_to(x, y);
+                }
+                dscom::MOUSE_KEY_UP => {
+                    if let Some(key) = key_mouse::mouse_to_engin(cmd[1]) {
+                        enigo.mouse_up(key);
                     }
-                    _ => {
-                        return;
+                }
+                dscom::MOUSE_KEY_DOWN => {
+                    if let Some(key) = key_mouse::mouse_to_engin(cmd[1]) {
+                        enigo.mouse_down(key);
                     }
                 }
-            }
+                dscom::MOUSE_WHEEL_UP => {
+                    enigo.mouse_scroll_y(-2);
+                }
+                dscom::MOUSE_WHEEL_DOWN => {
+                    enigo.mouse_scroll_y(2);
+                }
+                dscom::MOVE => {
+                    let x = ((cmd[1] as i32) << 8) | (cmd[2] as i32);
+                    let y = ((cmd[3] as i32) << 8) | (cmd[4] as i32);
+                    enigo.mouse_move_to(x, y);
+                }
+                _ => {
+                    return;
+                }
+            },
             OwnedMessage::Close(_) => {
                 println!("Front close !");
                 return;
             }
-            _=> {
-            }
+            _ => {}
         }
-        
     }
 }
 
@@ -114,11 +127,17 @@ fn get_data(data: OwnedMessage) -> Vec<u8> {
 }
 
 /// 初始化rgb块
-fn get_rgb_block(w: usize, h: usize, sw: usize, sh: usize, offset: usize) -> (Vec<Vec<u8>>, usize, usize) {
-    let srw = (w / sw) + if w % sw == 0 {0usize} else {1usize};
-    let srh = (h / sh) + if h % sh == 0 {0usize} else {1usize};
+fn get_rgb_block(
+    w: usize,
+    h: usize,
+    sw: usize,
+    sh: usize,
+    offset: usize,
+) -> (Vec<Vec<u8>>, usize, usize) {
+    let srw = (w / sw) + if w % sw == 0 { 0usize } else { 1usize };
+    let srh = (h / sh) + if h % sh == 0 { 0usize } else { 1usize };
     let num = srw * srh;
-    let mut buf = vec![vec![0u8;offset + sw*sh*3];num];
+    let mut buf = vec![vec![0u8; offset + sw * sh * 3]; num];
     for i in 0..(srw * srh) {
         buf[i][offset - 2] = (i >> 8) as u8;
         buf[i][offset - 1] = i as u8;
@@ -129,7 +148,7 @@ fn get_rgb_block(w: usize, h: usize, sw: usize, sh: usize, offset: usize) -> (Ve
 fn ws_send(mut stream: Writer<TcpStream>, data: Vec<u8>) -> (Writer<TcpStream>, Vec<u8>) {
     let om = OwnedMessage::Binary(data);
     stream.send_message(&om).unwrap();
-    return (stream, get_data(om))
+    return (stream, get_data(om));
 }
 
 fn screen_stream(mut stream: Writer<TcpStream>, running: Arc<AtomicBool>) {
@@ -139,7 +158,7 @@ fn screen_stream(mut stream: Writer<TcpStream>, running: Arc<AtomicBool>) {
     let block_len = srw * srh;
     let (mut b, _, _) = get_rgb_block(w, h, sw, sh, 2);
     // 发送w, h, sw, sh, mask
-    let mut meta = vec![0u8;9];
+    let mut meta = vec![0u8; 9];
     meta[0] = (w >> 8) as u8;
     meta[1] = w as u8;
     meta[2] = (h >> 8) as u8;