@@ -8,16 +8,16 @@ use std::sync::Arc;
88use std:: sync:: atomic:: Ordering ;
99
1010use bytes:: Bytes ;
11+ use tokio:: sync:: OwnedSemaphorePermit ;
1112use tokio_util:: sync:: CancellationToken ;
1213
13- use crate :: common:: udp_dial_timeout;
1414use crate :: protocol:: {
15- DATAGRAM_UDP_CLOSE , DATAGRAM_UDP_REQUEST , DATAGRAM_UDP_RESPONSE , decode_udp_datagram ,
15+ DATAGRAM_UDP_CLOSE , DATAGRAM_UDP_REQUEST , DATAGRAM_UDP_RESPONSE , decode_udp_datagram_parts ,
1616 new_udp_datagram_header,
1717} ;
1818
1919use super :: PortalSession ;
20- use super :: flow:: { PortalUdpFlow , UdpFlowKey } ;
20+ use super :: flow:: { PortalUdpFlow , QueuedDatagram , UdpFlowKey } ;
2121
2222impl PortalSession {
2323 /// Consumes pending and live QUIC datagrams for this authenticated session.
@@ -48,26 +48,35 @@ impl PortalSession {
4848 }
4949
5050 async fn handle_datagram ( self : & Arc < Self > , data : Bytes ) {
51- let ( frame_type , flow_id , target_addr , payload ) =
52- match decode_udp_datagram ( & data , & self . portal . credentials . protocol_spec ) {
53- Ok ( decoded) => decoded,
54- Err ( err) => {
55- self . portal . logger . debug ( format_args ! (
56- "portal::conn::datagram_loop: failed to decode datagram: {err}"
57- ) ) ;
58- return ;
59- }
60- } ;
51+ let decoded = match decode_udp_datagram_parts ( & data , & self . portal . credentials . protocol_spec )
52+ {
53+ Ok ( decoded) => decoded,
54+ Err ( err) => {
55+ self . portal . logger . debug ( format_args ! (
56+ "portal::conn::datagram_loop: failed to decode datagram: {err}"
57+ ) ) ;
58+ return ;
59+ }
60+ } ;
6161
62- if frame_type == DATAGRAM_UDP_CLOSE {
63- self . close_udp_flow ( flow_id, target_addr) . await ;
62+ if decoded. frame_type == DATAGRAM_UDP_CLOSE {
63+ self . close_udp_flow ( decoded. flow_id , decoded. target_addr )
64+ . await ;
6465 return ;
6566 }
66- if frame_type != DATAGRAM_UDP_REQUEST {
67+ if decoded . frame_type != DATAGRAM_UDP_REQUEST {
6768 return ;
6869 }
6970
70- self . handle_udp_request ( flow_id, target_addr, payload) . await ;
71+ let retained_bytes = data. len ( ) ;
72+ let payload = data. slice ( decoded. payload_offset ..) ;
73+ self . handle_udp_request (
74+ decoded. flow_id ,
75+ decoded. target_addr ,
76+ payload,
77+ retained_bytes,
78+ )
79+ . await ;
7180 }
7281
7382 async fn close_udp_flow ( & self , flow_id : u64 , target_addr : String ) {
@@ -82,87 +91,85 @@ impl PortalSession {
8291 self : & Arc < Self > ,
8392 flow_id : u64 ,
8493 target_addr : String ,
85- payload : & [ u8 ] ,
94+ payload : Bytes ,
95+ retained_bytes : usize ,
8696 ) {
87- let flow = match self . get_udp_flow ( flow_id, & target_addr) . await {
88- Ok ( flow) => flow,
89- Err ( err) => {
90- self . portal . logger . error ( format_args ! (
91- "portal::conn::handle_udp_request: failed to open UDP flow: {err}"
92- ) ) ;
93- return ;
94- }
97+ let Some ( budget) = self . try_reserve_udp_queue ( retained_bytes) else {
98+ self . warn_udp_drop ( "connection queue byte limit reached" ) ;
99+ return ;
95100 } ;
96- flow. touch ( ) . await ;
97-
98- if let Some ( limiter) = & self . portal . rate_limiter {
99- limiter. wait_read ( payload. len ( ) as i64 ) . await ;
100- }
101-
102- match flow. send_to_target ( payload) . await {
103- Ok ( n) => {
104- self . portal
105- . stats
106- . udp_rx
107- . fetch_add ( n as u64 , Ordering :: Relaxed ) ;
101+ let flow = match self . get_or_create_udp_flow ( flow_id, target_addr) . await {
102+ Ok ( Some ( flow) ) => flow,
103+ Ok ( None ) => {
104+ self . warn_udp_drop ( "per-connection flow limit reached" ) ;
105+ return ;
108106 }
109107 Err ( err) => {
110108 self . portal . logger . error ( format_args ! (
111- "portal::conn::handle_udp_request: failed to write target : {err}"
109+ "portal::conn::handle_udp_request: failed to create UDP flow : {err}"
112110 ) ) ;
113- flow . close ( ) . await ;
111+ return ;
114112 }
113+ } ;
114+ if !flow. enqueue ( QueuedDatagram :: new ( payload, budget) ) {
115+ self . warn_udp_drop ( "per-flow datagram queue is full" ) ;
115116 }
116117 }
117118
118- async fn get_udp_flow (
119+ async fn get_or_create_udp_flow (
119120 self : & Arc < Self > ,
120121 flow_id : u64 ,
121- target_addr : & str ,
122- ) -> anyhow:: Result < Arc < PortalUdpFlow > > {
122+ target_addr : String ,
123+ ) -> anyhow:: Result < Option < Arc < PortalUdpFlow > > > {
123124 let key = UdpFlowKey :: new ( flow_id, target_addr) ;
124-
125- // Reuse an open flow for repeated datagrams with the same client flow ID
126- // and target; closed flows are replaced below.
127- if let Some ( flow) = self . udp_flows . lock ( ) . await . get ( & key) . cloned ( )
128- && !flow. is_closed ( )
129- {
130- return Ok ( flow) ;
125+ let mut guard = self . udp_flows . lock ( ) . await ;
126+ if self . closed . load ( Ordering :: Acquire ) {
127+ return Ok ( None ) ;
128+ }
129+ if let Some ( flow) = guard. get ( & key) . cloned ( ) {
130+ if !flow. is_closed ( ) {
131+ return Ok ( Some ( flow) ) ;
132+ }
133+ guard. remove ( & key) ;
134+ }
135+ if guard. len ( ) >= self . portal . udp_flow_limits . max_flows {
136+ return Ok ( None ) ;
131137 }
132-
133- let socket = self
134- . portal
135- . outbound
136- . dial_udp ( target_addr, udp_dial_timeout ( ) )
137- . await ?;
138138 let response_header = new_udp_datagram_header (
139139 DATAGRAM_UDP_RESPONSE ,
140140 flow_id,
141- target_addr ,
141+ key . target ( ) ,
142142 & self . portal . credentials . protocol_spec ,
143143 )
144144 . map_err ( |e| {
145- anyhow:: anyhow!( "portal::conn::get_udp_flow: failed to build response header: {e}" )
145+ anyhow:: anyhow!(
146+ "portal::conn::get_or_create_udp_flow: failed to build response header: {e}"
147+ )
146148 } ) ?;
147- let flow = Arc :: new ( PortalUdpFlow :: new (
148- Arc :: downgrade ( self ) ,
149- key. clone ( ) ,
150- socket,
151- response_header,
152- ) ) ;
153-
154- let mut guard = self . udp_flows . lock ( ) . await ;
155- if let Some ( existing) = guard. get ( & key) . cloned ( ) {
156- // Another task may have inserted the flow while this task was
157- // dialing. Prefer the existing entry to keep ownership singular.
158- return Ok ( existing) ;
159- }
149+ let ( flow, receiver) = PortalUdpFlow :: new ( Arc :: downgrade ( self ) , key. clone ( ) ) ;
150+ let flow = Arc :: new ( flow) ;
160151 guard. insert ( key, flow. clone ( ) ) ;
152+ self . portal . stats . add_session ( true ) ;
161153 drop ( guard) ;
162154
163- self . portal . stats . add_session ( true ) ;
164- tokio:: spawn ( flow. clone ( ) . read_loop ( ) ) ;
165- tokio:: spawn ( flow. clone ( ) . idle_loop ( ) ) ;
166- Ok ( flow)
155+ tokio:: spawn ( flow. clone ( ) . run ( receiver, response_header) ) ;
156+ Ok ( Some ( flow) )
157+ }
158+
159+ fn try_reserve_udp_queue ( & self , bytes : usize ) -> Option < OwnedSemaphorePermit > {
160+ let permits = u32:: try_from ( bytes) . ok ( ) ?;
161+ self . udp_queue_budget
162+ . clone ( )
163+ . try_acquire_many_owned ( permits)
164+ . ok ( )
165+ }
166+
167+ fn warn_udp_drop ( & self , reason : & str ) {
168+ if !self . udp_overload_logged . swap ( true , Ordering :: AcqRel ) {
169+ self . portal . logger . warn ( format_args ! (
170+ "portal::conn::datagram_loop: dropping UDP datagrams for {}: {reason}" ,
171+ self . conn. remote_address( )
172+ ) ) ;
173+ }
167174 }
168175}
0 commit comments