Denis Drakhnia
11 months ago
9 changed files with 707 additions and 10 deletions
@ -0,0 +1,14 @@
@@ -0,0 +1,14 @@
|
||||
[package] |
||||
name = "xash3d-query" |
||||
version = "0.1.0" |
||||
license = "GPL-3.0-only" |
||||
authors = ["Denis Drakhnia <numas13@gmail.com>"] |
||||
edition = "2021" |
||||
rust-version = "1.56" |
||||
|
||||
[dependencies] |
||||
thiserror = "1.0.49" |
||||
getopts = "0.2.21" |
||||
serde = { version = "1.0.188", features = ["derive"] } |
||||
serde_json = "1.0.107" |
||||
xash3d-protocol = { path = "../protocol", version = "0.1.0" } |
||||
@ -0,0 +1,166 @@
@@ -0,0 +1,166 @@
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
// SPDX-FileCopyrightText: 2023 Denis Drakhnia <numas13@gmail.com>
|
||||
|
||||
use std::process; |
||||
|
||||
use getopts::Options; |
||||
|
||||
const BIN_NAME: &str = env!("CARGO_BIN_NAME"); |
||||
const PKG_NAME: &str = env!("CARGO_PKG_NAME"); |
||||
const PKG_VERSION: &str = env!("CARGO_PKG_VERSION"); |
||||
|
||||
const DEFAULT_HOST: &str = "mentality.rip"; |
||||
const DEFAULT_PORT: u16 = 27010; |
||||
|
||||
#[derive(Debug)] |
||||
pub struct Cli { |
||||
pub masters: Vec<Box<str>>, |
||||
pub args: Vec<String>, |
||||
pub master_timeout: u32, |
||||
pub server_timeout: u32, |
||||
pub protocol: Vec<u8>, |
||||
pub json: bool, |
||||
pub debug: bool, |
||||
} |
||||
|
||||
impl Default for Cli { |
||||
fn default() -> Cli { |
||||
Cli { |
||||
masters: vec![ |
||||
format!("{}:{}", DEFAULT_HOST, DEFAULT_PORT).into_boxed_str(), |
||||
//format!("{}:{}", DEFAULT_HOST, DEFAULT_PORT + 1).into_boxed_str(),
|
||||
], |
||||
args: Default::default(), |
||||
master_timeout: 2, |
||||
server_timeout: 2, |
||||
protocol: vec![xash3d_protocol::VERSION, xash3d_protocol::VERSION - 1], |
||||
json: false, |
||||
debug: false, |
||||
} |
||||
} |
||||
} |
||||
|
||||
fn print_usage(opts: Options) { |
||||
let brief = format!( |
||||
"\ |
||||
Usage: {} [options] <COMMAND> [ARGS] |
||||
|
||||
COMMANDS: |
||||
all fetch servers from all masters and fetch info for each server |
||||
info hosts... fetch info for each server |
||||
list fetch servers from all masters and print server addresses\ |
||||
", |
||||
BIN_NAME |
||||
); |
||||
print!("{}", opts.usage(&brief)); |
||||
} |
||||
|
||||
fn print_version() { |
||||
println!("{} v{}", PKG_NAME, PKG_VERSION); |
||||
} |
||||
|
||||
pub fn parse() -> Cli { |
||||
let mut cli = Cli::default(); |
||||
|
||||
let args: Vec<_> = std::env::args().collect(); |
||||
let mut opts = Options::new(); |
||||
opts.optflag("h", "help", "print usage help"); |
||||
opts.optflag("v", "version", "print program version"); |
||||
let help = format!( |
||||
"master address to connect [default: {}]", |
||||
cli.masters.join(",") |
||||
); |
||||
opts.optopt("m", "master", &help, "LIST"); |
||||
let help = format!( |
||||
"time to wait results from masters [default: {}]", |
||||
cli.master_timeout |
||||
); |
||||
opts.optopt("T", "master-timeout", &help, "SECONDS"); |
||||
let help = format!( |
||||
"time to wait results from servers [default: {}]", |
||||
cli.server_timeout |
||||
); |
||||
opts.optopt("t", "server-timeout", &help, "SECONDS"); |
||||
let protocols = cli |
||||
.protocol |
||||
.iter() |
||||
.map(|&i| format!("{}", i)) |
||||
.collect::<Vec<_>>() |
||||
.join(","); |
||||
let help = format!("protocol version [default: {}]", protocols); |
||||
opts.optopt("p", "protocol", &help, "VERSION"); |
||||
opts.optflag("j", "json", "output JSON"); |
||||
opts.optflag("d", "debug", "output debug"); |
||||
|
||||
let matches = match opts.parse(&args[1..]) { |
||||
Ok(m) => m, |
||||
Err(e) => { |
||||
eprintln!("{}", e); |
||||
process::exit(1); |
||||
} |
||||
}; |
||||
|
||||
if matches.opt_present("help") { |
||||
print_usage(opts); |
||||
process::exit(0); |
||||
} |
||||
|
||||
if matches.opt_present("version") { |
||||
print_version(); |
||||
process::exit(0); |
||||
} |
||||
|
||||
if let Some(s) = matches.opt_str("master") { |
||||
cli.masters.clear(); |
||||
|
||||
for mut i in s.split(',').map(String::from) { |
||||
if !i.contains(':') { |
||||
i.push_str(":27010"); |
||||
} |
||||
cli.masters.push(i.into_boxed_str()); |
||||
} |
||||
} |
||||
|
||||
match matches.opt_get("master") { |
||||
Ok(Some(t)) => cli.master_timeout = t, |
||||
Ok(None) => {} |
||||
Err(_) => { |
||||
eprintln!("Invalid master-timeout"); |
||||
process::exit(1); |
||||
} |
||||
} |
||||
|
||||
match matches.opt_get("server-timeout") { |
||||
Ok(Some(t)) => cli.server_timeout = t, |
||||
Ok(None) => {} |
||||
Err(_) => { |
||||
eprintln!("Invalid server-timeout"); |
||||
process::exit(1); |
||||
} |
||||
} |
||||
|
||||
if let Some(s) = matches.opt_str("protocol") { |
||||
cli.protocol.clear(); |
||||
|
||||
let mut error = false; |
||||
for i in s.split(',') { |
||||
match i.parse() { |
||||
Ok(i) => cli.protocol.push(i), |
||||
Err(_) => { |
||||
eprintln!("Invalid protocol version: {}", i); |
||||
error = true; |
||||
} |
||||
} |
||||
} |
||||
|
||||
if error { |
||||
process::exit(1); |
||||
} |
||||
} |
||||
|
||||
cli.json = matches.opt_present("json"); |
||||
cli.debug = matches.opt_present("debug"); |
||||
cli.args = matches.free; |
||||
|
||||
cli |
||||
} |
||||
@ -0,0 +1,447 @@
@@ -0,0 +1,447 @@
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
// SPDX-FileCopyrightText: 2023 Denis Drakhnia <numas13@gmail.com>
|
||||
|
||||
mod cli; |
||||
|
||||
use std::cmp; |
||||
use std::collections::{HashMap, HashSet}; |
||||
use std::fmt; |
||||
use std::io; |
||||
use std::net::{Ipv4Addr, SocketAddrV4, UdpSocket}; |
||||
use std::process; |
||||
use std::sync::mpsc; |
||||
use std::thread; |
||||
use std::time::{Duration, Instant}; |
||||
|
||||
use serde::Serialize; |
||||
use thiserror::Error; |
||||
use xash3d_protocol::types::Str; |
||||
use xash3d_protocol::{filter, game, master, server, Error as ProtocolError}; |
||||
|
||||
use crate::cli::Cli; |
||||
|
||||
#[derive(Error, Debug)] |
||||
enum Error { |
||||
#[error("Undefined command")] |
||||
UndefinedCommand, |
||||
#[error(transparent)] |
||||
Protocol(#[from] ProtocolError), |
||||
#[error(transparent)] |
||||
Io(#[from] io::Error), |
||||
} |
||||
|
||||
#[derive(Clone, Debug, Serialize)] |
||||
#[serde(tag = "type")] |
||||
enum ServerResultKind { |
||||
#[serde(rename = "ok")] |
||||
Ok { info: ServerInfo }, |
||||
#[serde(rename = "error")] |
||||
Error { message: String }, |
||||
#[serde(rename = "invalid")] |
||||
Invalid { message: String, response: String }, |
||||
#[serde(rename = "timeout")] |
||||
Timeout, |
||||
#[serde(rename = "protocol")] |
||||
Protocol, |
||||
} |
||||
|
||||
#[derive(Clone, Debug, Serialize)] |
||||
struct ServerResult { |
||||
address: String, |
||||
#[serde(flatten)] |
||||
kind: ServerResultKind, |
||||
} |
||||
|
||||
impl ServerResult { |
||||
fn new(address: String, kind: ServerResultKind) -> Self { |
||||
Self { |
||||
address: address.to_string(), |
||||
kind, |
||||
} |
||||
} |
||||
|
||||
fn ok<T: Into<ServerInfo>>(address: String, info: T) -> Self { |
||||
Self::new(address, ServerResultKind::Ok { info: info.into() }) |
||||
} |
||||
|
||||
fn timeout(address: String) -> Self { |
||||
Self::new(address, ServerResultKind::Timeout) |
||||
} |
||||
|
||||
fn protocol(address: String) -> Self { |
||||
Self::new(address, ServerResultKind::Protocol) |
||||
} |
||||
|
||||
fn error<T>(address: String, message: T) -> Self |
||||
where |
||||
T: fmt::Display, |
||||
{ |
||||
Self::new( |
||||
address, |
||||
ServerResultKind::Error { |
||||
message: message.to_string(), |
||||
}, |
||||
) |
||||
} |
||||
|
||||
fn invalid<T>(address: String, message: T, response: &[u8]) -> Self |
||||
where |
||||
T: fmt::Display, |
||||
{ |
||||
Self::new( |
||||
address, |
||||
ServerResultKind::Invalid { |
||||
message: message.to_string(), |
||||
response: Str(response).to_string(), |
||||
}, |
||||
) |
||||
} |
||||
} |
||||
|
||||
#[derive(Clone, Debug, Serialize)] |
||||
struct ServerInfo { |
||||
pub gamedir: String, |
||||
pub map: String, |
||||
pub host: String, |
||||
pub protocol: u8, |
||||
pub numcl: u8, |
||||
pub maxcl: u8, |
||||
pub dm: bool, |
||||
pub team: bool, |
||||
pub coop: bool, |
||||
pub password: bool, |
||||
} |
||||
|
||||
impl From<server::GetServerInfoResponse<&str>> for ServerInfo { |
||||
fn from(other: server::GetServerInfoResponse<&str>) -> Self { |
||||
Self { |
||||
gamedir: other.gamedir.to_owned(), |
||||
map: other.map.to_owned(), |
||||
host: other.host.to_owned(), |
||||
protocol: other.protocol, |
||||
numcl: other.numcl, |
||||
maxcl: other.maxcl, |
||||
dm: other.dm, |
||||
team: other.team, |
||||
coop: other.coop, |
||||
password: other.password, |
||||
} |
||||
} |
||||
} |
||||
|
||||
#[derive(Clone, Debug, Serialize)] |
||||
struct InfoResult<'a> { |
||||
protocol: &'a [u8], |
||||
master_timeout: u32, |
||||
server_timeout: u32, |
||||
masters: &'a [Box<str>], |
||||
servers: &'a [&'a ServerResult], |
||||
} |
||||
|
||||
#[derive(Clone, Debug, Serialize)] |
||||
struct ListResult<'a> { |
||||
master_timeout: u32, |
||||
masters: &'a [Box<str>], |
||||
servers: &'a [&'a str], |
||||
} |
||||
|
||||
enum Message { |
||||
Servers(Vec<String>), |
||||
ServerResult(ServerResult), |
||||
End, |
||||
} |
||||
|
||||
fn cmp_address(a: &str, b: &str) -> cmp::Ordering { |
||||
match (a.parse::<SocketAddrV4>(), b.parse::<SocketAddrV4>()) { |
||||
(Ok(a), Ok(b)) => a.cmp(&b), |
||||
_ => a.cmp(b), |
||||
} |
||||
} |
||||
|
||||
fn query_servers(host: &str, timeout: Duration, tx: &mpsc::Sender<Message>) -> Result<(), Error> { |
||||
let sock = UdpSocket::bind("0.0.0.0:0")?; |
||||
sock.connect(host)?; |
||||
|
||||
let mut buf = [0; 512]; |
||||
let p = game::QueryServers { |
||||
region: server::Region::RestOfTheWorld, |
||||
last: SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0), |
||||
filter: filter::Filter { |
||||
gamedir: b"valve", |
||||
clver: filter::Version::new(0, 20), |
||||
// TODO: filter
|
||||
..Default::default() |
||||
}, |
||||
}; |
||||
let n = p.encode(&mut buf)?; |
||||
sock.send(&buf[..n])?; |
||||
|
||||
let start_time = Instant::now(); |
||||
while let Some(timeout) = timeout.checked_sub(start_time.elapsed()) { |
||||
sock.set_read_timeout(Some(timeout))?; |
||||
let n = match sock.recv(&mut buf) { |
||||
Ok(n) => n, |
||||
Err(e) => match e.kind() { |
||||
io::ErrorKind::AddrInUse | io::ErrorKind::WouldBlock => break, |
||||
_ => Err(e)?, |
||||
}, |
||||
}; |
||||
if let Ok(packet) = master::QueryServersResponse::decode(&buf[..n]) { |
||||
tx.send(Message::Servers( |
||||
packet.iter().map(|i| i.to_string()).collect(), |
||||
)) |
||||
.unwrap(); |
||||
} else { |
||||
eprintln!("Unexpected packet from master {}", host); |
||||
} |
||||
} |
||||
|
||||
Ok(()) |
||||
} |
||||
|
||||
fn get_server_info( |
||||
addr: String, |
||||
versions: &[u8], |
||||
timeout: Duration, |
||||
) -> Result<ServerResult, Error> { |
||||
let sock = UdpSocket::bind("0.0.0.0:0")?; |
||||
sock.connect(&addr)?; |
||||
sock.set_read_timeout(Some(timeout))?; |
||||
|
||||
for &i in versions { |
||||
let p = game::GetServerInfo::new(i); |
||||
let mut buf = [0; 2048]; |
||||
let n = p.encode(&mut buf)?; |
||||
sock.send(&buf[..n])?; |
||||
|
||||
let n = match sock.recv(&mut buf) { |
||||
Ok(n) => n, |
||||
Err(e) => match e.kind() { |
||||
io::ErrorKind::AddrInUse | io::ErrorKind::WouldBlock => { |
||||
return Ok(ServerResult::timeout(addr)); |
||||
} |
||||
_ => Err(e)?, |
||||
}, |
||||
}; |
||||
|
||||
let response = &buf[..n]; |
||||
match server::GetServerInfoResponse::decode(response) { |
||||
Ok(packet) => { |
||||
return Ok(ServerResult::ok(addr, packet)); |
||||
} |
||||
Err(ProtocolError::InvalidProtocolVersion) => { |
||||
// try another protocol version
|
||||
} |
||||
Err(e) => { |
||||
return Ok(ServerResult::invalid(addr, e, response)); |
||||
} |
||||
} |
||||
} |
||||
|
||||
Ok(ServerResult::protocol(addr)) |
||||
} |
||||
|
||||
fn query_server_info(cli: &Cli, servers: &[String]) -> Result<(), Error> { |
||||
let (tx, rx) = mpsc::channel(); |
||||
let mut workers = 0; |
||||
|
||||
if servers.is_empty() { |
||||
for i in cli.masters.iter() { |
||||
let master = i.to_owned(); |
||||
let tx = tx.clone(); |
||||
let timeout = Duration::from_secs(cli.master_timeout as u64); |
||||
thread::spawn(move || { |
||||
if let Err(e) = query_servers(&master, timeout, &tx) { |
||||
eprintln!("master({}) error: {}", master, e); |
||||
} |
||||
tx.send(Message::End).unwrap(); |
||||
}); |
||||
workers += 1; |
||||
} |
||||
} else { |
||||
tx.send(Message::Servers(servers.to_vec())).unwrap(); |
||||
} |
||||
|
||||
let mut servers = HashMap::new(); |
||||
while let Ok(msg) = rx.recv() { |
||||
match msg { |
||||
Message::Servers(list) => { |
||||
for address in list { |
||||
let tx = tx.clone(); |
||||
let timeout = Duration::from_secs(cli.server_timeout as u64); |
||||
let versions = cli.protocol.clone(); |
||||
thread::spawn(move || { |
||||
let result = get_server_info(address.clone(), &versions, timeout) |
||||
.unwrap_or_else(|e| ServerResult::error(address, e)); |
||||
tx.send(Message::ServerResult(result)).unwrap(); |
||||
tx.send(Message::End).unwrap(); |
||||
}); |
||||
workers += 1; |
||||
} |
||||
} |
||||
Message::End => { |
||||
workers -= 1; |
||||
if workers == 0 { |
||||
break; |
||||
} |
||||
} |
||||
Message::ServerResult(result) => { |
||||
servers.insert(result.address.clone(), result); |
||||
} |
||||
} |
||||
} |
||||
|
||||
let mut servers: Vec<_> = servers.values().collect(); |
||||
servers.sort_by(|a, b| cmp_address(&a.address, &b.address)); |
||||
|
||||
if cli.json || cli.debug { |
||||
let result = InfoResult { |
||||
protocol: &cli.protocol, |
||||
master_timeout: cli.master_timeout, |
||||
server_timeout: cli.server_timeout, |
||||
masters: &cli.masters, |
||||
servers: &servers, |
||||
}; |
||||
|
||||
if cli.json { |
||||
println!("{}", serde_json::to_string_pretty(&result).unwrap()); |
||||
} else if cli.debug { |
||||
println!("{:#?}", result); |
||||
} else { |
||||
todo!() |
||||
} |
||||
} else { |
||||
for i in servers { |
||||
println!("server: {}", i.address); |
||||
|
||||
macro_rules! p { |
||||
($($key:ident: $value:expr),+ $(,)?) => { |
||||
$(println!(" {}: \"{}\"", stringify!($key), $value);)+ |
||||
}; |
||||
} |
||||
|
||||
match &i.kind { |
||||
ServerResultKind::Ok { info } => { |
||||
p! { |
||||
type: "ok", |
||||
host: info.host, |
||||
gamedir: info.gamedir, |
||||
map: info.map, |
||||
protocol: info.protocol, |
||||
numcl: info.numcl, |
||||
maxcl: info.maxcl, |
||||
dm: info.dm, |
||||
team: info.team, |
||||
coop: info.coop, |
||||
password: info.password, |
||||
} |
||||
} |
||||
ServerResultKind::Timeout => { |
||||
p! { |
||||
type: "timeout", |
||||
} |
||||
} |
||||
ServerResultKind::Protocol => { |
||||
p! { |
||||
type: "protocol", |
||||
} |
||||
} |
||||
ServerResultKind::Error { message } => { |
||||
p! { |
||||
type: "error", |
||||
message: message, |
||||
} |
||||
} |
||||
ServerResultKind::Invalid { message, response } => { |
||||
p! { |
||||
type: "invalid", |
||||
message: message, |
||||
response: response, |
||||
} |
||||
} |
||||
} |
||||
println!(); |
||||
} |
||||
} |
||||
|
||||
Ok(()) |
||||
} |
||||
|
||||
fn list_servers(cli: &Cli) -> Result<(), Error> { |
||||
let (tx, rx) = mpsc::channel(); |
||||
let mut workers = 0; |
||||
|
||||
for i in cli.masters.iter() { |
||||
let master = i.to_owned(); |
||||
let tx = tx.clone(); |
||||
let timeout = Duration::from_secs(cli.master_timeout as u64); |
||||
thread::spawn(move || { |
||||
if let Err(e) = query_servers(&master, timeout, &tx) { |
||||
eprintln!("master({}) error: {}", master, e); |
||||
} |
||||
tx.send(Message::End).unwrap(); |
||||
}); |
||||
workers += 1; |
||||
} |
||||
|
||||
let mut servers = HashSet::new(); |
||||
while let Ok(msg) = rx.recv() { |
||||
match msg { |
||||
Message::Servers(list) => { |
||||
servers.extend(list); |
||||
} |
||||
Message::End => { |
||||
workers -= 1; |
||||
if workers == 0 { |
||||
break; |
||||
} |
||||
} |
||||
_ => panic!(), |
||||
} |
||||
} |
||||
|
||||
let mut servers: Vec<_> = servers.iter().map(|i| i.as_str()).collect(); |
||||
servers.sort_by(|a, b| cmp_address(a, b)); |
||||
|
||||
if cli.json || cli.debug { |
||||
let result = ListResult { |
||||
master_timeout: cli.master_timeout, |
||||
masters: &cli.masters, |
||||
servers: &servers, |
||||
}; |
||||
|
||||
if cli.json { |
||||
println!("{}", serde_json::to_string_pretty(&result).unwrap()); |
||||
} else if cli.debug { |
||||
println!("{:#?}", result); |
||||
} else { |
||||
todo!() |
||||
} |
||||
} else { |
||||
for i in servers { |
||||
println!("{}", i); |
||||
} |
||||
} |
||||
|
||||
Ok(()) |
||||
} |
||||
|
||||
fn execute(cli: Cli) -> Result<(), Error> { |
||||
match cli.args.get(0).map(|s| s.as_str()).unwrap_or_default() { |
||||
"all" | "" => query_server_info(&cli, &[])?, |
||||
"info" => query_server_info(&cli, &cli.args[1..])?, |
||||
"list" => list_servers(&cli)?, |
||||
_ => return Err(Error::UndefinedCommand), |
||||
} |
||||
|
||||
Ok(()) |
||||
} |
||||
|
||||
fn main() { |
||||
let cli = cli::parse(); |
||||
|
||||
if let Err(e) = execute(cli) { |
||||
eprintln!("error: {}", e); |
||||
process::exit(1); |
||||
} |
||||
} |
||||
Loading…
Reference in new issue