mirror of
https://github.com/htrefil/rkvm.git
synced 2024-12-26 09:58:32 +01:00
Refactor server
This commit is contained in:
parent
437c3398f0
commit
b6c7ff6f85
7 changed files with 398 additions and 174 deletions
189
Cargo.lock
generated
189
Cargo.lock
generated
|
@ -92,6 +92,12 @@ version = "1.1.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.21.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a4a4ddaa51a5bc52a6948f74c06d20aaaddb71924eab79b8c97a8c556e942d6a"
|
||||
|
||||
[[package]]
|
||||
name = "bincode"
|
||||
version = "1.3.3"
|
||||
|
@ -131,6 +137,12 @@ version = "1.3.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d261e256854913907f67ed06efbc3338dfe6179796deefc1ff763fc1aee5535"
|
||||
|
||||
[[package]]
|
||||
name = "bytes"
|
||||
version = "1.4.0"
|
||||
|
@ -528,6 +540,15 @@ dependencies = [
|
|||
"windows-sys 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.61"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "445dde2150c55e483f3d8416706b97ec8e8237c307e5b7b4b8dd15e6af2a0730"
|
||||
dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.4.0"
|
||||
|
@ -775,6 +796,21 @@ version = "0.6.29"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1"
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.16.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3053cf52e236a3ed746dfc745aa9cacf1b791d846bdaf412f60a8d7d6e17c8fc"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"spin",
|
||||
"untrusted",
|
||||
"web-sys",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rkvm-certificate-gen"
|
||||
version = "0.1.0"
|
||||
|
@ -828,15 +864,17 @@ dependencies = [
|
|||
name = "rkvm-server"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"clap 4.2.2",
|
||||
"env_logger 0.8.4",
|
||||
"log",
|
||||
"rkvm-input",
|
||||
"rkvm-net",
|
||||
"rustls-pemfile",
|
||||
"serde",
|
||||
"slab",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tokio-rustls",
|
||||
"toml",
|
||||
]
|
||||
|
||||
|
@ -860,6 +898,37 @@ dependencies = [
|
|||
"windows-sys 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.21.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "07180898a28ed6a7f7ba2311594308f595e3dd2e3c3812fa0a80a47b45f17e5d"
|
||||
dependencies = [
|
||||
"log",
|
||||
"ring",
|
||||
"rustls-webpki",
|
||||
"sct",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-pemfile"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d194b56d58803a43635bdc398cd17e383d6f71f9182b9a192c127ca42494a59b"
|
||||
dependencies = [
|
||||
"base64",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-webpki"
|
||||
version = "0.100.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6207cd5ed3d8dca7816f8f3725513a34609c0c765bf652b8c3cb4cfd87db46b"
|
||||
dependencies = [
|
||||
"ring",
|
||||
"untrusted",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schannel"
|
||||
version = "0.1.21"
|
||||
|
@ -869,6 +938,16 @@ dependencies = [
|
|||
"windows-sys 0.42.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sct"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d53dcdb7c9f8158937a7981b48accfd39a43af418591a5d008c7b22b5e1b7ca4"
|
||||
dependencies = [
|
||||
"ring",
|
||||
"untrusted",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "2.8.2"
|
||||
|
@ -946,6 +1025,12 @@ dependencies = [
|
|||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spin"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
|
||||
|
||||
[[package]]
|
||||
name = "strsim"
|
||||
version = "0.8.0"
|
||||
|
@ -1035,6 +1120,26 @@ dependencies = [
|
|||
"unicode-width",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.40"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "978c9a314bd8dc99be594bc3c175faaa9794be04a5a5e153caba6915336cebac"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "1.0.40"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.15",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.27.0"
|
||||
|
@ -1074,6 +1179,16 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-rustls"
|
||||
version = "0.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e0d409377ff5b1e3ca6437aa86c1eb7d40c134bfec254e44c830defa92669db5"
|
||||
dependencies = [
|
||||
"rustls",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml"
|
||||
version = "0.5.11"
|
||||
|
@ -1101,6 +1216,12 @@ version = "0.1.10"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c0edd1e5b14653f783770bce4a4dabb4a5108a5370a5f5d8cfe8710c361f6c8b"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a"
|
||||
|
||||
[[package]]
|
||||
name = "utf8parse"
|
||||
version = "0.2.1"
|
||||
|
@ -1131,6 +1252,70 @@ version = "0.11.0+wasi-snapshot-preview1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.84"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31f8dcbc21f30d9b8f2ea926ecb58f6b91192c17e9d33594b3df58b2007ca53b"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.0",
|
||||
"wasm-bindgen-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-backend"
|
||||
version = "0.2.84"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "95ce90fd5bcc06af55a641a86428ee4229e44e07033963a2290a8e241607ccb9"
|
||||
dependencies = [
|
||||
"bumpalo",
|
||||
"log",
|
||||
"once_cell",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
"wasm-bindgen-shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro"
|
||||
version = "0.2.84"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4c21f77c0bedc37fd5dc21f897894a5ca01e7bb159884559461862ae90c0b4c5"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"wasm-bindgen-macro-support",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro-support"
|
||||
version = "0.2.84"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2aff81306fcac3c7515ad4e177f521b5c9a15f2b08f4e32d823066102f35a5f6"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
"wasm-bindgen-backend",
|
||||
"wasm-bindgen-shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-shared"
|
||||
version = "0.2.84"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0046fef7e28c3804e5e38bfa31ea2a0f73905319b677e57ebe37e49358989b5d"
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.61"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e33b99f4b23ba3eec1a53ac264e35a755f00e966e0065077d6027c0f575b0b97"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "which"
|
||||
version = "3.1.1"
|
||||
|
|
|
@ -1,2 +1,2 @@
|
|||
[workspace]
|
||||
members = ["rkvm-client", "rkvm-server", "rkvm-input", "rkvm-net", "rkvm-certificate-gen"]
|
||||
members = ["rkvm-client", "rkvm-server", "rkvm-input", "rkvm-net", "rkvm-certificate-gen"]
|
||||
|
|
|
@ -6,8 +6,7 @@ use std::time::Duration;
|
|||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
|
||||
// Is it bold to assume there won't be more than 65536 protocol versions?
|
||||
pub const PROTOCOL_VERSION: u16 = 1;
|
||||
pub const MESSAGE_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
pub const PROTOCOL_VERSION: u16 = 2;
|
||||
|
||||
pub async fn read_version<R>(mut reader: R) -> Result<u16, Error>
|
||||
where
|
||||
|
|
|
@ -15,9 +15,11 @@ serde = { version = "1.0.117", features = ["derive"] }
|
|||
toml = "0.5.7"
|
||||
log = "0.4.11"
|
||||
env_logger = "0.8.1"
|
||||
tokio-native-tls = "0.3.0"
|
||||
anyhow = "1.0.33"
|
||||
clap = { version = "4.2.2", features = ["derive"] }
|
||||
tokio-rustls = "0.24.0"
|
||||
rustls-pemfile = "1.0.2"
|
||||
thiserror = "1.0.40"
|
||||
slab = "0.4.8"
|
||||
|
||||
[package.metadata.rpm]
|
||||
package = "rkvm-server"
|
||||
|
|
|
@ -1,15 +1,13 @@
|
|||
use rkvm_input::Key;
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashSet;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub struct Config {
|
||||
pub listen_address: SocketAddr,
|
||||
pub switch_keys: HashSet<Key>,
|
||||
pub identity_path: PathBuf,
|
||||
#[serde(default)]
|
||||
pub identity_password: String,
|
||||
pub listen: SocketAddr,
|
||||
pub switch_key: Key,
|
||||
pub certificate: PathBuf,
|
||||
pub key: PathBuf,
|
||||
}
|
||||
|
|
|
@ -1,159 +1,22 @@
|
|||
mod config;
|
||||
mod tls;
|
||||
|
||||
use anyhow::{Context, Error};
|
||||
use config::Config;
|
||||
use rkvm_input::{Direction, Event, EventManager, Key, KeyKind};
|
||||
use log::LevelFilter;
|
||||
use rkvm_net::{self, Message, PROTOCOL_VERSION};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::convert::Infallible;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process;
|
||||
use clap::Parser;
|
||||
use config::Config;
|
||||
use log::LevelFilter;
|
||||
use rkvm_input::{Direction, Event, EventManager, Key, KeyKind};
|
||||
use rkvm_net::{self, Message};
|
||||
use slab::Slab;
|
||||
use std::io::{Error, ErrorKind};
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use std::process::ExitCode;
|
||||
use tokio::fs;
|
||||
use tokio::io::{AsyncRead, AsyncWrite};
|
||||
use tokio::io::{AsyncWriteExt, BufStream};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
|
||||
use tokio::time;
|
||||
use tokio_native_tls::native_tls::{Identity, TlsAcceptor};
|
||||
|
||||
async fn handle_connection<T>(
|
||||
mut stream: T,
|
||||
mut receiver: UnboundedReceiver<Event>,
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
T: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
rkvm_net::write_version(&mut stream, PROTOCOL_VERSION).await?;
|
||||
|
||||
let version = rkvm_net::read_version(&mut stream).await?;
|
||||
if version != PROTOCOL_VERSION {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Incompatible protocol version (got {}, expecting {})",
|
||||
version,
|
||||
PROTOCOL_VERSION
|
||||
));
|
||||
}
|
||||
|
||||
loop {
|
||||
// Send a keep alive message in intervals of half of the timeout just to be on the safe side.
|
||||
let message = match time::timeout(rkvm_net::MESSAGE_TIMEOUT / 2, receiver.recv()).await {
|
||||
Ok(Some(message)) => Message::Event(message),
|
||||
Ok(None) => return Ok(()),
|
||||
Err(_) => Message::KeepAlive,
|
||||
};
|
||||
|
||||
time::timeout(
|
||||
rkvm_net::MESSAGE_TIMEOUT,
|
||||
rkvm_net::write_message(&mut stream, &message),
|
||||
)
|
||||
.await
|
||||
.context("Write timeout")??;
|
||||
}
|
||||
}
|
||||
|
||||
async fn run(
|
||||
listen_address: SocketAddr,
|
||||
switch_keys: &HashSet<Key>,
|
||||
identity_path: &Path,
|
||||
identity_password: &str,
|
||||
) -> Result<Infallible, Error> {
|
||||
let identity = fs::read(identity_path)
|
||||
.await
|
||||
.context("Failed to read identity")?;
|
||||
let identity =
|
||||
Identity::from_pkcs12(&identity, identity_password).context("Failed to parse identity")?;
|
||||
let acceptor: tokio_native_tls::TlsAcceptor = TlsAcceptor::new(identity)
|
||||
.context("Failed to create TLS acceptor")
|
||||
.map(Into::into)?;
|
||||
let listener = TcpListener::bind(listen_address).await?;
|
||||
|
||||
log::info!("Listening on {}", listen_address);
|
||||
|
||||
let (client_sender, mut client_receiver) = mpsc::unbounded_channel();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (stream, address) = match listener.accept().await {
|
||||
Ok(sa) => sa,
|
||||
Err(err) => {
|
||||
let _ = client_sender.send(Err(err));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let stream = match acceptor.accept(stream).await {
|
||||
Ok(stream) => stream,
|
||||
Err(err) => {
|
||||
log::error!("{}: TLS error: {}", address, err);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let (sender, receiver) = mpsc::unbounded_channel();
|
||||
if client_sender.send(Ok(sender)).is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
log::info!("{}: connected", address);
|
||||
let message = handle_connection(stream, receiver)
|
||||
.await
|
||||
.err()
|
||||
.map(|err| format!(" ({})", err))
|
||||
.unwrap_or_else(String::new);
|
||||
log::info!("{}: disconnected{}", address, message);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
let mut clients: Vec<UnboundedSender<Event>> = Vec::new();
|
||||
let mut current = 0;
|
||||
let mut manager = EventManager::new().await?;
|
||||
let mut key_states: HashMap<_, _> = switch_keys
|
||||
.iter()
|
||||
.copied()
|
||||
.map(|key| (key, false))
|
||||
.collect();
|
||||
loop {
|
||||
tokio::select! {
|
||||
event = manager.read() => {
|
||||
let event = event?;
|
||||
if let Event::Key { direction, kind: KeyKind::Key(key) } = event {
|
||||
if let Some(state) = key_states.get_mut(&key) {
|
||||
*state = direction == Direction::Down;
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: This won't work with multiple keys.
|
||||
if key_states.iter().filter(|(_, state)| **state).count() == key_states.len() {
|
||||
for state in key_states.values_mut() {
|
||||
*state = false;
|
||||
}
|
||||
|
||||
current = (current + 1) % (clients.len() + 1);
|
||||
log::info!("Switching to client {}", current);
|
||||
continue;
|
||||
}
|
||||
|
||||
if current != 0 {
|
||||
let idx = current - 1;
|
||||
if clients[idx].send(event).is_ok() {
|
||||
continue;
|
||||
}
|
||||
|
||||
clients.remove(idx);
|
||||
current = 0;
|
||||
}
|
||||
|
||||
manager.write(event).await?;
|
||||
}
|
||||
sender = client_receiver.recv() => {
|
||||
clients.push(sender.unwrap()?);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
use tokio::signal;
|
||||
use tokio::sync::mpsc::{self, Sender};
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[structopt(name = "rkvm-server", about = "The rkvm server application")]
|
||||
|
@ -163,43 +26,138 @@ struct Args {
|
|||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
async fn main() -> ExitCode {
|
||||
env_logger::builder()
|
||||
.format_timestamp(None)
|
||||
.filter(None, LevelFilter::Info)
|
||||
.parse_default_env()
|
||||
.init();
|
||||
|
||||
let args = Args::parse();
|
||||
let config = match fs::read_to_string(&args.config_path).await {
|
||||
Ok(config) => config,
|
||||
Err(err) => {
|
||||
log::error!("Error loading config: {}", err);
|
||||
process::exit(1);
|
||||
log::error!("Error reading config: {}", err);
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
};
|
||||
|
||||
let config: Config = match toml::from_str(&config) {
|
||||
let config = match toml::from_str::<Config>(&config) {
|
||||
Ok(config) => config,
|
||||
Err(err) => {
|
||||
log::error!("Error parsing config: {}", err);
|
||||
process::exit(1);
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
};
|
||||
|
||||
let acceptor = match tls::configure(&config.certificate, &config.key).await {
|
||||
Ok(acceptor) => acceptor,
|
||||
Err(err) => {
|
||||
log::error!("Error configuring TLS: {}", err);
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
result = run(config.listen_address, &config.switch_keys, &config.identity_path, &config.identity_password) => {
|
||||
result = run(config.listen, acceptor, config.switch_key) => {
|
||||
if let Err(err) = result {
|
||||
log::error!("Error: {:#}", err);
|
||||
process::exit(1);
|
||||
log::error!("Error running server: {}", err);
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
}
|
||||
result = tokio::signal::ctrl_c() => {
|
||||
// This is needed to properly clean evdev stuff up.
|
||||
result = signal::ctrl_c() => {
|
||||
if let Err(err) = result {
|
||||
log::error!("Error setting up signal handler: {}", err);
|
||||
process::exit(1);
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
|
||||
log::info!("Exiting on signal");
|
||||
}
|
||||
}
|
||||
|
||||
ExitCode::SUCCESS
|
||||
}
|
||||
|
||||
async fn run(listen: SocketAddr, acceptor: TlsAcceptor, switch_key: Key) -> Result<(), Error> {
|
||||
let listener = TcpListener::bind(&listen).await?;
|
||||
log::info!("Listening on {}", listen);
|
||||
|
||||
let mut clients = Slab::<Sender<_>>::new();
|
||||
let mut current = 0;
|
||||
let mut manager = EventManager::new().await?;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = listener.accept() => {
|
||||
// Remove dead clients.
|
||||
clients.retain(|_, client| !client.is_closed());
|
||||
if !clients.contains(current) {
|
||||
current = 0;
|
||||
}
|
||||
|
||||
let (stream, addr) = result?;
|
||||
let acceptor = acceptor.clone();
|
||||
|
||||
let (sender, mut receiver) = mpsc::channel::<Event>(1);
|
||||
clients.insert(sender);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let stream = match acceptor.accept(stream).await {
|
||||
Ok(stream) => stream,
|
||||
Err(err) => {
|
||||
log::error!("{}: TLS accept error: {}", addr, err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
log::info!("{}: Connected", addr);
|
||||
|
||||
let result = async {
|
||||
let mut stream = BufStream::with_capacity(1024, 1024, stream);
|
||||
|
||||
rkvm_net::write_version(&mut stream, rkvm_net::PROTOCOL_VERSION).await?;
|
||||
stream.flush().await?;
|
||||
|
||||
let version = rkvm_net::read_version(&mut stream).await?;
|
||||
if version != rkvm_net::PROTOCOL_VERSION {
|
||||
return Err(Error::new(ErrorKind::InvalidData, "Invalid client protocol version"));
|
||||
}
|
||||
|
||||
loop {
|
||||
let event = match receiver.recv().await {
|
||||
Some(event) => event,
|
||||
None => break,
|
||||
};
|
||||
|
||||
rkvm_net::write_message(&mut stream, &Message::Event(event)).await?;
|
||||
stream.flush().await?;
|
||||
}
|
||||
|
||||
Ok::<_, Error>(())
|
||||
}
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(()) => log::info!("{}: Disconnected", addr),
|
||||
Err(err) => log::error!("{}: Disconnected: {}", addr, err),
|
||||
}
|
||||
});
|
||||
}
|
||||
result = manager.read() => {
|
||||
let event = result?;
|
||||
if let Event::Key { direction: Direction::Down, kind: KeyKind::Key(key) } = event {
|
||||
if key == switch_key {
|
||||
current = (current + 1) % (clients.len() + 1);
|
||||
log::info!("Switching to client {}", current);
|
||||
}
|
||||
}
|
||||
|
||||
if current != 0 && clients[current].send(event).await.is_err() {
|
||||
current = 0;
|
||||
manager.write(event).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
82
rkvm-server/src/tls.rs
Normal file
82
rkvm-server/src/tls.rs
Normal file
|
@ -0,0 +1,82 @@
|
|||
use rustls_pemfile::Item;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::{io, iter};
|
||||
use thiserror::Error;
|
||||
use tokio::fs;
|
||||
use tokio_rustls::rustls::{self, Certificate, PrivateKey, ServerConfig};
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum Error {
|
||||
#[error(transparent)]
|
||||
Rustls(#[from] rustls::Error),
|
||||
#[error(transparent)]
|
||||
Io(#[from] io::Error),
|
||||
#[error("Multiple private keys provided")]
|
||||
MultipleKeys,
|
||||
#[error("No suitable private keys provided")]
|
||||
NoKeys,
|
||||
}
|
||||
|
||||
pub async fn configure(certificate: &Path, key: &Path) -> Result<TlsAcceptor, Error> {
|
||||
enum LoadedItem {
|
||||
Certificate(Vec<u8>),
|
||||
Key(Vec<u8>),
|
||||
}
|
||||
|
||||
let certificate = fs::read_to_string(certificate).await?;
|
||||
let key = fs::read_to_string(key).await?;
|
||||
|
||||
let certificates_iter = iter::from_fn({
|
||||
let mut buffer = certificate.as_bytes();
|
||||
|
||||
move || rustls_pemfile::read_one(&mut buffer).transpose()
|
||||
})
|
||||
.filter_map(|item| match item {
|
||||
Ok(Item::X509Certificate(data)) => Some(Ok(LoadedItem::Certificate(data))),
|
||||
Err(err) => Some(Err(err)),
|
||||
_ => None,
|
||||
});
|
||||
|
||||
let keys_iter = iter::from_fn({
|
||||
let mut buffer = key.as_bytes();
|
||||
|
||||
move || rustls_pemfile::read_one(&mut buffer).transpose()
|
||||
})
|
||||
.filter_map(|item| match item {
|
||||
Ok(Item::RSAKey(data)) | Ok(Item::PKCS8Key(data)) | Ok(Item::ECKey(data)) => {
|
||||
Some(Ok(LoadedItem::Key(data)))
|
||||
}
|
||||
Err(err) => Some(Err(err)),
|
||||
_ => None,
|
||||
});
|
||||
|
||||
let mut certificates = Vec::new();
|
||||
let mut key = None;
|
||||
|
||||
for item in certificates_iter.chain(keys_iter) {
|
||||
let item = item?;
|
||||
|
||||
match item {
|
||||
LoadedItem::Certificate(data) => certificates.push(Certificate(data)),
|
||||
LoadedItem::Key(data) => {
|
||||
if key.is_some() {
|
||||
return Err(Error::MultipleKeys);
|
||||
}
|
||||
|
||||
key = Some(PrivateKey(data));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let key = key.ok_or(Error::NoKeys)?;
|
||||
|
||||
ServerConfig::builder()
|
||||
.with_safe_defaults()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(certificates, key)
|
||||
.map(Arc::new)
|
||||
.map(Into::into)
|
||||
.map_err(Into::into)
|
||||
}
|
Loading…
Reference in a new issue