netget/client/diameter/
mod.rs

1pub mod actions;
2use crate::llm::actions::client_trait::{Client, ClientActionResult};
3use crate::protocol::{ConnectContext, Event};
4use crate::server::diameter::{codec::*, control, reader, Config};
5use crate::state::{
6    client_handles::{ClientCommand, ClientSendOutcome},
7    AccessLogOwner, ClientStatus,
8};
9use crate::{console_error, console_info};
10pub use actions::DiameterClientProtocol;
11use anyhow::{bail, ensure, Context, Result};
12use serde_json::{json, Value};
13use std::{collections::VecDeque, future::Future, pin::Pin, sync::Arc, time::Duration};
14use tokio::{
15    net::{tcp::OwnedWriteHalf, TcpStream},
16    time::Instant,
17};
18pub const MAX_QUEUED_EVENTS: usize = 32;
19pub const MAX_HANDLER_ACTIONS: usize = 32;
20pub const MAX_FOLLOWUP_DEPTH: usize = 8;
21type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
22fn handler(ctx: ConnectContext, event: Event, deadline: Duration) -> Handler {
23    Box::pin(async move {
24        let instruction = ctx
25            .state
26            .get_instruction_for_client(ctx.client_id)
27            .await
28            .unwrap_or_default();
29        let configured = ctx
30            .state
31            .get_client_event_handler_config(ctx.client_id)
32            .await
33            .is_some_and(|c| c.find_handler(event.id()).is_some());
34        if instruction.is_empty() && !configured {
35            let event_id = event.id().to_owned();
36            ctx.state
37                .record_access_log(
38                    AccessLogOwner::Client(ctx.client_id.as_u32()),
39                    "DIAMETER",
40                    None,
41                    &event_id,
42                    event.data,
43                    vec![],
44                )
45                .await;
46            return Ok(crate::llm::ClientLlmResult {
47                actions: vec![],
48                memory_updates: None,
49            });
50        }
51        let memory = ctx
52            .state
53            .get_memory_for_client(ctx.client_id)
54            .await
55            .unwrap_or_default();
56        tokio::time::timeout(
57            deadline,
58            crate::client::llm_budget::call_llm_for_client(
59                &ctx.llm_client,
60                &ctx.state,
61                ctx.client_id.to_string(),
62                &instruction,
63                &memory,
64                Some(&event),
65                &DiameterClientProtocol,
66                &ctx.status_tx,
67            ),
68        )
69        .await
70        .context("Diameter event handler deadline")?
71    })
72}
73struct Pending {
74    packet: Packet,
75    request: Value,
76    command: Option<ClientCommand>,
77    depth: usize,
78    deadline: Instant,
79}
80struct Closing {
81    packet: Packet,
82    command: Option<ClientCommand>,
83    deadline: Instant,
84}
85pub struct DiameterClient;
86impl DiameterClient {
87    pub async fn connect(mut ctx: ConnectContext) -> Result<std::net::SocketAddr> {
88        ensure!(
89            ctx.remote_addr.len() <= 1024 && !ctx.remote_addr.contains(['\r', '\n', '@', '/']),
90            "Diameter endpoint host:port only"
91        );
92        let p = ctx
93            .startup_params
94            .as_ref()
95            .context("Diameter identity parameters required")?;
96        let cfg = Arc::new(Config {
97            identity: crate::server::diameter::actions::identity(p)?,
98            io_timeout: timeout(
99                p.get_optional_u64("io_timeout_seconds")?,
100                DEFAULT_IO_SECONDS,
101            )?,
102            handler_timeout: timeout(
103                p.get_optional_u64("handler_timeout_seconds")?,
104                DEFAULT_HANDLER_SECONDS,
105            )?,
106            watchdog_interval: timeout(
107                p.get_optional_u64("watchdog_interval_seconds")?,
108                DEFAULT_WATCHDOG_SECONDS,
109            )?,
110            llm_fallback: false,
111        });
112        ctx.startup_params = None;
113        let socket = tokio::time::timeout(cfg.io_timeout, TcpStream::connect(&ctx.remote_addr))
114            .await
115            .context("Diameter TCP connect deadline")??;
116        let local = socket.local_addr()?;
117        let connected_addr = socket.peer_addr()?;
118        let (mut half, mut write) = socket.into_split();
119        let mut cer = Packet::request(CER, 0)?;
120        capability_fields(&mut cer, &cfg.identity, local.ip());
121        write_packet(&mut write, &cer).await?;
122        let answer = read_packet(&mut half, cfg.io_timeout).await?;
123        ensure!(answer.matches(&cer), "Diameter CEA correlation");
124        let peer = capabilities(&answer)?;
125        ctx.state
126            .with_client_mut(ctx.client_id, |client| {
127                let now = crate::utils::clock::Instant::now();
128                client.connection = Some(crate::state::client::ClientConnectionState {
129                    id: ctx.client_id,
130                    remote_addr: ctx.remote_addr.clone(),
131                    connected_addr: Some(connected_addr),
132                    local_addr: Some(local),
133                    bytes_sent: cer.encode().map(|v| v.len()).unwrap_or(0) as u64,
134                    bytes_received: answer.encode().map(|v| v.len()).unwrap_or(0) as u64,
135                    packets_sent: 1,
136                    packets_received: 1,
137                    last_activity: now,
138                    status: ClientStatus::Connected,
139                    status_changed_at: now,
140                    protocol_info: crate::state::server::ProtocolConnectionInfo::empty(),
141                });
142            })
143            .await;
144        let commands =
145            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
146                .await;
147        ctx.state
148            .update_client_status(ctx.client_id, ClientStatus::Connected)
149            .await;
150        console_info!(
151            ctx.status_tx,
152            "Diameter NASREQ capabilities negotiated with {}",
153            ctx.remote_addr
154        );
155        ctx.state
156            .clone()
157            .spawn_client_task(ctx.client_id, run(ctx, cfg, peer, half, write, commands))
158            .await;
159        Ok(local)
160    }
161}
162async fn stats(ctx: &ConnectContext, bytes: u64, received: bool) {
163    ctx.state
164        .with_client_mut(ctx.client_id, |client| {
165            if let Some(connection) = client.connection.as_mut() {
166                if received {
167                    connection.bytes_received += bytes;
168                    connection.packets_received += 1;
169                } else {
170                    connection.bytes_sent += bytes;
171                    connection.packets_sent += 1;
172                }
173                connection.last_activity = crate::utils::clock::Instant::now();
174            }
175        })
176        .await;
177}
178async fn send(ctx: &ConnectContext, write: &mut OwnedWriteHalf, packet: &Packet) -> Result<()> {
179    let bytes = write_packet(write, packet).await?;
180    stats(ctx, bytes as u64, false).await;
181    Ok(())
182}
183fn prepared(value: Value, depth: usize) -> Result<Option<Request>> {
184    ensure!(depth <= MAX_FOLLOWUP_DEPTH, "Diameter follow-up depth");
185    match DiameterClientProtocol.execute_action(value)? {
186        ClientActionResult::Disconnect => Ok(None),
187        ClientActionResult::Custom { name, data } if name == "diameter_command" => {
188            Ok(Some(actions::parse_action(data)?))
189        }
190        _ => bail!("Diameter action result"),
191    }
192}
193fn request(
194    request: Request,
195    cfg: &Config,
196    peer: &Identity,
197    command: Option<ClientCommand>,
198    depth: usize,
199) -> Result<Pending> {
200    let session = format!("{};{}", cfg.identity.host, uuid::Uuid::new_v4());
201    let packet = request.packet(&cfg.identity, peer, &session)?;
202    let mut safe = request.credential_free();
203    safe["session_id"] = json!(session);
204    Ok(Pending {
205        packet,
206        request: safe,
207        command,
208        depth,
209        deadline: Instant::now() + cfg.io_timeout,
210    })
211}
212async fn finish(ctx: &ConnectContext, command: ClientCommand, result: Result<ClientSendOutcome>) {
213    let outcome = match &result {
214        Ok(v) => serde_json::to_value(v).unwrap_or(Value::Null),
215        Err(_) => json!({"error":"Diameter operation failed or cancelled; no replay"}),
216    };
217    ctx.state
218        .record_access_log(
219            AccessLogOwner::Client(ctx.client_id.as_u32()),
220            "DIAMETER",
221            None,
222            "injected_action",
223            crate::utils::redact::redact_sensitive(&command.action),
224            vec![outcome],
225        )
226        .await;
227    crate::client::command_support::reply(command, result);
228}
229async fn begin_close(
230    ctx: &ConnectContext,
231    cfg: &Config,
232    write: &mut OwnedWriteHalf,
233    pending: &mut Option<Pending>,
234    command: Option<ClientCommand>,
235) -> Result<Closing> {
236    if let Some(p) = pending.as_mut() {
237        if let Some(command) = p.command.take() {
238            finish(
239                ctx,
240                command,
241                Err(anyhow::anyhow!("Diameter request cancelled")),
242            )
243            .await;
244        }
245    }
246    let mut packet = Packet::request(DPR, 0)?;
247    packet.origin(&cfg.identity);
248    packet.avps.push(Avp::number(DISCONNECT_CAUSE, 0));
249    send(ctx, write, &packet).await?;
250    Ok(Closing {
251        packet,
252        command,
253        deadline: Instant::now() + cfg.io_timeout,
254    })
255}
256async fn run(
257    ctx: ConnectContext,
258    cfg: Arc<Config>,
259    peer: Identity,
260    half: tokio::net::tcp::OwnedReadHalf,
261    mut write: OwnedWriteHalf,
262    mut commands: tokio::sync::mpsc::Receiver<ClientCommand>,
263) {
264    let connected = Event::new(
265        &actions::CONNECTED_EVENT,
266        json!({"peer":{"origin_host":peer.host,"origin_realm":peer.realm,"remote_addr":ctx.remote_addr}}),
267    );
268    let mut current = Some((handler(ctx.clone(), connected, cfg.handler_timeout), 0usize));
269    let mut events = VecDeque::new();
270    let mut queued: VecDeque<(Value, usize)> = VecDeque::new();
271    let mut pending: Option<Pending> = None;
272    let mut closing: Option<Closing> = None;
273    let mut watchdog: Option<(Packet, Instant)> = None;
274    let mut next_watchdog = Instant::now() + cfg.watchdog_interval;
275    let mut reading = reader(half, cfg.io_timeout);
276    let mut failed = false;
277    'peer: loop {
278        if closing.is_none() {
279            if current.is_none() {
280                if let Some((event, depth)) = events.pop_front() {
281                    current = Some((handler(ctx.clone(), event, cfg.handler_timeout), depth));
282                }
283            }
284            if pending.is_none() {
285                if let Some((action, depth)) = queued.pop_front() {
286                    match prepared(action, depth) {
287                        Ok(Some(r)) => match request(r, &cfg, &peer, None, depth) {
288                            Ok(work) => {
289                                if send(&ctx, &mut write, &work.packet).await.is_err() {
290                                    failed = true;
291                                    break;
292                                }
293                                pending = Some(work);
294                            }
295                            Err(_) => {
296                                failed = true;
297                                break;
298                            }
299                        },
300                        Ok(None) => {
301                            match begin_close(&ctx, &cfg, &mut write, &mut pending, None).await {
302                                Ok(v) => {
303                                    closing = Some(v);
304                                    drop(current.take());
305                                }
306                                Err(_) => {
307                                    failed = true;
308                                    break;
309                                }
310                            }
311                        }
312                        Err(_) => console_error!(
313                            ctx.status_tx,
314                            "Diameter handler action refused decision=fail_closed_action"
315                        ),
316                    }
317                    continue;
318                }
319            }
320        }
321        tokio::select! {
322         (result,half)=&mut reading=>{
323             let packet = match result {
324                 Ok(p) => p,
325                 Err(_) => {
326                     failed = true;
327                     break;
328                 }
329             };
330             stats(
331                 &ctx,
332                 packet.encode().map(|v| v.len()).unwrap_or(0) as u64,
333                 true,
334             )
335             .await;
336             reading = reader(half, cfg.io_timeout);
337             next_watchdog = Instant::now() + cfg.watchdog_interval;
338             if packet.is_request() {
339                 if !matches!(packet.command, DWR | DPR) || control(&packet, &peer).is_err() {
340                     failed = true;
341                     break;
342                 }
343                 let mut answer = packet.answer(2001);
344                 answer.origin(&cfg.identity);
345                 if send(&ctx, &mut write, &answer).await.is_err() {
346                     failed = true;
347                     break;
348                 }
349                 if packet.command == DPR {
350                     if let Some(v) = closing.take() {
351                         if let Some(command) = v.command {
352                             finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
353                         }
354                     }
355                     break;
356                 }
357                 continue;
358             }
359             if let Some(work) = closing.as_ref() {
360                 if packet.matches(&work.packet) && control(&packet, &peer).is_ok() {
361                     if let Some(command) = closing.take().unwrap().command {
362                         finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
363                     }
364                     break;
365                 }
366                 if pending.as_ref().is_some_and(|p| packet.matches(&p.packet)) {
367                     pending = None;
368                     continue;
369                 }
370             }
371             if watchdog.as_ref().is_some_and(|w| packet.matches(&w.0)) {
372                 if control(&packet, &peer).is_err() {
373                     failed = true;
374                     break;
375                 }
376                 watchdog = None;
377                 continue;
378             }
379             let Some(work) = pending.take() else {
380                 failed = true;
381                 break;
382             };
383             let response = match Response::from_packet(&packet, &work.packet, &peer) {
384                 Ok(v) => v,
385                 Err(_) => {
386                     pending = Some(work);
387                     failed = true;
388                     break;
389                 }
390             };
391             if let Some(command) = work.command {
392                 finish(
393                     &ctx,
394                     command,
395                     Ok(ClientSendOutcome::Executed {
396                         detail: format!("Diameter NASREQ result {}", response.result_code),
397                     }),
398                 )
399                 .await;
400             }
401             if events.len() >= MAX_QUEUED_EVENTS {
402                 failed = true;
403                 break;
404             }
405             events.push_back((
406                 Event::new(
407                     &actions::AA_EVENT,
408                     json!({"request":work.request,"reply":response}),
409                 ),
410                 work.depth,
411             ));
412         }
413         result=async{current.as_mut().unwrap().0.as_mut().await},if current.is_some()&&closing.is_none()=>{
414             let (_, depth) = current.take().unwrap();
415             match result {
416                 Ok(result) => {
417                     if result.actions.len() > MAX_HANDLER_ACTIONS
418                         || queued.len() + result.actions.len() > MAX_HANDLER_ACTIONS
419                         || result.actions.iter().any(|v| !within_json_budget(v))
420                     {
421                         for v in result.actions {
422                             crate::utils::json_budget::drop_iteratively(v);
423                         }
424                         failed = true;
425                         break;
426                     }
427                     if result
428                         .actions
429                         .iter()
430                         .any(|v| DiameterClientProtocol.execute_action(v.clone()).is_err())
431                     {
432                         for v in result.actions {
433                             crate::utils::json_budget::drop_iteratively(v);
434                         }
435                         console_error!(
436                             ctx.status_tx,
437                             "Diameter handler action invalid decision=fail_closed_action"
438                         );
439                         continue;
440                     }
441                     if result.actions.iter().any(|v| v["type"] == "disconnect") {
442                         for v in result.actions {
443                             crate::utils::json_budget::drop_iteratively(v);
444                         }
445                         match begin_close(&ctx, &cfg, &mut write, &mut pending, None).await {
446                             Ok(v) => closing = Some(v),
447                             Err(_) => {
448                                 failed = true;
449                                 break;
450                             }
451                         }
452                         continue;
453                     }
454                     if let Some(memory) = result.memory_updates {
455                         ctx.state.set_memory_for_client(ctx.client_id, memory).await;
456                     }
457                     queued.extend(result.actions.into_iter().map(|v| (v, depth + 1)));
458                 }
459                 Err(_) => console_error!(
460                     ctx.status_tx,
461                     "Diameter handler failed decision=fail_closed_handler"
462                 ),
463             }
464         }
465         command=commands.recv(),if closing.is_none()=>{
466             let Some(mut command) = command else {
467                 break;
468             };
469             if !within_json_budget(&command.action) {
470                 crate::utils::json_budget::drop_iteratively(std::mem::take(&mut command.action));
471                 finish(
472                     &ctx,
473                     command,
474                     Ok(ClientSendOutcome::Rejected {
475                         error: "Diameter JSON size/node/depth budget".into(),
476                     }),
477                 )
478                 .await;
479                 continue;
480             }
481             match prepared(command.action.clone(), 0) {
482                 Ok(None) => {
483                     drop(current.take());
484                     match begin_close(&ctx, &cfg, &mut write, &mut pending, Some(command)).await {
485                         Ok(v) => closing = Some(v),
486                         Err(_) => {
487                             failed = true;
488                             break 'peer;
489                         }
490                     }
491                 }
492                 Ok(Some(r)) if pending.is_none() => match request(r, &cfg, &peer, Some(command), 0) {
493                     Ok(work) => {
494                         if send(&ctx, &mut write, &work.packet).await.is_err() {
495                             pending = Some(work);
496                             failed = true;
497                             break;
498                         }
499                         pending = Some(work);
500                     }
501                     Err(_) => {
502                         failed = true;
503                         break;
504                     }
505                 },
506                 Ok(Some(_)) =>
507                     finish(
508                         &ctx,
509                         command,
510                         Ok(ClientSendOutcome::Rejected {
511                             error: "One Diameter AAA request already in flight".into(),
512                         }),
513                     )
514                     .await,
515                 Err(_) =>
516                     finish(
517                         &ctx,
518                         command,
519                         Ok(ClientSendOutcome::Rejected {
520                             error: "Diameter action invalid or over bound".into(),
521                         }),
522                     )
523                     .await,
524             }
525         }
526         _=tokio::time::sleep_until(pending.as_ref().map(|v|v.deadline).unwrap_or_else(Instant::now)),if pending.is_some()&&closing.is_none()=>{
527             failed = true;
528             break;
529         }
530         _=tokio::time::sleep_until(closing.as_ref().map(|v|v.deadline).unwrap_or_else(Instant::now)),if closing.is_some()=>{
531             failed = true;
532             break;
533         }
534         _=tokio::time::sleep_until(next_watchdog),if watchdog.is_none()&&closing.is_none()=>{
535             let mut p = match Packet::request(DWR, 0) {
536                 Ok(v) => v,
537                 Err(_) => {
538                     failed = true;
539                     break;
540                 }
541             };
542             p.origin(&cfg.identity);
543             if send(&ctx, &mut write, &p).await.is_err() {
544                 failed = true;
545                 break;
546             }
547             watchdog = Some((p, Instant::now() + cfg.io_timeout));
548         }
549         _=tokio::time::sleep_until(watchdog.as_ref().map(|v|v.1).unwrap_or_else(Instant::now)),if watchdog.is_some()&&closing.is_none()=>{
550             failed = true;
551             break;
552         }
553        }
554        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
555    }
556    // Dropping this retained read future closes the owned TCP half even when a
557    // partially read frame was interrupted. No detached reader or retry remains.
558    drop(reading);
559    drop(write);
560    drop(current.take());
561    for (v, _) in queued {
562        crate::utils::json_budget::drop_iteratively(v);
563    }
564    events.clear();
565    if let Some(p) = pending {
566        if let Some(command) = p.command {
567            finish(
568                &ctx,
569                command,
570                Err(anyhow::anyhow!("Diameter operation failed or cancelled")),
571            )
572            .await;
573        }
574    }
575    if let Some(v) = closing {
576        if let Some(command) = v.command {
577            finish(
578                &ctx,
579                command,
580                Err(anyhow::anyhow!("Diameter disconnect failed")),
581            )
582            .await;
583        }
584    }
585    ctx.state.remove_client_handle(ctx.client_id).await;
586    ctx.state
587        .update_client_status(ctx.client_id, ClientStatus::Disconnected)
588        .await;
589    if failed {
590        console_error!(
591            ctx.status_tx,
592            "Diameter peer failed decision=fail_closed_peer_session"
593        );
594        if let Ok(result) = handler(
595            ctx.clone(),
596            Event::new(
597                &actions::ERROR_EVENT,
598                json!({"error":"Diameter transport, framing, correlation or deadline failure; no replay"}),
599            ),
600            cfg.handler_timeout,
601        )
602        .await
603        {
604            for v in result.actions {
605                crate::utils::json_budget::drop_iteratively(v);
606            }
607        }
608    }
609    let _ = ctx.status_tx.send("__UPDATE_UI__".into());
610}
netget/client/diameter/mod.rs

netget/client/diameter/
mod.rs
