Browse Source

query: do not use threads

ipv6
Denis Drakhnia 1 year ago
parent
commit
62fe4d0b6c
  1. 4
      protocol/src/master.rs
  2. 351
      query/src/main.rs

4
protocol/src/master.rs

@ -76,6 +76,10 @@ impl<'a> QueryServersResponse<&'a [u8]> {
SocketAddrV4::new(ip, port) SocketAddrV4::new(ip, port)
}) })
} }
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
} }
impl<I> QueryServersResponse<I> impl<I> QueryServersResponse<I>

351
query/src/main.rs

@ -7,10 +7,9 @@ use std::cmp;
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::fmt; use std::fmt;
use std::io; use std::io;
use std::net::{Ipv4Addr, SocketAddrV4, UdpSocket}; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
use std::process; use std::process;
use std::sync::{mpsc, Arc}; use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use serde::{Serialize, Serializer}; use serde::{Serialize, Serializer};
@ -50,34 +49,34 @@ enum ServerResultKind {
#[derive(Clone, Debug, Serialize)] #[derive(Clone, Debug, Serialize)]
struct ServerResult { struct ServerResult {
address: String, address: SocketAddrV4,
ping: f32, ping: f32,
#[serde(flatten)] #[serde(flatten)]
kind: ServerResultKind, kind: ServerResultKind,
} }
impl ServerResult { impl ServerResult {
fn new(address: String, ping: f32, kind: ServerResultKind) -> Self { fn new(address: SocketAddrV4, ping: f32, kind: ServerResultKind) -> Self {
Self { Self {
address: address.to_string(), address,
ping, ping,
kind, kind,
} }
} }
fn ok(address: String, ping: f32, info: ServerInfo) -> Self { fn ok(address: SocketAddrV4, ping: f32, info: ServerInfo) -> Self {
Self::new(address, ping, ServerResultKind::Ok { info }) Self::new(address, ping, ServerResultKind::Ok { info })
} }
fn timeout(address: String) -> Self { fn timeout(address: SocketAddrV4) -> Self {
Self::new(address, 0.0, ServerResultKind::Timeout) Self::new(address, 0.0, ServerResultKind::Timeout)
} }
fn protocol(address: String, ping: f32) -> Self { fn protocol(address: SocketAddrV4, ping: f32) -> Self {
Self::new(address, ping, ServerResultKind::Protocol) Self::new(address, ping, ServerResultKind::Protocol)
} }
fn error<T>(address: String, message: T) -> Self fn error<T>(address: SocketAddrV4, message: T) -> Self
where where
T: fmt::Display, T: fmt::Display,
{ {
@ -90,7 +89,7 @@ impl ServerResult {
) )
} }
fn invalid<T>(address: String, ping: f32, message: T, response: &[u8]) -> Self fn invalid<T>(address: SocketAddrV4, ping: f32, message: T, response: &[u8]) -> Self
where where
T: fmt::Display, T: fmt::Display,
{ {
@ -152,7 +151,7 @@ struct ListResult<'a> {
master_timeout: u32, master_timeout: u32,
masters: &'a [Box<str>], masters: &'a [Box<str>],
filter: &'a str, filter: &'a str,
servers: &'a [&'a str], servers: &'a [SocketAddrV4],
} }
fn serialize_colored<S>(s: &str, ser: S) -> Result<S::Ok, S::Error> fn serialize_colored<S>(s: &str, ser: S) -> Result<S::Ok, S::Error>
@ -203,159 +202,243 @@ impl fmt::Display for Colored<'_> {
} }
} }
enum Message { fn get_socket_addrs<'a>(iter: impl Iterator<Item = &'a str>) -> Result<Vec<SocketAddrV4>, Error> {
Servers(Vec<String>), use std::net::ToSocketAddrs;
ServerResult(ServerResult),
End, let mut out = Vec::with_capacity(iter.size_hint().0);
for i in iter {
match i
.to_socket_addrs()?
.find(|i| matches!(i, SocketAddr::V4(_)))
{
Some(SocketAddr::V4(addr)) => out.push(addr),
_ => eprintln!("warn: failed to resolve address for {}", i),
}
} }
fn cmp_address(a: &str, b: &str) -> cmp::Ordering { Ok(out)
match (a.parse::<SocketAddrV4>(), b.parse::<SocketAddrV4>()) {
(Ok(a), Ok(b)) => a.cmp(&b),
_ => a.cmp(b),
} }
struct ServerQuery {
start: Instant,
protocol: usize,
} }
fn query_servers( impl ServerQuery {
host: &str, fn ping(&self) -> f32 {
cli: &Cli, self.start.elapsed().as_micros() as f32 / 1000.0
timeout: Duration, }
tx: &mpsc::Sender<Message>, }
) -> Result<(), Error> {
let sock = UdpSocket::bind("0.0.0.0:0")?;
sock.connect(host)?;
impl ServerQuery {
fn new(protocol: usize) -> Self {
Self {
start: Instant::now(),
protocol,
}
}
}
struct Scan<'a> {
cli: &'a Cli,
masters: Vec<SocketAddrV4>,
sock: UdpSocket,
}
impl<'a> Scan<'a> {
fn new(cli: &'a Cli) -> Result<Self, Error> {
Ok(Self {
cli,
masters: get_socket_addrs(cli.masters.iter().map(|i| i.as_ref()))?,
sock: UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0))?,
})
}
fn is_master(&self, addr: &SocketAddrV4) -> bool {
self.masters.iter().any(|i| i == addr)
}
fn query_servers(&self) -> Result<(), Error> {
let mut buf = [0; 512]; let mut buf = [0; 512];
let p = game::QueryServers { let packet = game::QueryServers {
region: server::Region::RestOfTheWorld, region: server::Region::RestOfTheWorld,
last: SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0), last: SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0),
filter: cli.filter.as_str(), filter: self.cli.filter.as_str(),
}; };
let n = p.encode(&mut buf)?; let n = packet.encode(&mut buf)?;
sock.send(&buf[..n])?; let packet = &buf[..n];
for i in &self.masters {
self.sock.send_to(packet, i)?;
}
Ok(())
}
fn servers(&self) -> Result<HashSet<SocketAddrV4>, Error> {
self.query_servers()?;
let mut set = HashSet::with_capacity(256);
let mut buf = [0; 2048];
let timeout = Duration::from_secs(self.cli.master_timeout as u64);
let start_time = Instant::now(); let start_time = Instant::now();
while let Some(timeout) = timeout.checked_sub(start_time.elapsed()) { while let Some(timeout) = timeout.checked_sub(start_time.elapsed()) {
sock.set_read_timeout(Some(timeout))?; self.sock.set_read_timeout(Some(timeout))?;
let n = match sock.recv(&mut buf) {
Ok(n) => n, let (n, from) = match self.sock.recv_from(&mut buf) {
Ok(x) => x,
Err(e) => match e.kind() { Err(e) => match e.kind() {
io::ErrorKind::AddrInUse | io::ErrorKind::WouldBlock => break, io::ErrorKind::AddrInUse => break,
io::ErrorKind::WouldBlock => break,
_ => Err(e)?, _ => Err(e)?,
}, },
}; };
let from = match from {
SocketAddr::V4(x) => x,
_ => todo!(),
};
if self.is_master(&from) {
if let Ok(packet) = master::QueryServersResponse::decode(&buf[..n]) { if let Ok(packet) = master::QueryServersResponse::decode(&buf[..n]) {
tx.send(Message::Servers( set.extend(packet.iter());
packet.iter().map(|i| i.to_string()).collect(),
))
.unwrap();
} else { } else {
eprintln!("Unexpected packet from master {}", host); eprintln!("warn: invalid packet from master {}", from);
} }
} }
Ok(())
} }
fn get_server_info( Ok(set)
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))?;
let mut ping = 0.0; fn server_info(
for &i in versions { &self,
let p = game::GetServerInfo::new(i); list: &[SocketAddrV4],
) -> Result<HashMap<SocketAddrV4, ServerResult>, Error> {
let mut set = HashSet::new();
let mut active = HashMap::new();
let mut out = HashMap::new();
let mut buf = [0; 2048]; let mut buf = [0; 2048];
let n = p.encode(&mut buf)?;
let start = Instant::now();
sock.send(&buf[..n])?;
let n = match sock.recv(&mut buf) { let now = Instant::now();
Ok(n) => n, let master_timeout = Duration::from_secs(self.cli.master_timeout as u64);
Err(e) => match e.kind() { let server_timeout = Duration::from_secs(self.cli.server_timeout as u64);
io::ErrorKind::AddrInUse | io::ErrorKind::WouldBlock => { let master_end = now + master_timeout;
return Ok(ServerResult::timeout(addr)); let mut server_end = now + server_timeout;
if list.is_empty() {
self.query_servers()?;
} else {
let mut buf = [0; 512];
let n = game::GetServerInfo::new(self.cli.protocol[0]).encode(&mut buf)?;
for addr in list.iter().filter(|i| set.insert(**i)) {
match self.sock.send_to(&buf[..n], addr) {
Ok(_) => {
let query = ServerQuery::new(0);
server_end = query.start + server_timeout;
active.insert(*addr, query);
}
Err(e) => {
let res = ServerResult::error(*addr, e);
out.insert(*addr, res);
}
} }
}
}
loop {
let time = cmp::max(master_end, server_end);
match time.checked_duration_since(Instant::now()) {
Some(t) => self.sock.set_read_timeout(Some(t))?,
None => break,
}
let (n, from) = match self.sock.recv_from(&mut buf) {
Ok(x) => x,
Err(e) => match e.kind() {
io::ErrorKind::AddrInUse => break,
io::ErrorKind::WouldBlock => break,
_ => Err(e)?, _ => Err(e)?,
}, },
}; };
ping = start.elapsed().as_micros() as f32 / 1000.0; let from = match from {
SocketAddr::V4(x) => x,
_ => todo!(),
};
let raw = &buf[..n];
let response = &buf[..n]; if self.is_master(&from) {
match server::GetServerInfoResponse::decode(response) { if let Ok(packet) = master::QueryServersResponse::decode(raw) {
Ok(packet) => { for addr in packet.iter().filter(|i| set.insert(*i)) {
let info = ServerInfo::from(packet); let mut buf = [0; 512];
return Ok(ServerResult::ok(addr, ping, info)); let n = game::GetServerInfo::new(self.cli.protocol[0]).encode(&mut buf)?;
}
Err(ProtocolError::InvalidProtocolVersion) => { match self.sock.send_to(&buf[..n], addr) {
// try another protocol version Ok(_) => {
let query = ServerQuery::new(0);
server_end = query.start + server_timeout;
active.insert(addr, query);
} }
Err(e) => { Err(e) => {
return Ok(ServerResult::invalid(addr, ping, e, response)); let res = ServerResult::error(addr, e);
out.insert(addr, res);
} }
} }
} }
Ok(ServerResult::protocol(addr, ping))
} }
} else if let Some(query) = active.remove(&from) {
match server::GetServerInfoResponse::decode(raw) {
Ok(packet) => {
let info = ServerInfo::from(packet);
let res = ServerResult::ok(from, query.ping(), info);
out.insert(from, res);
}
Err(ProtocolError::InvalidProtocolVersion) => {
let next_protocol = query.protocol + 1;
if let Some(protocol) = self.cli.protocol.get(next_protocol) {
let mut buf = [0; 512];
let n = game::GetServerInfo::new(*protocol).encode(&mut buf)?;
fn query_server_info(cli: &Arc<Cli>, servers: &[String]) -> Result<(), Error> { match self.sock.send_to(&buf[..n], from) {
let (tx, rx) = mpsc::channel(); Ok(_) => {
let mut workers = 0; active.insert(from, ServerQuery::new(next_protocol));
}
if servers.is_empty() { Err(e) => {
for i in cli.masters.iter() { let res = ServerResult::error(from, e);
let master = i.to_owned(); out.insert(from, res);
let tx = tx.clone(); }
let timeout = Duration::from_secs(cli.master_timeout as u64);
let cli = cli.clone();
thread::spawn(move || {
if let Err(e) = query_servers(&master, &cli, timeout, &tx) {
eprintln!("master({}) error: {}", master, e);
}
tx.send(Message::End).unwrap();
});
workers += 1;
} }
} else { } else {
tx.send(Message::Servers(servers.to_vec())).unwrap(); let res = ServerResult::protocol(from, query.ping());
out.insert(from, res);
} }
}
let mut servers = HashMap::new(); Err(e) => {
while let Ok(msg) = rx.recv() { let res = ServerResult::invalid(from, query.ping(), e, raw);
match msg { out.insert(from, res);
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); for (addr, _) in active {
let res = ServerResult::timeout(addr);
out.insert(addr, res);
} }
Ok(out)
} }
} }
fn query_server_info(cli: &Arc<Cli>, servers: &[String]) -> Result<(), Error> {
let scan = Scan::new(cli)?;
let servers = get_socket_addrs(servers.iter().map(|i| i.as_str()))?;
let servers = scan.server_info(&servers)?;
let mut servers: Vec<_> = servers.values().collect(); let mut servers: Vec<_> = servers.values().collect();
servers.sort_by(|a, b| cmp_address(&a.address, &b.address)); servers.sort_by(|a, b| a.address.cmp(&b.address));
if cli.json || cli.debug { if cli.json || cli.debug {
let result = InfoResult { let result = InfoResult {
@ -432,41 +515,9 @@ fn query_server_info(cli: &Arc<Cli>, servers: &[String]) -> Result<(), Error> {
} }
fn list_servers(cli: &Arc<Cli>) -> Result<(), Error> { fn list_servers(cli: &Arc<Cli>) -> Result<(), Error> {
let (tx, rx) = mpsc::channel(); let scan = Scan::new(cli)?;
let mut workers = 0; let mut servers: Vec<_> = scan.servers()?.into_iter().collect();
servers.sort();
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);
let cli = cli.clone();
thread::spawn(move || {
if let Err(e) = query_servers(&master, &cli, 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 { if cli.json || cli.debug {
let result = ListResult { let result = ListResult {

Loading…
Cancel
Save