netget/server/a2s/
mod.rs

1//! Steam/Source server query (A2S) over UDP. Rust owns the challenge exchange, encoding and
2//! splitting; handlers decide what the server reports.
3pub mod actions;
4pub mod wire;
5
6use crate::llm::action_helper::call_llm;
7use crate::llm::actions::protocol_trait::ActionResult;
8use crate::logging::emit::Log;
9use crate::protocol::{Event, SpawnContext};
10use anyhow::Result;
11use serde_json::{json, Value};
12use std::{hash::BuildHasher, net::SocketAddr, sync::Arc};
13use tokio::net::UdpSocket;
14use wire::Kind;
15
16pub const DEFAULT_INFO_CHALLENGE: bool = false;
17/// Queries handled at once; datagrams beyond this are dropped.
18pub const MAX_IN_FLIGHT: usize = 64;
19
20pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
21    let info_challenge = ctx
22        .startup_params
23        .as_ref()
24        .map(|p| p.get_optional_bool("info_challenge"))
25        .transpose()?
26        .flatten()
27        .unwrap_or(DEFAULT_INFO_CHALLENGE);
28    let socket = Arc::new(UdpSocket::bind(ctx.legacy_listen_addr()).await?);
29    let addr = socket.local_addr()?;
30    Log::new(Some(&ctx.status_tx)).info(format!("A2S listening on udp {addr}"));
31    let child = ctx.clone();
32    let task = tokio::spawn(serve(child, socket, info_challenge));
33    ctx.state.register_server_task(ctx.server_id, task).await;
34    Ok(addr)
35}
36
37/// A per-server keyed hash of the client's address: no state per client, unforgeable
38/// without receiving at that address. Never 0xFFFFFFFF, which asks for a challenge.
39fn challenge_for(key: &std::collections::hash_map::RandomState, peer: SocketAddr) -> u32 {
40    let value = key.hash_one(peer) as u32;
41    if value == wire::NO_CHALLENGE {
42        0x5A5A_5A5A
43    } else {
44        value
45    }
46}
47
48fn outcome(ctx: &SpawnContext, peer: SocketAddr, query: &str, decision: &str) {
49    let summary = format!("A2S datagram from {peer} operation={query} decision={decision}");
50    let log = Log::new(Some(&ctx.status_tx));
51    if decision.starts_with("fail_closed_") || decision == "model_silent" {
52        log.error(summary);
53    } else {
54        log.info(summary);
55    }
56}
57
58/// The handler's encoded answer; `None` when it refused, failed or answered the wrong query.
59async fn answer(ctx: &SpawnContext, kind: Kind, peer: SocketAddr) -> Option<Vec<u8>> {
60    let event = Event::new(
61        &actions::QUERY_EVENT,
62        json!({"query": kind.name(), "remote_addr": peer.to_string()}),
63    );
64    let result = match call_llm(
65        &ctx.llm_client,
66        &ctx.state,
67        ctx.server_id,
68        None,
69        &event,
70        &actions::A2sProtocol,
71    )
72    .await
73    {
74        Ok(result) => result,
75        Err(_) => {
76            outcome(ctx, peer, kind.name(), "fail_closed_llm_error");
77            return None;
78        }
79    };
80    if !result.failures.is_empty() {
81        outcome(ctx, peer, kind.name(), "fail_closed_invalid_reply");
82        return None;
83    }
84    let mut found: Option<(String, Value)> = None;
85    let mut pending = result.protocol_results;
86    while let Some(item) = pending.pop() {
87        match item {
88            ActionResult::Custom { name, data } if name.starts_with("a2s_") => {
89                if found.is_some() {
90                    outcome(ctx, peer, kind.name(), "fail_closed_invalid_reply");
91                    return None;
92                }
93                found = Some((name, data));
94            }
95            ActionResult::Multiple(items) => pending.extend(items),
96            _ => {}
97        }
98    }
99    let Some((name, data)) = found else {
100        outcome(ctx, peer, kind.name(), "model_silent");
101        return None;
102    };
103    if name == "a2s_refuse" {
104        outcome(ctx, peer, kind.name(), "model_reject");
105        return None;
106    }
107    if name != actions::answer_name(kind) {
108        outcome(ctx, peer, kind.name(), "fail_closed_invalid_reply");
109        return None;
110    }
111    let encoded = match kind {
112        Kind::Info => wire::encode_info(&data),
113        Kind::Players => wire::encode_players(&data),
114        Kind::Rules => wire::encode_rules(&data),
115    };
116    match encoded {
117        Ok(payload) => {
118            outcome(ctx, peer, kind.name(), "model_answer");
119            Some(payload)
120        }
121        Err(_) => {
122            outcome(ctx, peer, kind.name(), "fail_closed_invalid_reply");
123            None
124        }
125    }
126}
127
128async fn serve(ctx: SpawnContext, socket: Arc<UdpSocket>, info_challenge: bool) {
129    let key = std::collections::hash_map::RandomState::new();
130    let in_flight = Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT));
131    let mut buf = vec![0u8; wire::MAX_REQUEST + 1];
132    let mut split_id: u32 = 1;
133    loop {
134        let Ok((n, peer)) = socket.recv_from(&mut buf).await else {
135            continue;
136        };
137        let (kind, challenge) = match wire::parse_request(&buf[..n]) {
138            Ok(request) => request,
139            Err(e) => {
140                Log::new(Some(&ctx.status_tx))
141                    .debug(format!("A2S dropped a datagram from {peer}: {e}"));
142                continue;
143            }
144        };
145        let expected = challenge_for(&key, peer);
146        let needs_challenge = kind != Kind::Info || info_challenge;
147        if needs_challenge && challenge != Some(expected) {
148            let _ = socket
149                .send_to(&wire::encode_challenge(expected), peer)
150                .await;
151            continue;
152        }
153        let Ok(permit) = in_flight.clone().try_acquire_owned() else {
154            Log::new(Some(&ctx.status_tx)).warn(format!(
155                "A2S dropped a query from {peer}: too many in flight"
156            ));
157            continue;
158        };
159        let id = split_id;
160        split_id = split_id.wrapping_add(1) & 0x7FFF_FFFF;
161        let child = ctx.clone();
162        let reply = socket.clone();
163        ctx.state
164            .spawn_server_task(ctx.server_id, async move {
165                let _permit = permit;
166                let Some(payload) = answer(&child, kind, peer).await else {
167                    return;
168                };
169                match wire::packetize(&payload, id) {
170                    Ok(packets) => {
171                        for packet in packets {
172                            let _ = reply.send_to(&packet, peer).await;
173                        }
174                    }
175                    Err(e) => Log::new(Some(&child.status_tx))
176                        .warn(format!("A2S answer to {peer} not sent: {e}")),
177                }
178            })
179            .await;
180    }
181}
netget/server/a2s/mod.rs

netget/server/a2s/
mod.rs
