weyne/rustdesk
1
0
Fork 0
mirror of https://github.com/weyne85/rustdesk.git synced 2025-10-29 17:00:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into lan_discovery

Browse Source
This commit is contained in:
RustDesk
2022-01-10 17:34:51 +08:00
committed by GitHub
parent 35462eecfe e2a879692d
commit dfeb9a29c1
68 changed files with 3775 additions and 1605 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
libs/enigo/examples/key.rs
View File

@@ -6,7 +6,7 @@ fn main() {
thread::sleep(Duration::from_secs(2));
let mut enigo = Enigo::new();
enigo.key_down(Key::Layout('a'));
enigo.key_down(Key::Layout('a')).ok();
thread::sleep(Duration::from_secs(1));
enigo.key_up(Key::Layout('a'));
}

2
libs/enigo/examples/keyboard.rs
View File

@@ -10,7 +10,7 @@ fn main() {
enigo.key_sequence("Hello World! here is a lot of text ❤️");
// select all
enigo.key_down(Key::Control);
enigo.key_down(Key::Control).ok();
enigo.key_click(Key::Layout('a'));
enigo.key_up(Key::Control);
}

2
libs/enigo/examples/mouse.rs
View File

@@ -11,7 +11,7 @@ fn main() {
enigo.mouse_move_to(500, 200);
thread::sleep(wait_time);
enigo.mouse_down(MouseButton::Left);
enigo.mouse_down(MouseButton::Left).ok();
thread::sleep(wait_time);
enigo.mouse_move_relative(100, 100);

2
libs/enigo/examples/timer.rs
View File

@@ -16,7 +16,7 @@ fn main() {
println!("{:?}", time);
// select all
enigo.key_down(Key::Control);
enigo.key_down(Key::Control).ok();
enigo.key_click(Key::Layout('a'));
enigo.key_up(Key::Control);
}

12
libs/hbb_common/Cargo.toml
View File

@@ -7,16 +7,16 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
protobuf = { version = "3.0.0-pre", git = "https://github.com/stepancheg/rust-protobuf" }
tokio = { version = "1.10", features = ["full"] }
protobuf = "3.0.0-alpha.2"
tokio = { version = "1.15", features = ["full"] }
tokio-util = { version = "0.6", features = ["full"] }
futures = "0.3"
bytes = "1.0"
bytes = "1.1"
log = "0.4"
env_logger = "0.9"
socket2 = { version = "0.3", features = ["reuseport"] }
zstd = "0.9"
quinn = {version = "0.6", optional = true }
quinn = {version = "0.8", optional = true }
anyhow = "1.0"
futures-util = "0.3"
directories-next = "2.0"
@@ -28,6 +28,8 @@ confy = { git = "https://github.com/open-trade/confy" }
dirs-next = "2.0"
filetime = "0.2"
sodiumoxide = "0.2"
regex = "1.4"
tokio-socks = { git = "https://github.com/fufesou/tokio-socks" }
[target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
mac_address = "1.1"
@@ -36,7 +38,7 @@ mac_address = "1.1"
quic = ["quinn"]
[build-dependencies]
protobuf-codegen-pure = { version = "3.0.0-pre", git = "https://github.com/stepancheg/rust-protobuf" }
protobuf-codegen-pure = "3.0.0-alpha.2"
[target.'cfg(target_os = "windows")'.dependencies]
winapi = { version = "0.3", features = ["winuser"] }

2
libs/hbb_common/protos/message.proto
View File

@@ -66,6 +66,7 @@ message MouseEvent {
}
enum ControlKey {
Unknown = 0;
Alt = 1;
Backspace = 2;
CapsLock = 3;
@@ -342,7 +343,6 @@ message PublicKey {
message SignedId {
bytes id = 1;
bytes pk = 2;
}
message AudioFormat {

86
libs/hbb_common/src/config.rs
View File

@@ -6,14 +6,13 @@ use sodiumoxide::crypto::sign;
use std::{
collections::HashMap,
fs,
net::SocketAddr,
net::{IpAddr, Ipv4Addr, SocketAddr},
path::{Path, PathBuf},
sync::{Arc, Mutex, RwLock},
time::SystemTime,
};
pub const APP_NAME: &str = "RustDesk";
pub const BIND_INTERFACE: &str = "0.0.0.0";
pub const RENDEZVOUS_TIMEOUT: u64 = 12_000;
pub const CONNECT_TIMEOUT: u64 = 18_000;
pub const COMPRESS_LEVEL: i32 = 3;
@@ -55,7 +54,11 @@ pub const RENDEZVOUS_SERVERS: &'static [&'static str] = &[
pub const RENDEZVOUS_PORT: i32 = 21116;
pub const RELAY_PORT: i32 = 21117;
pub const SERVER_UDP_PORT: u16 = 21001; // udp
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum NetworkType {
Direct,
ProxySocks,
}
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct Config {
@@ -73,6 +76,16 @@ pub struct Config {
keys_confirmed: HashMap<String, bool>,
}
#[derive(Debug, Default, PartialEq, Serialize, Deserialize, Clone)]
pub struct Socks5Server {
#[serde(default)]
pub proxy: String,
#[serde(default)]
pub username: String,
#[serde(default)]
pub password: String,
}
// more variable configs
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct Config2 {
@@ -87,6 +100,9 @@ pub struct Config2 {
#[serde(default)]
serial: i32,
#[serde(default)]
socks: Option<Socks5Server>,
// the other scalar value must before this
#[serde(default)]
pub options: HashMap<String, String>,
@@ -276,8 +292,8 @@ impl Config {
return "".into();
}
#[allow(unreachable_code)]
pub fn log_path() -> PathBuf {
#[allow(unreachable_code)]
#[cfg(target_os = "macos")]
{
if let Some(path) = dirs_next::home_dir().as_mut() {
@@ -329,10 +345,10 @@ impl Config {
#[inline]
pub fn get_any_listen_addr() -> SocketAddr {
format!("{}:0", BIND_INTERFACE).parse().unwrap()
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
}
pub fn get_rendezvous_server() -> SocketAddr {
pub fn get_rendezvous_server() -> String {
let mut rendezvous_server = Self::get_option("custom-rendezvous-server");
if rendezvous_server.is_empty() {
rendezvous_server = CONFIG2.write().unwrap().rendezvous_server.clone();
@@ -346,11 +362,7 @@ impl Config {
if !rendezvous_server.contains(":") {
rendezvous_server = format!("{}:{}", rendezvous_server, RENDEZVOUS_PORT);
}
if let Ok(addr) = crate::to_socket_addr(&rendezvous_server) {
addr
} else {
Self::get_any_listen_addr()
}
rendezvous_server
}
pub fn get_rendezvous_servers() -> Vec<String> {
@@ -490,6 +502,9 @@ impl Config {
pub fn set_key_pair(pair: (Vec<u8>, Vec<u8>)) {
let mut config = CONFIG.write().unwrap();
if config.key_pair == pair {
return;
}
config.key_pair = pair;
config.store();
}
@@ -522,6 +537,9 @@ impl Config {
pub fn set_options(v: HashMap<String, String>) {
let mut config = CONFIG2.write().unwrap();
if config.options == v {
return;
}
config.options = v;
config.store();
}
@@ -621,6 +639,26 @@ impl Config {
pub fn get_remote_id() -> String {
CONFIG2.read().unwrap().remote_id.clone()
}
pub fn set_socks(socks: Option<Socks5Server>) {
let mut config = CONFIG2.write().unwrap();
if config.socks == socks {
return;
}
config.socks = socks;
config.store();
}
pub fn get_socks() -> Option<Socks5Server> {
CONFIG2.read().unwrap().socks.clone()
}
pub fn get_network_type() -> NetworkType {
match &CONFIG2.read().unwrap().socks {
None => NetworkType::Direct,
Some(_) => NetworkType::ProxySocks,
}
}
}
const PEERS: &str = "peers";
@@ -690,6 +728,32 @@ impl PeerConfig {
}
}
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct Fav {
#[serde(default)]
pub peers: Vec<String>,
}
impl Fav {
pub fn load() -> Fav {
let _ = CONFIG.read().unwrap(); // for lock
match confy::load_path(&Config::file_("_fav")) {
Ok(fav) => fav,
Err(err) => {
log::error!("Failed to load fav: {}", err);
Default::default()
}
}
}
pub fn store(peers: Vec<String>) {
let f = Fav { peers };
if let Err(err) = confy::store_path(Config::file_("_fav"), f) {
log::error!("Failed to store fav: {}", err);
}
}
}
#[cfg(test)]
mod tests {
use super::*;

23
libs/hbb_common/src/lib.rs
View File

@@ -13,12 +13,13 @@ pub use protobuf;
use std::{
fs::File,
io::{self, BufRead},
net::{Ipv4Addr, SocketAddr, SocketAddrV4, ToSocketAddrs},
net::{Ipv4Addr, SocketAddr, SocketAddrV4},
path::Path,
time::{self, SystemTime, UNIX_EPOCH},
};
pub use tokio;
pub use tokio_util;
pub mod socket_client;
pub mod tcp;
pub mod udp;
pub use env_logger;
@@ -30,7 +31,11 @@ pub use anyhow::{self, bail};
pub use futures_util;
pub mod config;
pub mod fs;
pub use regex;
pub use sodiumoxide;
pub use tokio_socks;
pub use tokio_socks::IntoTargetAddr;
pub use tokio_socks::TargetAddr;
#[cfg(feature = "quic")]
pub type Stream = quic::Connection;
@@ -151,14 +156,6 @@ pub fn get_version_from_url(url: &str) -> String {
"".to_owned()
}
pub fn to_socket_addr(host: &str) -> ResultType<SocketAddr> {
let addrs: Vec<SocketAddr> = host.to_socket_addrs()?.collect();
if addrs.is_empty() {
bail!("Failed to solve {}", host);
}
Ok(addrs[0])
}
pub fn gen_version() {
let mut file = File::create("./src/version.rs").unwrap();
for line in read_lines("Cargo.toml").unwrap() {
@@ -183,6 +180,14 @@ where
Ok(io::BufReader::new(file).lines())
}
pub fn get_version_number(v: &str) -> i64 {
let mut n = 0;
for x in v.split(".") {
n = n * 1000 + x.parse::<i64>().unwrap_or(0);
}
n
}
#[cfg(test)]
mod tests {
use super::*;

91
libs/hbb_common/src/socket_client.rs Normal file
View File

@@ -0,0 +1,91 @@
use crate::{
config::{Config, NetworkType},
tcp::FramedStream,
udp::FramedSocket,
ResultType,
};
use anyhow::Context;
use std::net::SocketAddr;
use tokio::net::ToSocketAddrs;
use tokio_socks::{IntoTargetAddr, TargetAddr};
fn to_socket_addr(host: &str) -> ResultType<SocketAddr> {
use std::net::ToSocketAddrs;
host.to_socket_addrs()?.next().context("Failed to solve")
}
pub fn get_target_addr(host: &str) -> ResultType<TargetAddr<'static>> {
let addr = match Config::get_network_type() {
NetworkType::Direct => to_socket_addr(&host)?.into_target_addr()?,
NetworkType::ProxySocks => host.into_target_addr()?,
}
.to_owned();
Ok(addr)
}
pub fn test_if_valid_server(host: &str) -> String {
let mut host = host.to_owned();
if !host.contains(":") {
host = format!("{}:{}", host, 0);
}
match Config::get_network_type() {
NetworkType::Direct => match to_socket_addr(&host) {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
NetworkType::ProxySocks => match &host.into_target_addr() {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
}
}
pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
target: T,
local: SocketAddr,
ms_timeout: u64,
) -> ResultType<FramedStream> {
let target_addr = target.into_target_addr()?;
if let Some(conf) = Config::get_socks() {
FramedStream::connect(
conf.proxy.as_str(),
target_addr,
local,
conf.username.as_str(),
conf.password.as_str(),
ms_timeout,
)
.await
} else {
let addr = std::net::ToSocketAddrs::to_socket_addrs(&target_addr)?
.next()
.context("Invalid target addr")?;
Ok(FramedStream::new(addr, local, ms_timeout).await?)
}
}
pub async fn new_udp<T: ToSocketAddrs>(local: T, ms_timeout: u64) -> ResultType<FramedSocket> {
match Config::get_socks() {
None => Ok(FramedSocket::new(local).await?),
Some(conf) => {
let socket = FramedSocket::new_proxy(
conf.proxy.as_str(),
local,
conf.username.as_str(),
conf.password.as_str(),
ms_timeout,
)
.await?;
Ok(socket)
}
}
}
pub async fn rebind_udp<T: ToSocketAddrs>(local: T) -> ResultType<Option<FramedSocket>> {
match Config::get_network_type() {
NetworkType::Direct => Ok(Some(FramedSocket::new(local).await?)),
_ => Ok(None),
}
}

161
libs/hbb_common/src/tcp.rs
View File

@@ -4,16 +4,31 @@ use futures::{SinkExt, StreamExt};
use protobuf::Message;
use sodiumoxide::crypto::secretbox::{self, Key, Nonce};
use std::{
io::{Error, ErrorKind},
io::{self, Error, ErrorKind},
net::SocketAddr,
ops::{Deref, DerefMut},
pin::Pin,
task::{Context, Poll},
};
use tokio::net::{lookup_host, TcpListener, TcpSocket, TcpStream, ToSocketAddrs};
use tokio::{
io::{AsyncRead, AsyncWrite, ReadBuf},
net::{lookup_host, TcpListener, TcpSocket, ToSocketAddrs},
};
use tokio_socks::{tcp::Socks5Stream, IntoTargetAddr, ToProxyAddrs};
use tokio_util::codec::Framed;
pub struct FramedStream(Framed<TcpStream, BytesCodec>, Option<(Key, u64, u64)>);
pub trait TcpStreamTrait: AsyncRead + AsyncWrite + Unpin {}
pub struct DynTcpStream(Box<dyn TcpStreamTrait + Send>);
pub struct FramedStream(
Framed<DynTcpStream, BytesCodec>,
SocketAddr,
Option<(Key, u64, u64)>,
u64,
);
impl Deref for FramedStream {
type Target = Framed<TcpStream, BytesCodec>;
type Target = Framed<DynTcpStream, BytesCodec>;
fn deref(&self) -> &Self::Target {
&self.0
@@ -26,6 +41,20 @@ impl DerefMut for FramedStream {
}
}
impl Deref for DynTcpStream {
type Target = Box<dyn TcpStreamTrait + Send>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl DerefMut for DynTcpStream {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
fn new_socket(addr: std::net::SocketAddr, reuse: bool) -> Result<TcpSocket, std::io::Error> {
let socket = match addr {
std::net::SocketAddr::V4(..) => TcpSocket::new_v4()?,
@@ -44,8 +73,8 @@ fn new_socket(addr: std::net::SocketAddr, reuse: bool) -> Result<TcpSocket, std:
}
impl FramedStream {
pub async fn new<T: ToSocketAddrs, T2: ToSocketAddrs>(
remote_addr: T,
pub async fn new<T1: ToSocketAddrs, T2: ToSocketAddrs>(
remote_addr: T1,
local_addr: T2,
ms_timeout: u64,
) -> ResultType<Self> {
@@ -56,23 +85,86 @@ impl FramedStream {
new_socket(local_addr, true)?.connect(remote_addr),
)
.await??;
return Ok(Self(Framed::new(stream, BytesCodec::new()), None));
let addr = stream.local_addr()?;
return Ok(Self(
Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
addr,
None,
0,
));
}
}
bail!("could not resolve to any address");
}
pub fn from(stream: TcpStream) -> Self {
Self(Framed::new(stream, BytesCodec::new()), None)
pub async fn connect<'a, 't, P, T1, T2>(
proxy: P,
target: T1,
local: T2,
username: &'a str,
password: &'a str,
ms_timeout: u64,
) -> ResultType<Self>
where
P: ToProxyAddrs,
T1: IntoTargetAddr<'t>,
T2: ToSocketAddrs,
{
if let Some(local) = lookup_host(&local).await?.next() {
if let Some(proxy) = proxy.to_proxy_addrs().next().await {
let stream =
super::timeout(ms_timeout, new_socket(local, true)?.connect(proxy?)).await??;
let stream = if username.trim().is_empty() {
super::timeout(
ms_timeout,
Socks5Stream::connect_with_socket(stream, target),
)
.await??
} else {
super::timeout(
ms_timeout,
Socks5Stream::connect_with_password_and_socket(
stream, target, username, password,
),
)
.await??
};
let addr = stream.local_addr()?;
return Ok(Self(
Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
addr,
None,
0,
));
};
};
bail!("could not resolve to any address");
}
pub fn local_addr(&self) -> SocketAddr {
self.1
}
pub fn set_send_timeout(&mut self, ms: u64) {
self.3 = ms;
}
pub fn from(stream: impl TcpStreamTrait + Send + 'static, addr: SocketAddr) -> Self {
Self(
Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
addr,
None,
0,
)
}
pub fn set_raw(&mut self) {
self.0.codec_mut().set_raw();
self.1 = None;
self.2 = None;
}
pub fn is_secured(&self) -> bool {
self.1.is_some()
self.2.is_some()
}
#[inline]
@@ -83,24 +175,29 @@ impl FramedStream {
#[inline]
pub async fn send_raw(&mut self, msg: Vec<u8>) -> ResultType<()> {
let mut msg = msg;
if let Some(key) = self.1.as_mut() {
if let Some(key) = self.2.as_mut() {
key.1 += 1;
let nonce = Self::get_nonce(key.1);
msg = secretbox::seal(&msg, &nonce, &key.0);
}
self.0.send(bytes::Bytes::from(msg)).await?;
self.send_bytes(bytes::Bytes::from(msg)).await?;
Ok(())
}
#[inline]
pub async fn send_bytes(&mut self, bytes: Bytes) -> ResultType<()> {
self.0.send(bytes).await?;
if self.3 > 0 {
super::timeout(self.3, self.0.send(bytes)).await??;
} else {
self.0.send(bytes).await?;
}
Ok(())
}
#[inline]
pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> {
let mut res = self.0.next().await;
if let Some(key) = self.1.as_mut() {
if let Some(key) = self.2.as_mut() {
if let Some(Ok(bytes)) = res.as_mut() {
key.2 += 1;
let nonce = Self::get_nonce(key.2);
@@ -128,7 +225,7 @@ impl FramedStream {
}
pub fn set_key(&mut self, key: Key) {
self.1 = Some((key, 0, 0));
self.2 = Some((key, 0, 0));
}
fn get_nonce(seqnum: u64) -> Nonce {
@@ -152,3 +249,35 @@ pub async fn new_listener<T: ToSocketAddrs>(addr: T, reuse: bool) -> ResultType<
bail!("could not resolve to any address");
}
}
impl Unpin for DynTcpStream {}
impl AsyncRead for DynTcpStream {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
AsyncRead::poll_read(Pin::new(&mut self.0), cx, buf)
}
}
impl AsyncWrite for DynTcpStream {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
AsyncWrite::poll_write(Pin::new(&mut self.0), cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
AsyncWrite::poll_flush(Pin::new(&mut self.0), cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
AsyncWrite::poll_shutdown(Pin::new(&mut self.0), cx)
}
}
impl<R: AsyncRead + AsyncWrite + Unpin> TcpStreamTrait for R {}

127
libs/hbb_common/src/udp.rs
View File

@@ -1,24 +1,17 @@
use crate::{bail, ResultType};
use bytes::BytesMut;
use anyhow::anyhow;
use bytes::{Bytes, BytesMut};
use futures::{SinkExt, StreamExt};
use protobuf::Message;
use socket2::{Domain, Socket, Type};
use std::{
io::Error,
net::SocketAddr,
ops::{Deref, DerefMut},
};
use tokio::{net::ToSocketAddrs, net::UdpSocket};
use std::net::SocketAddr;
use tokio::net::{ToSocketAddrs, UdpSocket};
use tokio_socks::{udp::Socks5UdpFramed, IntoTargetAddr, TargetAddr, ToProxyAddrs};
use tokio_util::{codec::BytesCodec, udp::UdpFramed};
pub struct FramedSocket(UdpFramed<BytesCodec>);
impl Deref for FramedSocket {
type Target = UdpFramed<BytesCodec>;
fn deref(&self) -> &Self::Target {
&self.0
}
pub enum FramedSocket {
Direct(UdpFramed<BytesCodec>),
ProxySocks(Socks5UdpFramed),
}
fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
@@ -38,52 +31,110 @@ fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
Ok(socket)
}
impl DerefMut for FramedSocket {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl FramedSocket {
pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> {
let socket = UdpSocket::bind(addr).await?;
Ok(Self(UdpFramed::new(socket, BytesCodec::new())))
Ok(Self::Direct(UdpFramed::new(socket, BytesCodec::new())))
}
#[allow(clippy::never_loop)]
pub async fn new_reuse<T: std::net::ToSocketAddrs>(addr: T) -> ResultType<Self> {
for addr in addr.to_socket_addrs()? {
return Ok(Self(UdpFramed::new(
UdpSocket::from_std(new_socket(addr, true)?.into_udp_socket())?,
let socket = new_socket(addr, true)?.into_udp_socket();
return Ok(Self::Direct(UdpFramed::new(
UdpSocket::from_std(socket)?,
BytesCodec::new(),
)));
}
bail!("could not resolve to any address");
}
pub async fn new_proxy<'a, 't, P: ToProxyAddrs, T: ToSocketAddrs>(
proxy: P,
local: T,
username: &'a str,
password: &'a str,
ms_timeout: u64,
) -> ResultType<Self> {
let framed = if username.trim().is_empty() {
super::timeout(ms_timeout, Socks5UdpFramed::connect(proxy, Some(local))).await??
} else {
super::timeout(
ms_timeout,
Socks5UdpFramed::connect_with_password(proxy, Some(local), username, password),
)
.await??
};
log::trace!(
"Socks5 udp connected, local addr: {:?}, target addr: {}",
framed.local_addr(),
framed.socks_addr()
);
Ok(Self::ProxySocks(framed))
}
#[inline]
pub async fn send(&mut self, msg: &impl Message, addr: SocketAddr) -> ResultType<()> {
self.0
.send((bytes::Bytes::from(msg.write_to_bytes().unwrap()), addr))
.await?;
pub async fn send(
&mut self,
msg: &impl Message,
addr: impl IntoTargetAddr<'_>,
) -> ResultType<()> {
let addr = addr.into_target_addr()?.to_owned();
let send_data = Bytes::from(msg.write_to_bytes()?);
let _ = match self {
Self::Direct(f) => match addr {
TargetAddr::Ip(addr) => f.send((send_data, addr)).await?,
_ => unreachable!(),
},
Self::ProxySocks(f) => f.send((send_data, addr)).await?,
};
Ok(())
}
// https://stackoverflow.com/a/68733302/1926020
#[inline]
pub async fn send_raw(
&mut self,
msg: &'static [u8],
addr: impl IntoTargetAddr<'static>,
) -> ResultType<()> {
let addr = addr.into_target_addr()?.to_owned();
let _ = match self {
Self::Direct(f) => match addr {
TargetAddr::Ip(addr) => f.send((Bytes::from(msg), addr)).await?,
_ => unreachable!(),
},
Self::ProxySocks(f) => f.send((Bytes::from(msg), addr)).await?,
};
Ok(())
}
#[inline]
pub async fn send_raw(&mut self, msg: &'static [u8], addr: SocketAddr) -> ResultType<()> {
self.0.send((bytes::Bytes::from(msg), addr)).await?;
Ok(())
pub async fn next(&mut self) -> Option<ResultType<(BytesMut, TargetAddr<'static>)>> {
match self {
Self::Direct(f) => match f.next().await {
Some(Ok((data, addr))) => {
Some(Ok((data, addr.into_target_addr().ok()?.to_owned())))
}
Some(Err(e)) => Some(Err(anyhow!(e))),
None => None,
},
Self::ProxySocks(f) => match f.next().await {
Some(Ok((data, _))) => Some(Ok((data.data, data.dst_addr))),
Some(Err(e)) => Some(Err(anyhow!(e))),
None => None,
},
}
}
#[inline]
pub async fn next(&mut self) -> Option<Result<(BytesMut, SocketAddr), Error>> {
self.0.next().await
}
#[inline]
pub async fn next_timeout(&mut self, ms: u64) -> Option<Result<(BytesMut, SocketAddr), Error>> {
pub async fn next_timeout(
&mut self,
ms: u64,
) -> Option<ResultType<(BytesMut, TargetAddr<'static>)>> {
if let Ok(res) =
tokio::time::timeout(std::time::Duration::from_millis(ms), self.0.next()).await
tokio::time::timeout(std::time::Duration::from_millis(ms), self.next()).await
{
res
} else {

1
libs/scrap/examples/record-screen.rs
View File

@@ -43,7 +43,6 @@ struct Args {
flag_time: Option<u64>,
flag_fps: u64,
flag_bv: u32,
flag_ba: u32,
}
#[derive(Debug, serde::Deserialize)]

4
libs/scrap/src/common/convert.rs
View File

@@ -9,8 +9,8 @@ extern "C" {
src_stride_argb: c_int,
dst_argb: *mut u8,
dst_stride_argb: c_int,
width: c_int,
height: c_int,
src_width: c_int,
src_height: c_int,
mode: c_int,
) -> c_int;

92
libs/scrap/src/dxgi/mod.rs
View File

@@ -3,32 +3,20 @@ pub mod gdi;
pub use gdi::CapturerGDI;
use winapi::{
shared::dxgi::{
CreateDXGIFactory1, IDXGIAdapter1, IDXGIFactory1, IDXGIResource, IDXGISurface,
IID_IDXGIFactory1, IID_IDXGISurface, DXGI_MAP_READ, DXGI_OUTPUT_DESC,
DXGI_RESOURCE_PRIORITY_MAXIMUM,
shared::{
dxgi::*,
dxgi1_2::*,
dxgitype::*,
minwindef::{DWORD, FALSE, TRUE, UINT},
ntdef::LONG,
windef::HMONITOR,
winerror::*,
// dxgiformat::{DXGI_FORMAT, DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_FORMAT_420_OPAQUE},
},
shared::dxgi1_2::IDXGIOutputDuplication,
// shared::dxgiformat::{DXGI_FORMAT, DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_FORMAT_420_OPAQUE},
shared::dxgi1_2::{IDXGIOutput1, IID_IDXGIOutput1},
shared::dxgitype::DXGI_MODE_ROTATION,
shared::minwindef::{DWORD, FALSE, TRUE, UINT},
shared::ntdef::LONG,
shared::windef::HMONITOR,
shared::winerror::{
DXGI_ERROR_ACCESS_LOST, DXGI_ERROR_INVALID_CALL, DXGI_ERROR_NOT_CURRENTLY_AVAILABLE,
DXGI_ERROR_SESSION_DISCONNECTED, DXGI_ERROR_UNSUPPORTED, DXGI_ERROR_WAIT_TIMEOUT,
E_ACCESSDENIED, E_INVALIDARG, S_OK,
um::{
d3d11::*, d3dcommon::D3D_DRIVER_TYPE_UNKNOWN, unknwnbase::IUnknown, wingdi::*,
winnt::HRESULT, winuser::*,
},
um::d3d11::{
D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, IID_ID3D11Texture2D,
D3D11_CPU_ACCESS_READ, D3D11_SDK_VERSION, D3D11_USAGE_STAGING,
},
um::d3dcommon::D3D_DRIVER_TYPE_UNKNOWN,
um::unknwnbase::IUnknown,
um::wingdi::*,
um::winnt::HRESULT,
um::winuser::*,
};
pub struct ComPtr<T>(*mut T);
@@ -54,12 +42,11 @@ pub struct Capturer {
duplication: ComPtr<IDXGIOutputDuplication>,
fastlane: bool,
surface: ComPtr<IDXGISurface>,
data: *const u8,
len: usize,
width: usize,
height: usize,
use_yuv: bool,
yuv: Vec<u8>,
rotated: Vec<u8>,
gdi_capturer: Option<CapturerGDI>,
gdi_buffer: Vec<u8>,
}
@@ -158,10 +145,9 @@ impl Capturer {
width: display.width() as usize,
height: display.height() as usize,
display,
data: ptr::null(),
len: 0,
use_yuv,
yuv: Vec::new(),
rotated: Vec::new(),
gdi_capturer,
gdi_buffer: Vec::new(),
})
@@ -181,10 +167,9 @@ impl Capturer {
self.gdi_capturer.take();
}
unsafe fn load_frame(&mut self, timeout: UINT) -> io::Result<()> {
unsafe fn load_frame(&mut self, timeout: UINT) -> io::Result<(*const u8, i32)> {
let mut frame = ptr::null_mut();
let mut info = mem::MaybeUninit::uninit().assume_init();
self.data = ptr::null();
wrap_hresult((*self.duplication.0).AcquireNextFrame(timeout, &mut info, &mut frame))?;
let frame = ComPtr(frame);
@@ -200,9 +185,7 @@ impl Capturer {
self.surface = ComPtr(self.ohgodwhat(frame.0)?);
wrap_hresult((*self.surface.0).Map(&mut rect, DXGI_MAP_READ))?;
}
self.data = rect.pBits;
self.len = self.height * rect.Pitch as usize;
Ok(())
Ok((rect.pBits, rect.Pitch))
}
// copy from GPU memory to system memory
@@ -257,8 +240,42 @@ impl Capturer {
}
} else {
self.unmap();
self.load_frame(timeout)?;
slice::from_raw_parts(self.data, self.len)
let r = self.load_frame(timeout)?;
let rotate = match self.display.rotation() {
DXGI_MODE_ROTATION_IDENTITY | DXGI_MODE_ROTATION_UNSPECIFIED => 0,
DXGI_MODE_ROTATION_ROTATE90 => 90,
DXGI_MODE_ROTATION_ROTATE180 => 180,
DXGI_MODE_ROTATION_ROTATE270 => 270,
_ => {
return Err(io::Error::new(
io::ErrorKind::Other,
"Unknown roration".to_string(),
));
}
};
if rotate == 0 {
slice::from_raw_parts(r.0, r.1 as usize * self.height)
} else {
self.rotated.resize(self.width * self.height * 4, 0);
crate::common::ARGBRotate(
r.0,
r.1,
self.rotated.as_mut_ptr(),
4 * self.width as i32,
if rotate == 180 {
self.width
} else {
self.height
} as _,
if rotate != 180 {
self.width
} else {
self.height
} as _,
rotate,
);
&self.rotated[..]
}
}
};
Ok({
@@ -478,7 +495,10 @@ pub struct Display {
gdi: bool,
}
// https://github.com/dchapyshev/aspia/blob/59233c5d01a4d03ed6de19b03ce77d61a6facf79/source/base/desktop/win/screen_capture_utils.cc
// optimized for updated region
// https://github.com/dchapyshev/aspia/blob/master/source/base/desktop/win/dxgi_output_duplicator.cc
// rotation
// https://github.com/bryal/dxgcap-rs/blob/master/src/lib.rs
impl Display {
pub fn width(&self) -> LONG {