netget/server/diameter/
mod.rs

1pub mod actions;
2pub mod codec;
3use crate::llm::actions::{executor::ExecutionResult, protocol_trait::ActionResult};
4use crate::logging::emit::Log;
5use crate::protocol::{Event, SpawnContext};
6use crate::server::{
7    accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS},
8    connection::ConnectionId,
9};
10use crate::state::{
11    server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo},
12    AccessLogOwner,
13};
14use anyhow::{bail, ensure, Context, Result};
15use codec::*;
16use serde_json::{json, Value};
17use std::{future::Future, net::SocketAddr, pin::Pin, sync::Arc, time::Duration};
18use tokio::{
19    net::{
20        tcp::{OwnedReadHalf, OwnedWriteHalf},
21        TcpListener, TcpStream,
22    },
23    time::Instant,
24};
25
26pub struct Config {
27    pub identity: Identity,
28    pub io_timeout: Duration,
29    pub handler_timeout: Duration,
30    pub watchdog_interval: Duration,
31    pub llm_fallback: bool,
32}
33impl Config {
34    pub fn from_params(p: &crate::protocol::StartupParams) -> Result<Self> {
35        Ok(Self {
36            identity: actions::identity(p)?,
37            io_timeout: timeout(
38                p.get_optional_u64("io_timeout_seconds")?,
39                DEFAULT_IO_SECONDS,
40            )?,
41            handler_timeout: timeout(
42                p.get_optional_u64("handler_timeout_seconds")?,
43                DEFAULT_HANDLER_SECONDS,
44            )?,
45            watchdog_interval: timeout(
46                p.get_optional_u64("watchdog_interval_seconds")?,
47                DEFAULT_WATCHDOG_SECONDS,
48            )?,
49            llm_fallback: p
50                .get_optional_bool("llm_fallback")?
51                .unwrap_or(DEFAULT_LLM_FALLBACK),
52        })
53    }
54}
55pub type Reader = Pin<Box<dyn Future<Output = (Result<Packet>, OwnedReadHalf)> + Send>>;
56pub fn reader(mut half: OwnedReadHalf, deadline: Duration) -> Reader {
57    Box::pin(async move {
58        let packet = read_packet(&mut half, deadline).await;
59        (packet, half)
60    })
61}
62pub struct DiameterServer;
63impl DiameterServer {
64    pub async fn spawn(mut ctx: SpawnContext) -> Result<SocketAddr> {
65        let cfg = Arc::new(Config::from_params(
66            ctx.startup_params
67                .as_ref()
68                .context("Diameter identity parameters required")?,
69        )?);
70        ctx.startup_params = None;
71        let listener =
72            TcpListener::bind(ctx.socket_addr().context("Diameter TCP bind address")?).await?;
73        let local = listener.local_addr()?;
74        let sid = ctx.server_id;
75        Log::new(Some(&ctx.status_tx))
76            .info(format!("Diameter experimental NASREQ listening on {local}"));
77        ctx.state
78            .clone()
79            .spawn_server_task(sid, async move {
80                let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
81                loop {
82                    let (socket, peer, permit) = match accept_bounded(
83                        &listener,
84                        &limiter,
85                        b"",
86                        "DIAMETER",
87                        Some(&ctx.status_tx),
88                    )
89                    .await
90                    {
91                        Ok(v) => v,
92                        Err(_) => break,
93                    };
94                    let cid = ConnectionId::new(ctx.state.get_next_unified_id().await);
95                    let now = crate::utils::clock::Instant::now();
96                    ctx.state
97                        .add_connection_to_server(
98                            sid,
99                            ConnectionState {
100                                id: cid,
101                                remote_addr: peer,
102                                local_addr: local,
103                                bytes_sent: 0,
104                                bytes_received: 0,
105                                packets_sent: 0,
106                                packets_received: 0,
107                                last_activity: now,
108                                status: ConnectionStatus::Active,
109                                status_changed_at: now,
110                                protocol_info: ProtocolConnectionInfo::empty(),
111                            },
112                        )
113                        .await;
114                    let child = ctx.clone();
115                    let config = cfg.clone();
116                    ctx.state
117                        .spawn_server_task(sid, async move {
118                            let _permit = permit;
119                            if session(&child, cid, socket, peer, &config).await.is_err() {
120                                decision(&child, cid, "fail_closed_peer_session");
121                            }
122                            child
123                                .state
124                                .update_connection_status(sid, cid, ConnectionStatus::Closed)
125                                .await;
126                            child.state.remove_connection_from_server(sid, cid).await;
127                            let _ = child.status_tx.send("__UPDATE_UI__".into());
128                        })
129                        .await;
130                }
131            })
132            .await;
133        Ok(local)
134    }
135}
136fn decision(ctx: &SpawnContext, id: ConnectionId, tag: &str) {
137    Log::new(Some(&ctx.status_tx)).info(format!("Diameter connection={id} decision={tag}"));
138}
139async fn send(
140    ctx: &SpawnContext,
141    id: ConnectionId,
142    w: &mut OwnedWriteHalf,
143    p: &Packet,
144) -> Result<()> {
145    let n = write_packet(w, p).await?;
146    ctx.state
147        .update_connection_stats(ctx.server_id, id, None, Some(n as u64), None, Some(1))
148        .await;
149    Ok(())
150}
151async fn received(ctx: &SpawnContext, id: ConnectionId, p: &Packet) {
152    ctx.state
153        .update_connection_stats(
154            ctx.server_id,
155            id,
156            Some(p.encode().map(|b| b.len()).unwrap_or(0) as u64),
157            None,
158            Some(1),
159            None,
160        )
161        .await;
162}
163pub fn control(p: &Packet, peer: &Identity) -> Result<()> {
164    ensure!(
165        p.application == 0 && p.flags & 0x40 == 0 && p.flags & 0x20 == 0,
166        "Diameter peer-control header"
167    );
168    peer.same_peer(p)?;
169    let allowed = if p.command == DPR {
170        &[HOST, REALM, ORIGIN_STATE, DISCONNECT_CAUSE, RESULT][..]
171    } else {
172        &[HOST, REALM, ORIGIN_STATE, RESULT][..]
173    };
174    ensure!(
175        p.unsupported_mandatory(allowed).is_none(),
176        "Diameter unsupported mandatory peer-control AVP"
177    );
178    if let Some(a) = p.optional(ORIGIN_STATE)? {
179        a.integer()?;
180    }
181    if p.is_request() && p.command == DPR {
182        ensure!(p.num(DISCONNECT_CAUSE)? <= 2, "Diameter Disconnect-Cause");
183    }
184    if !p.is_request() {
185        ensure!(p.num(RESULT)? == 2001, "Diameter peer-control rejected");
186    }
187    Ok(())
188}
189/// All execution failures precede any positive wire verdict. Constructed result
190/// trees and raw actions are consumed iteratively before JSON copies or decoding.
191pub fn chosen_reply(result: ExecutionResult) -> Result<Option<Reply>> {
192    let mut invalid = !result.failures.is_empty() || result.raw_actions.len() > 32;
193    for v in result.raw_actions {
194        invalid |= !within_json_budget(&v);
195        crate::utils::json_budget::drop_iteratively(v);
196    }
197    let mut pending = result.protocol_results;
198    let mut found: Option<Value> = None;
199    let mut seen = 0usize;
200    while let Some(result) = pending.pop() {
201        seen = seen.saturating_add(1);
202        invalid |= seen > 64;
203        match result {
204            ActionResult::Multiple(v) => pending.extend(v),
205            ActionResult::Custom { name, data } => {
206                if name == "respond_diameter_aa" && found.is_none() && within_json_budget(&data) {
207                    found = Some(data);
208                } else {
209                    invalid = true;
210                    crate::utils::json_budget::drop_iteratively(data);
211                }
212            }
213            ActionResult::NoAction => {}
214            _ => invalid = true,
215        }
216    }
217    if invalid {
218        if let Some(v) = found {
219            crate::utils::json_budget::drop_iteratively(v);
220        }
221        bail!("Diameter handler failure/type/count/budget");
222    }
223    found
224        .map(|v| {
225            let reply: Reply = serde_json::from_value(v)?;
226            reply.validate()?;
227            Ok(reply)
228        })
229        .transpose()
230}
231type Handler = Pin<Box<dyn Future<Output = Result<Option<Reply>>> + Send>>;
232fn handler(ctx: SpawnContext, cid: ConnectionId, event: Event, cfg: Arc<Config>) -> Handler {
233    Box::pin(async move {
234        let configured = ctx
235            .state
236            .get_event_handler_config(ctx.server_id)
237            .await
238            .is_some_and(|c| c.find_handler(event.id()).is_some());
239        if !configured && !cfg.llm_fallback {
240            let event_id = event.id().to_owned();
241            ctx.state
242                .record_access_log(
243                    AccessLogOwner::Server(ctx.server_id.as_u32()),
244                    "DIAMETER",
245                    Some(cid.as_u32()),
246                    &event_id,
247                    event.data,
248                    vec![],
249                )
250                .await;
251            return Ok(None);
252        }
253        let result = tokio::time::timeout(
254            cfg.handler_timeout,
255            crate::llm::action_helper::call_llm(
256                &ctx.llm_client,
257                &ctx.state,
258                ctx.server_id,
259                Some(cid),
260                &event,
261                &actions::DiameterProtocol,
262            ),
263        )
264        .await
265        .context("Diameter AAA handler deadline")??;
266        chosen_reply(result)
267    })
268}
269async fn refuse(
270    ctx: &SpawnContext,
271    cid: ConnectionId,
272    w: &mut OwnedWriteHalf,
273    p: &Packet,
274    id: &Identity,
275    code: u32,
276    failed: Option<&Avp>,
277    tag: &str,
278) -> Result<()> {
279    decision(ctx, cid, tag);
280    send(ctx, cid, w, &error_answer(p, id, code, failed)?).await
281}
282async fn session(
283    ctx: &SpawnContext,
284    cid: ConnectionId,
285    socket: TcpStream,
286    source: SocketAddr,
287    cfg: &Arc<Config>,
288) -> Result<()> {
289    let local_ip = socket.local_addr()?.ip();
290    let (mut half, mut w) = socket.into_split();
291    let first = read_packet(&mut half, cfg.io_timeout).await?;
292    received(ctx, cid, &first).await;
293    ensure!(
294        first.is_request() && first.command == CER,
295        "Diameter CER required first"
296    );
297    let peer = match capabilities(&first) {
298        Ok(peer) => peer,
299        Err(_) => {
300            let failed = first.unsupported_mandatory(BASE_ALLOWED);
301            let security = first.avps.iter().any(|a| {
302                a.code == SECURITY && a.vendor.is_none() && a.integer().ok().is_some_and(|v| v != 0)
303            });
304            let code =
305                if failed.is_some() {
306                    5001
307                } else if security {
308                    5017
309                } else if !first.avps.iter().any(|a| {
310                    a.code == AUTH_APP && a.vendor.is_none() && a.integer().ok() == Some(1)
311                }) {
312                    5010
313                } else {
314                    5012
315                };
316            decision(ctx, cid, "fail_closed_capabilities");
317            let mut answer = error_answer(&first, &cfg.identity, code, failed)?;
318            answer.avps.retain(|a| ![HOST, REALM].contains(&a.code));
319            capability_fields(&mut answer, &cfg.identity, local_ip);
320            send(ctx, cid, &mut w, &answer).await?;
321            return Ok(());
322        }
323    };
324    let mut cea = first.answer(2001);
325    capability_fields(&mut cea, &cfg.identity, local_ip);
326    send(ctx, cid, &mut w, &cea).await?;
327    decision(ctx, cid, "peer_capabilities_accepted");
328    let mut reading = reader(half, cfg.io_timeout);
329    let mut current: Option<(Packet, Handler)> = None;
330    let mut watchdog: Option<(Packet, Instant)> = None;
331    let mut next_watchdog = Instant::now() + cfg.watchdog_interval;
332    loop {
333        tokio::select! {
334         (result,half)=&mut reading=>{
335             let p = result?;
336             received(ctx, cid, &p).await;
337             reading = reader(half, cfg.io_timeout);
338             next_watchdog = Instant::now() + cfg.watchdog_interval;
339             if !p.is_request() {
340                 let Some((request, _)) = watchdog.take() else {
341                     bail!("Diameter unexpected answer");
342                 };
343                 ensure!(p.matches(&request), "Diameter watchdog correlation");
344                 control(&p, &peer)?;
345                 continue;
346             }
347             match p.command {
348                 DWR | DPR => {
349                     control(&p, &peer)?;
350                     let mut answer = p.answer(2001);
351                     answer.origin(&cfg.identity);
352                     send(ctx, cid, &mut w, &answer).await?;
353                     if p.command == DPR {
354                         decision(ctx, cid, "peer_disconnect_accepted");
355                         break;
356                     }
357                 }
358                 CER => {
359                     refuse(
360                         ctx,
361                         cid,
362                         &mut w,
363                         &p,
364                         &cfg.identity,
365                         5012,
366                         None,
367                         "fail_closed_repeated_capabilities",
368                     )
369                     .await?;
370                     break;
371                 }
372                 AA if p.application == 1 => {
373                     if let Some(failed) = p.unsupported_mandatory(REQUEST_ALLOWED) {
374                         refuse(
375                             ctx,
376                             cid,
377                             &mut w,
378                             &p,
379                             &cfg.identity,
380                             5001,
381                             Some(failed),
382                             "fail_closed_mandatory_avp",
383                         )
384                         .await?;
385                         continue;
386                     }
387                     let (request, session) = match Request::from_packet(&p, &cfg.identity, &peer) {
388                         Ok(v) => v,
389                         Err(_) => {
390                             refuse(
391                                 ctx,
392                                 cid,
393                                 &mut w,
394                                 &p,
395                                 &cfg.identity,
396                                 5012,
397                                 None,
398                                 "fail_closed_request_semantics",
399                             )
400                             .await?;
401                             continue;
402                         }
403                     };
404                     if current.is_some() {
405                         refuse(
406                             ctx,
407                             cid,
408                             &mut w,
409                             &p,
410                             &cfg.identity,
411                             3004,
412                             None,
413                             "fail_closed_handler_capacity",
414                         )
415                         .await?;
416                         continue;
417                     }
418                     let mut data = serde_json::to_value(request)?;
419                     data["session_id"] = json!(session);
420                     data["origin_host"] = json!(peer.host);
421                     data["origin_realm"] = json!(peer.realm);
422                     data["source_addr"] = json!(source.to_string());
423                     let event = Event::new(&actions::AA_EVENT, json!({"request":data}));
424                     current = Some((p, handler(ctx.clone(), cid, event, cfg.clone())));
425                 }
426                 _ => {
427                     let code = if p.application != 0 && p.application != 1 {
428                         3007
429                     } else {
430                         3001
431                     };
432                     refuse(
433                         ctx,
434                         cid,
435                         &mut w,
436                         &p,
437                         &cfg.identity,
438                         code,
439                         None,
440                         "fail_closed_unsupported_command_application",
441                     )
442                     .await?;
443                 }
444             }
445         }
446         result=async{current.as_mut().unwrap().1.as_mut().await},if current.is_some()=>{
447             let (request, _) = current.take().unwrap();
448             match result {
449                 Ok(reply) => {
450                     let reply = reply.unwrap_or_default();
451                     decision(
452                         ctx,
453                         cid,
454                         match reply.verdict {
455                             Verdict::Accept => "handler_accept",
456                             Verdict::Reject => "handler_reject_or_no_action",
457                             Verdict::Error => "handler_error",
458                         },
459                     );
460                     let answer = reply.packet(&request, &cfg.identity)?;
461                     send(ctx, cid, &mut w, &answer).await?;
462                 }
463                 Err(error) => {
464                     let code = match crate::utils::WireFailure::classify(&error) {
465                         crate::utils::WireFailure::Overloaded => 3004,
466                         crate::utils::WireFailure::Unavailable => 5012,
467                     };
468                     refuse(
469                         ctx,
470                         cid,
471                         &mut w,
472                         &request,
473                         &cfg.identity,
474                         code,
475                         None,
476                         "fail_closed_handler_failure",
477                     )
478                     .await?;
479                 }
480             }
481             next_watchdog = Instant::now() + cfg.watchdog_interval;
482         }
483         _=tokio::time::sleep_until(next_watchdog),if watchdog.is_none()=>{
484             let mut request = Packet::request(DWR, 0)?;
485             request.origin(&cfg.identity);
486             send(ctx, cid, &mut w, &request).await?;
487             watchdog = Some((request, Instant::now() + cfg.io_timeout));
488         }
489         _=tokio::time::sleep_until(watchdog.as_ref().map(|v|v.1).unwrap_or_else(Instant::now)),if watchdog.is_some()=>{
490             decision(ctx, cid, "fail_closed_watchdog_deadline");
491             bail!("Diameter watchdog answer deadline");
492         }
493        }
494    }
495    Ok(())
496}
netget/server/diameter/mod.rs

netget/server/diameter/
mod.rs
