cloudpub_common/
routing.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::time::{Duration, Instant};
4use tracing::{debug, info};
5
6#[derive(Debug, Clone)]
7pub struct UdpSession<T> {
8    pub external_addr: SocketAddr,
9    pub local_data: T,
10    pub last_activity: Instant,
11}
12
13pub struct UdpRoutingTable<T> {
14    // Maps external addresses to session info
15    external_to_session: HashMap<SocketAddr, UdpSession<T>>,
16    // Maps local addresses to external addresses (for client-side routing)
17    local_to_external: HashMap<SocketAddr, SocketAddr>,
18    // Session timeout
19    session_timeout: Duration,
20}
21
22impl<T: Clone> UdpRoutingTable<T> {
23    pub fn new(session_timeout_secs: u64) -> Self {
24        Self {
25            external_to_session: HashMap::new(),
26            local_to_external: HashMap::new(),
27            session_timeout: Duration::from_secs(session_timeout_secs),
28        }
29    }
30
31    pub fn register_session(&mut self, external_addr: SocketAddr, local_data: T) {
32        let session = UdpSession {
33            external_addr,
34            local_data,
35            last_activity: Instant::now(),
36        };
37
38        self.external_to_session.insert(external_addr, session);
39        debug!("Registered UDP session for {}", external_addr);
40    }
41
42    pub fn register_client_session(
43        &mut self,
44        external_addr: SocketAddr,
45        local_addr: SocketAddr,
46        local_data: T,
47    ) {
48        let session = UdpSession {
49            external_addr,
50            local_data,
51            last_activity: Instant::now(),
52        };
53
54        self.external_to_session.insert(external_addr, session);
55        self.local_to_external.insert(local_addr, external_addr);
56
57        debug!(
58            "Registered UDP client session: {} <-> {}",
59            external_addr, local_addr
60        );
61    }
62
63    pub fn update_activity(&mut self, addr: &SocketAddr, is_external: bool) {
64        if is_external {
65            if let Some(session) = self.external_to_session.get_mut(addr) {
66                session.last_activity = Instant::now();
67            }
68        } else {
69            // Update via local address
70            if let Some(external_addr) = self.local_to_external.get(addr) {
71                if let Some(session) = self.external_to_session.get_mut(external_addr) {
72                    session.last_activity = Instant::now();
73                }
74            }
75        }
76    }
77
78    pub fn get_session(&self, external_addr: &SocketAddr) -> Option<&UdpSession<T>> {
79        self.external_to_session.get(external_addr)
80    }
81
82    pub fn get_local_data(&self, external_addr: &SocketAddr) -> Option<&T> {
83        self.external_to_session
84            .get(external_addr)
85            .map(|s| &s.local_data)
86    }
87
88    pub fn get_external_addr(&self, local_addr: &SocketAddr) -> Option<SocketAddr> {
89        self.local_to_external.get(local_addr).copied()
90    }
91
92    pub fn cleanup_expired_sessions(&mut self) -> usize {
93        let now = Instant::now();
94
95        let mut expired_external: Vec<SocketAddr> = Vec::new();
96        let mut expired_local: Vec<SocketAddr> = Vec::new();
97
98        for (external_addr, session) in &self.external_to_session {
99            if now.duration_since(session.last_activity) > self.session_timeout {
100                expired_external.push(*external_addr);
101                // Find corresponding local addresses to clean up
102                for (local_addr, ext_addr) in &self.local_to_external {
103                    if *ext_addr == *external_addr {
104                        expired_local.push(*local_addr);
105                    }
106                }
107            }
108        }
109
110        for addr in &expired_external {
111            self.external_to_session.remove(addr);
112            debug!("Cleaned up expired UDP session for external {}", addr);
113        }
114
115        for addr in &expired_local {
116            self.local_to_external.remove(addr);
117        }
118
119        let cleaned = expired_external.len();
120        if cleaned > 0 {
121            info!("Cleaned up {} expired UDP sessions", cleaned);
122        }
123        cleaned
124    }
125
126    pub fn session_count(&self) -> usize {
127        self.external_to_session.len()
128    }
129}
cloudpub_common/routing.rs

cloudpub_common/
routing.rs
