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 'hwcodec' into master

Browse Source
This commit is contained in:
21pages
2022-06-14 11:46:03 +08:00
parent 8753b28a96 42c7c5982c
commit 847c4acb07
17 changed files with 1782 additions and 638 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
src/client.rs
View File

@@ -12,6 +12,11 @@ use cpal::{
Device, Host, StreamConfig,
};
use magnum_opus::{Channels::*, Decoder as AudioDecoder};
use scrap::{
codec::{Decoder, DecoderCfg},
VpxDecoderConfig, VpxVideoCodecId,
};
use sha2::{Digest, Sha256};
use uuid::Uuid;
@@ -30,7 +35,6 @@ use hbb_common::{
tokio::time::Duration,
AddrMangle, ResultType, Stream,
};
use scrap::{Decoder, Image, VideoCodecId};
pub use super::lang::*;
pub mod file_trait;
@@ -717,7 +721,12 @@ pub struct VideoHandler {
impl VideoHandler {
pub fn new(latency_controller: Arc<Mutex<LatencyController>>) -> Self {
VideoHandler {
decoder: Decoder::new(VideoCodecId::VP9, (num_cpus::get() / 2) as _).unwrap(),
decoder: Decoder::new(DecoderCfg {
vpx: VpxDecoderConfig {
codec: VpxVideoCodecId::VP9,
num_threads: (num_cpus::get() / 2) as _,
},
}),
latency_controller,
rgb: Default::default(),
}
@@ -731,33 +740,18 @@ impl VideoHandler {
.update_video(vf.timestamp);
}
match &vf.union {
Some(video_frame::Union::vp9s(vp9s)) => self.handle_vp9s(vp9s),
Some(frame) => self.decoder.handle_video_frame(frame, &mut self.rgb),
_ => Ok(false),
}
}
pub fn handle_vp9s(&mut self, vp9s: &VP9s) -> ResultType<bool> {
let mut last_frame = Image::new();
for vp9 in vp9s.frames.iter() {
for frame in self.decoder.decode(&vp9.data)? {
drop(last_frame);
last_frame = frame;
}
}
for frame in self.decoder.flush()? {
drop(last_frame);
last_frame = frame;
}
if last_frame.is_null() {
Ok(false)
} else {
last_frame.rgb(1, true, &mut self.rgb);
Ok(true)
}
}
pub fn reset(&mut self) {
self.decoder = Decoder::new(VideoCodecId::VP9, 1).unwrap();
self.decoder = Decoder::new(DecoderCfg {
vpx: VpxDecoderConfig {
codec: VpxVideoCodecId::VP9,
num_threads: 1,
},
});
}
}
@@ -952,6 +946,11 @@ impl LoginConfigHandler {
msg.disable_clipboard = BoolOption::Yes.into();
n += 1;
}
// TODO: add option
let state = Decoder::video_codec_state();
msg.video_codec_state = hbb_common::protobuf::MessageField::some(state);
n += 1;
if n > 0 {
Some(msg)
} else {
@@ -1170,9 +1169,11 @@ where
let latency_controller = LatencyController::new();
let latency_controller_cl = latency_controller.clone();
// Create video_handler out of the thread below to ensure that the handler exists before client start.
// It will take a few tenths of a second for the first time, and then tens of milliseconds.
let mut video_handler = VideoHandler::new(latency_controller);
std::thread::spawn(move || {
let mut video_handler = VideoHandler::new(latency_controller);
loop {
if let Ok(data) = video_receiver.recv() {
match data {

29
src/ipc.rs
View File

@@ -1,6 +1,8 @@
use crate::rendezvous_mediator::RendezvousMediator;
#[cfg(not(any(target_os = "android", target_os = "ios")))]
pub use clipboard::ClipbaordFile;
#[cfg(feature = "hwcodec")]
use hbb_common::config::HwCodecConfig;
use hbb_common::{
allow_err, bail, bytes,
bytes_codec::BytesCodec,
@@ -73,7 +75,7 @@ pub enum FS {
WriteOffset {
id: i32,
file_num: i32,
offset_blk: u32
offset_blk: u32,
},
CheckDigest {
id: i32,
@@ -127,6 +129,8 @@ pub enum Data {
ClipbaordFile(ClipbaordFile),
ClipboardFileEnabled(bool),
PrivacyModeState((i32, PrivacyModeState)),
#[cfg(feature = "hwcodec")]
HwCodecConfig(Option<HashMap<String, String>>),
}
#[tokio::main(flavor = "current_thread")]
@@ -336,6 +340,12 @@ async fn handle(data: Data, stream: &mut Connection) {
.await
);
}
#[cfg(feature = "hwcodec")]
Data::HwCodecConfig(Some(config)) => {
for (k, v) in config {
HwCodecConfig::set_option(k, v);
}
}
_ => {}
}
}
@@ -635,3 +645,20 @@ pub async fn set_socks(value: config::Socks5Server) -> ResultType<()> {
.await?;
Ok(())
}
#[cfg(feature = "hwcodec")]
#[tokio::main]
pub async fn check_hwcodec_config() {
if let Some(config) = scrap::hwcodec::check_config() {
match connect(1000, "").await {
Ok(mut conn) => {
if conn.send(&Data::HwCodecConfig(Some(config))).await.is_err() {
log::error!("Failed to send hwcodec config by ipc");
}
}
Err(err) => {
log::info!("Failed to connect ipc: {:?}", err);
}
}
}
}

23
src/main.rs
View File

@@ -70,6 +70,15 @@ fn main() {
}
if args.is_empty() {
std::thread::spawn(move || start_server(false));
#[cfg(feature = "hwcodec")]
if let Ok(exe) = std::env::current_exe() {
std::thread::spawn(move || {
std::process::Command::new(exe)
.arg("--check-hwcodec-config")
.status()
.ok()
});
}
} else {
#[cfg(windows)]
{
@@ -109,6 +118,9 @@ fn main() {
} else if args[0] == "--extract" {
#[cfg(feature = "with_rc")]
hbb_common::allow_err!(crate::rc::extract_resources(&args[1]));
} else if args[0] == "--check-hwcodec-config" {
#[cfg(feature = "hwcodec")]
ipc::check_hwcodec_config();
return;
}
}
@@ -197,7 +209,7 @@ fn main() {
.about("RustDesk command line tool")
.args_from_usage(&args)
.get_matches();
use hbb_common::{env_logger::*, config::LocalConfig};
use hbb_common::{config::LocalConfig, env_logger::*};
init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
if let Some(p) = matches.value_of("port-forward") {
let options: Vec<String> = p.split(":").map(|x| x.to_owned()).collect();
@@ -225,6 +237,13 @@ fn main() {
}
let key = matches.value_of("key").unwrap_or("").to_owned();
let token = LocalConfig::get_option("access_token");
cli::start_one_port_forward(options[0].clone(), port, remote_host, remote_port, key, token);
cli::start_one_port_forward(
options[0].clone(),
port,
remote_host,
remote_port,
key,
token,
);
}
}

17
src/server/connection.rs
View File

@@ -396,6 +396,7 @@ impl Connection {
video_service::notify_video_frame_feched(id, None);
video_service::update_test_latency(id, 0);
video_service::update_image_quality(id, None);
scrap::codec::Encoder::update_video_encoder(id, scrap::codec::EncoderUpdate::Remove);
if let Err(err) = conn.try_port_forward_loop(&mut rx_from_cm).await {
conn.on_close(&err.to_string(), false);
}
@@ -780,6 +781,22 @@ impl Connection {
if let Some(message::Union::login_request(lr)) = msg.union {
if let Some(o) = lr.option.as_ref() {
self.update_option(o).await;
if let Some(q) = o.video_codec_state.clone().take() {
scrap::codec::Encoder::update_video_encoder(
self.inner.id(),
scrap::codec::EncoderUpdate::State(q),
);
} else {
scrap::codec::Encoder::update_video_encoder(
self.inner.id(),
scrap::codec::EncoderUpdate::DisableHwIfNotExist,
);
}
} else {
scrap::codec::Encoder::update_video_encoder(
self.inner.id(),
scrap::codec::EncoderUpdate::DisableHwIfNotExist,
);
}
self.video_ack_required = lr.video_ack_required;
if self.authorized {

99
src/server/video_service.rs
View File

@@ -23,7 +23,11 @@ use hbb_common::tokio::sync::{
mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
Mutex as TokioMutex,
};
use scrap::{Capturer, Config, Display, EncodeFrame, Encoder, Frame, VideoCodecId, STRIDE_ALIGN};
use scrap::{
codec::{Encoder, EncoderCfg, HwEncoderConfig},
vpxcodec::{VpxEncoderConfig, VpxVideoCodecId},
Capturer, Display, Frame,
};
use std::{
collections::HashSet,
io::{ErrorKind::WouldBlock, Result},
@@ -201,9 +205,11 @@ fn check_display_changed(
}
// Capturer object is expensive, avoiding to create it frequently.
fn create_capturer(privacy_mode_id: i32, display: Display) -> ResultType<Box<dyn TraitCapturer>> {
let use_yuv = true;
fn create_capturer(
privacy_mode_id: i32,
display: Display,
use_yuv: bool,
) -> ResultType<Box<dyn TraitCapturer>> {
#[cfg(not(windows))]
let c: Option<Box<dyn TraitCapturer>> = None;
#[cfg(windows)]
@@ -292,7 +298,7 @@ pub fn test_create_capturer(privacy_mode_id: i32, timeout_millis: u64) -> bool {
let test_begin = Instant::now();
while test_begin.elapsed().as_millis() < timeout_millis as _ {
if let Ok((_, _, display)) = get_current_display() {
if let Ok(_) = create_capturer(privacy_mode_id, display) {
if let Ok(_) = create_capturer(privacy_mode_id, display, true) {
return true;
}
}
@@ -336,6 +342,36 @@ fn run(sp: GenericService) -> ResultType<()> {
num_cpus::get(),
);
let q = get_image_quality();
let (bitrate, rc_min_quantizer, rc_max_quantizer, speed) = get_quality(width, height, q);
log::info!("bitrate={}, rc_min_quantizer={}", bitrate, rc_min_quantizer);
let encoder_cfg = match Encoder::current_hw_encoder_name() {
Some(codec_name) => EncoderCfg::HW(HwEncoderConfig {
codec_name,
width,
height,
bitrate_ratio: q >> 8,
}),
None => EncoderCfg::VPX(VpxEncoderConfig {
width: width as _,
height: height as _,
timebase: [1, 1000], // Output timestamp precision
bitrate,
codec: VpxVideoCodecId::VP9,
rc_min_quantizer,
rc_max_quantizer,
speed,
num_threads: (num_cpus::get() / 2) as _,
}),
};
let mut encoder;
match Encoder::new(encoder_cfg) {
Ok(x) => encoder = x,
Err(err) => bail!("Failed to create encoder: {}", err),
}
let privacy_mode_id = *PRIVACY_MODE_CONN_ID.lock().unwrap();
#[cfg(not(windows))]
let captuerer_privacy_mode_id = privacy_mode_id;
@@ -355,26 +391,7 @@ fn run(sp: GenericService) -> ResultType<()> {
} else {
log::info!("In privacy mode, the peer side cannot watch the screen");
}
let mut c = create_capturer(captuerer_privacy_mode_id, display)?;
let q = get_image_quality();
let (bitrate, rc_min_quantizer, rc_max_quantizer, speed) = get_quality(width, height, q);
log::info!("bitrate={}, rc_min_quantizer={}", bitrate, rc_min_quantizer);
let cfg = Config {
width: width as _,
height: height as _,
timebase: [1, 1000], // Output timestamp precision
bitrate,
codec: VideoCodecId::VP9,
rc_min_quantizer,
rc_max_quantizer,
speed,
};
let mut vpx;
match Encoder::new(&cfg, (num_cpus::get() / 2) as _) {
Ok(x) => vpx = x,
Err(err) => bail!("Failed to create encoder: {}", err),
}
let mut c = create_capturer(captuerer_privacy_mode_id, display, encoder.use_yuv())?;
if *SWITCH.lock().unwrap() {
log::debug!("Broadcasting display switch");
@@ -464,7 +481,7 @@ fn run(sp: GenericService) -> ResultType<()> {
Ok(frame) => {
let time = now - start;
let ms = (time.as_secs() * 1000 + time.subsec_millis() as u64) as i64;
let send_conn_ids = handle_one_frame(&sp, &frame, ms, &mut vpx)?;
let send_conn_ids = handle_one_frame(&sp, &frame, ms, &mut encoder)?;
frame_controller.set_send(now, send_conn_ids);
#[cfg(windows)]
{
@@ -531,7 +548,6 @@ fn run(sp: GenericService) -> ResultType<()> {
Ok(())
}
#[inline]
fn check_privacy_mode_changed(sp: &GenericService, privacy_mode_id: i32) -> ResultType<()> {
let privacy_mode_id_2 = *PRIVACY_MODE_CONN_ID.lock().unwrap();
if privacy_mode_id != privacy_mode_id_2 {
@@ -547,6 +563,7 @@ fn check_privacy_mode_changed(sp: &GenericService, privacy_mode_id: i32) -> Resu
}
#[inline]
#[cfg(any(target_os = "android", target_os = "ios"))]
fn create_msg(vp9s: Vec<VP9>) -> Message {
let mut msg_out = Message::new();
let mut vf = VideoFrame::new();
@@ -559,22 +576,12 @@ fn create_msg(vp9s: Vec<VP9>) -> Message {
msg_out
}
#[inline]
fn create_frame(frame: &EncodeFrame) -> VP9 {
VP9 {
data: frame.data.to_vec(),
key: frame.key,
pts: frame.pts,
..Default::default()
}
}
#[inline]
fn handle_one_frame(
sp: &GenericService,
frame: &[u8],
ms: i64,
vpx: &mut Encoder,
encoder: &mut Encoder,
) -> ResultType<HashSet<i32>> {
sp.snapshot(|sps| {
// so that new sub and old sub share the same encoder after switch
@@ -585,20 +592,8 @@ fn handle_one_frame(
})?;
let mut send_conn_ids: HashSet<i32> = Default::default();
let mut frames = Vec::new();
for ref frame in vpx
.encode(ms, frame, STRIDE_ALIGN)
.with_context(|| "Failed to encode")?
{
frames.push(create_frame(frame));
}
for ref frame in vpx.flush().with_context(|| "Failed to flush")? {
frames.push(create_frame(frame));
}
// to-do: flush periodically, e.g. 1 second
if frames.len() > 0 {
send_conn_ids = sp.send_video_frame(create_msg(frames));
if let Ok(msg) = encoder.encode_to_message(frame, ms) {
send_conn_ids = sp.send_video_frame(msg);
}
Ok(send_conn_ids)
}