|
| 1 | +#![allow(unused_crate_dependencies)] // False positives because there are both a library and a binary. |
| 2 | + |
| 3 | +use std::{io::Write, str::FromStr}; |
| 4 | + |
| 5 | +use anyhow::Context; |
| 6 | +use ironrdp_pdu::rdp::capability_sets::{CmdFlags, EntropyBits}; |
| 7 | +use ironrdp_server::{ |
| 8 | + bench::encoder::UpdateEncoder, BitmapUpdate, DesktopSize, DisplayUpdate, PixelFormat, PixelOrder, |
| 9 | + RdpServerDisplayUpdates, |
| 10 | +}; |
| 11 | +use tokio::{fs::File, io::AsyncReadExt}; |
| 12 | + |
| 13 | +#[tokio::main(flavor = "current_thread")] |
| 14 | +async fn main() -> Result<(), anyhow::Error> { |
| 15 | + setup_logging()?; |
| 16 | + let mut pargs = pico_args::Arguments::from_env(); |
| 17 | + |
| 18 | + let width = pargs.opt_value_from_str("--width")?.unwrap_or(3840); |
| 19 | + let height = pargs.opt_value_from_str("--height")?.unwrap_or(2400); |
| 20 | + let codec = pargs.opt_value_from_str("--codec")?.unwrap_or(OptCodec::default()); |
| 21 | + |
| 22 | + let filename: String = pargs.free_from_str().context("missing RGB input filename")?; |
| 23 | + let file = File::open(&filename) |
| 24 | + .await |
| 25 | + .with_context(|| format!("Failed to open file: {}", filename))?; |
| 26 | + |
| 27 | + let mut flags = CmdFlags::all(); |
| 28 | + |
| 29 | + #[allow(unused)] |
| 30 | + let (remotefx, qoicodec) = match codec { |
| 31 | + OptCodec::RemoteFX => (Some((EntropyBits::Rlgr3, 0)), None::<u8>), |
| 32 | + OptCodec::Bitmap => { |
| 33 | + flags -= CmdFlags::SET_SURFACE_BITS; |
| 34 | + (None, None) |
| 35 | + } |
| 36 | + OptCodec::None => (None, None), |
| 37 | + #[cfg(feature = "qoi")] |
| 38 | + OptCodec::QOI => (None, Some(0)), |
| 39 | + }; |
| 40 | + let mut encoder = UpdateEncoder::new( |
| 41 | + flags, |
| 42 | + remotefx, |
| 43 | + #[cfg(feature = "qoi")] |
| 44 | + qoicodec, |
| 45 | + ); |
| 46 | + |
| 47 | + let mut total_raw = 0u64; |
| 48 | + let mut total_enc = 0u64; |
| 49 | + let mut n_updates = 0u64; |
| 50 | + let mut updates = DisplayUpdates::new(file, DesktopSize { width, height }); |
| 51 | + while let Some(DisplayUpdate::Bitmap(up)) = updates.next_update().await { |
| 52 | + total_raw += up.data.len() as u64; |
| 53 | + let frag = encoder.bitmap(up)?; |
| 54 | + let len = frag.data.len() as u64; |
| 55 | + total_enc += len; |
| 56 | + n_updates += 1; |
| 57 | + print!("."); |
| 58 | + std::io::stdout().flush().unwrap(); |
| 59 | + } |
| 60 | + println!(); |
| 61 | + |
| 62 | + let ratio = total_enc as f64 / total_raw as f64; |
| 63 | + let percent = 100.0 - ratio * 100.0; |
| 64 | + println!("Encoder: {:?}", encoder); |
| 65 | + println!("Nb updates: {:?}", n_updates); |
| 66 | + println!( |
| 67 | + "Sum of bytes: {}/{} ({:.2}%)", |
| 68 | + bytesize::ByteSize(total_enc), |
| 69 | + bytesize::ByteSize(total_raw), |
| 70 | + percent, |
| 71 | + ); |
| 72 | + Ok(()) |
| 73 | +} |
| 74 | + |
| 75 | +struct DisplayUpdates { |
| 76 | + file: File, |
| 77 | + desktop_size: DesktopSize, |
| 78 | +} |
| 79 | + |
| 80 | +impl DisplayUpdates { |
| 81 | + fn new(file: File, desktop_size: DesktopSize) -> Self { |
| 82 | + Self { file, desktop_size } |
| 83 | + } |
| 84 | +} |
| 85 | + |
| 86 | +#[async_trait::async_trait] |
| 87 | +impl RdpServerDisplayUpdates for DisplayUpdates { |
| 88 | + async fn next_update(&mut self) -> Option<DisplayUpdate> { |
| 89 | + let stride = self.desktop_size.width as usize * 4; |
| 90 | + let frame_size = stride * self.desktop_size.height as usize; |
| 91 | + let mut buf = vec![0u8; frame_size]; |
| 92 | + if self.file.read_exact(&mut buf).await.is_err() { |
| 93 | + return None; |
| 94 | + } |
| 95 | + |
| 96 | + let up = DisplayUpdate::Bitmap(BitmapUpdate { |
| 97 | + top: 0, |
| 98 | + left: 0, |
| 99 | + width: self.desktop_size.width.try_into().unwrap(), |
| 100 | + height: self.desktop_size.height.try_into().unwrap(), |
| 101 | + format: PixelFormat::RgbX32, |
| 102 | + order: PixelOrder::TopToBottom, |
| 103 | + data: buf, |
| 104 | + stride, |
| 105 | + }); |
| 106 | + Some(up) |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +fn setup_logging() -> anyhow::Result<()> { |
| 111 | + use tracing::metadata::LevelFilter; |
| 112 | + use tracing_subscriber::prelude::*; |
| 113 | + use tracing_subscriber::EnvFilter; |
| 114 | + |
| 115 | + let fmt_layer = tracing_subscriber::fmt::layer().compact(); |
| 116 | + |
| 117 | + let env_filter = EnvFilter::builder() |
| 118 | + .with_default_directive(LevelFilter::WARN.into()) |
| 119 | + .with_env_var("IRONRDP_LOG") |
| 120 | + .from_env_lossy(); |
| 121 | + |
| 122 | + tracing_subscriber::registry() |
| 123 | + .with(fmt_layer) |
| 124 | + .with(env_filter) |
| 125 | + .try_init() |
| 126 | + .context("failed to set tracing global subscriber")?; |
| 127 | + |
| 128 | + Ok(()) |
| 129 | +} |
| 130 | + |
| 131 | +enum OptCodec { |
| 132 | + #[cfg(feature = "qoi")] |
| 133 | + QOI, |
| 134 | + RemoteFX, |
| 135 | + Bitmap, |
| 136 | + None, |
| 137 | +} |
| 138 | + |
| 139 | +impl Default for OptCodec { |
| 140 | + fn default() -> Self { |
| 141 | + #[cfg(feature = "qoi")] |
| 142 | + { |
| 143 | + Self::QOI |
| 144 | + } |
| 145 | + |
| 146 | + #[cfg(not(feature = "qoi"))] |
| 147 | + { |
| 148 | + Self::RemoteFX |
| 149 | + } |
| 150 | + } |
| 151 | +} |
| 152 | + |
| 153 | +impl FromStr for OptCodec { |
| 154 | + type Err = anyhow::Error; |
| 155 | + |
| 156 | + fn from_str(s: &str) -> Result<Self, Self::Err> { |
| 157 | + match s { |
| 158 | + #[cfg(feature = "qoi")] |
| 159 | + "qoi" => Ok(Self::QOI), |
| 160 | + "remotefx" => Ok(Self::RemoteFX), |
| 161 | + "bitmap" => Ok(Self::Bitmap), |
| 162 | + "none" => Ok(Self::None), |
| 163 | + _ => Err(anyhow::anyhow!("unknown codec: {}", s)), |
| 164 | + } |
| 165 | + } |
| 166 | +} |
0 commit comments