From 6ded4b89641f6e05f0ef907d8dfe1cce31b95f77 Mon Sep 17 00:00:00 2001 From: Zac Bergquist Date: Wed, 30 Oct 2024 13:58:04 -0600 Subject: [PATCH] Store RDP licenses in a LRU cache instead of a HashMap This ensures memory doesn't grow unbounded if we end up caching a large number of licenses. --- Cargo.lock | 41 ++++++++++++++- lib/srv/desktop/rdp/rdpclient/Cargo.toml | 1 + lib/srv/desktop/rdp/rdpclient/src/lib.rs | 65 ++++++++++++++++-------- 3 files changed, 86 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2201695e591d9..0261bb969c4ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f" + [[package]] name = "asn1-rs" version = "0.5.1" @@ -452,6 +458,12 @@ dependencies = [ "termcolor", ] +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + [[package]] name = "errno" version = "0.3.1" @@ -504,6 +516,12 @@ dependencies = [ "miniz_oxide 0.6.2", ] +[[package]] +name = "foldhash" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f81ec6369c545a7d40e4589b5597581fa1c441fe1cce96dd1de43159910a36a2" + [[package]] name = "foreign-types" version = "0.5.0" @@ -609,6 +627,17 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +[[package]] +name = "hashbrown" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + [[package]] name = "heapless" version = "0.7.16" @@ -675,7 +704,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1885e79c1fc4b10f0e172c475f458b7f7b93061064d98c3293e98c5ba0c8b399" dependencies = [ "autocfg", - "hashbrown", + "hashbrown 0.12.3", ] [[package]] @@ -809,6 +838,15 @@ version = "0.4.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.0", +] + [[package]] name = "md-5" version = "0.8.0" @@ -1208,6 +1246,7 @@ dependencies = [ "iso7816-tlv", "libc", "log", + "lru", "num-derive", "num-traits", "png", diff --git a/lib/srv/desktop/rdp/rdpclient/Cargo.toml b/lib/srv/desktop/rdp/rdpclient/Cargo.toml index e5c6fe042c832..d7f696f1be9f7 100644 --- a/lib/srv/desktop/rdp/rdpclient/Cargo.toml +++ b/lib/srv/desktop/rdp/rdpclient/Cargo.toml @@ -27,6 +27,7 @@ rdp-rs = { git = "https://github.com/gravitational/rdp-rs", rev = "2b0d99cc60c7b uuid = { version = "1.4.1", features = ["v4"] } utf16string = "0.2.0" png = "0.17.10" +lru = "0.12.5" [build-dependencies] cbindgen = "0.25.0" diff --git a/lib/srv/desktop/rdp/rdpclient/src/lib.rs b/lib/srv/desktop/rdp/rdpclient/src/lib.rs index 4f4b4b110509d..b39cba8b7e449 100644 --- a/lib/srv/desktop/rdp/rdpclient/src/lib.rs +++ b/lib/srv/desktop/rdp/rdpclient/src/lib.rs @@ -50,13 +50,13 @@ extern crate log; extern crate num_derive; use errors::try_error; +use lru::LruCache; use rand::Rng; use rand::SeedableRng; use rdp::core::event::*; use rdp::core::gcc::KeyboardLayout; use rdp::core::global; use rdp::core::global::ServerError; -use rdp::core::license::MemoryLicenseStore; use rdp::core::mcs; use rdp::core::sec; use rdp::core::tpkt; @@ -74,6 +74,7 @@ use std::io::ErrorKind; use std::io::{Cursor, Read, Write}; use std::net; use std::net::{TcpStream, ToSocketAddrs}; +use std::num::NonZeroUsize; use std::os::raw::c_char; use std::os::unix::io::AsRawFd; use std::sync::OnceLock; @@ -260,7 +261,7 @@ struct ConnectParams { show_desktop_wallpaper: bool, } -static LICENSE_STORE: OnceLock> = OnceLock::new(); +static LICENSE_STORE: OnceLock = OnceLock::new(); fn connect_rdp_inner(go_ref: usize, params: ConnectParams) -> Result { // Connect and authenticate. @@ -312,8 +313,7 @@ fn connect_rdp_inner(go_ref: usize, params: ConnectParams) -> Result { - inner: Mutex, +const MAX_SAVED_LICENSES: usize = 256; + +/// LruLicenseStore stores licenses in an in-memory LRU cache. +struct LruLicenseStore { + licenses: Mutex>>, } -impl SyncLicenseStore { - pub fn new(license_store: L) -> Self { +impl LruLicenseStore { + pub fn new() -> Self { Self { - inner: Mutex::new(license_store), + licenses: Mutex::new(LruCache::new( + NonZeroUsize::new(MAX_SAVED_LICENSES).unwrap(), + )), } } } -impl LicenseStore for &SyncLicenseStore { +impl LicenseStore for &LruLicenseStore { fn write_license( &mut self, major: u16, @@ -2089,10 +2091,16 @@ impl LicenseStore for &SyncLicenseStore { license: &[u8], ) { info!("Saving {major}.{minor} license from {issuer}"); - self.inner - .lock() - .unwrap() - .write_license(major, minor, company, issuer, product_id, license); + self.licenses.lock().unwrap().put( + LicenseStoreKey { + major, + minor, + company: company.to_owned(), + issuer: issuer.to_owned(), + product_id: product_id.to_owned(), + }, + license.to_vec(), + ); } fn read_license( @@ -2104,20 +2112,37 @@ impl LicenseStore for &SyncLicenseStore { product_id: &str, ) -> Option> { let license = self - .inner + .licenses .lock() .unwrap() - .read_license(major, minor, company, issuer, product_id); + .get(&LicenseStoreKey { + major, + minor, + company: company.to_owned(), + issuer: issuer.to_owned(), + product_id: product_id.to_owned(), + }) + .cloned(); + if license.is_some() { info!("Found existing {major}.{minor} license from {issuer}"); } else { info!("No existing {major}.{minor} license from {issuer}"); } - license + return license; } } +#[derive(PartialEq, Eq, Hash)] +struct LicenseStoreKey { + major: u16, + minor: u16, + company: String, + issuer: String, + product_id: String, +} + // These functions are defined on the Go side. Look for functions with '//export funcname' // comments. extern "C" {