use std::collections::{HashMap, HashSet, VecDeque}; use std::io; use std::net::IpAddr; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize}; use std::task::{Context, Poll, Waker}; use std::time::Instant; use bytes::{Bytes, BytesMut}; use parking_lot::Mutex; use tokio::io::ReadBuf; use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use crate::config::{WebCarrier, WebLimitsConfig, WebRuntimeProfile, WebTimeoutsConfig}; use crate::web::frame::{self, FrameType}; use crate::web::manager::{ CarrierClientClass, CarrierLearningContext, ProfileKey, TokenHash, WebProcessRuntime, }; // Backend tasks own generation admission and authenticated MTProxy relay lifetimes. mod backend; // Downlink queues own cursor replay, flow control, and memory reservations. mod downlink; // Response ownership keeps detached batches charged until the last body clone drops. mod resident; // Read-only control-plane snapshots stay isolated from carrier operations. mod status; pub(crate) use status::WebSessionStatus; // Lane carrier state isolates request sequencing and downlink replay per logical stream. mod lanes; // Lane batch staging transfers queue ownership without escaping process budgets. mod lane_downlink; // Lane uplink creation remains transactional across validation and queue reservations. mod lane_uplink; // WebSocket carrier state owns pre-OPEN lane reservations and failure isolation. mod websocket; pub(crate) use websocket::WebSocketLaneReservation; pub(crate) use websocket::WebSocketProbeReservation; // Carrier commit and health evidence share one session-locked state machine. mod negotiation; // Session closure and carrier-attempt transitions share one cancellation boundary. mod lifecycle; // Uplink batches own exactly-once sequencing and client-frame validation. mod uplink; /// Conservative allocator and container overhead charged to every queued item. pub(crate) const QUEUE_ITEM_COST: usize = 256; #[derive(Clone, Copy, PartialEq, Eq)] enum PendingClass { Uplink, Downlink, Control, } struct InboundChunk { bytes: Bytes, offset: usize, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct StreamIdentity { pub(crate) id: u32, pub(crate) instance: u64, } /// Exact server-local identity of one carrier-lane incarnation. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct CarrierLaneIdentity { /// Numeric lane identifier carried on the wire. pub(crate) lane_id: u32, /// Monotonic server-local incarnation of that numeric lane. pub(crate) instance: u64, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct WebSocketLaneClaim { lane: CarrierLaneIdentity, peer_port: u16, connection_id: Option, } struct StreamState { instance: u64, inbound: VecDeque, receive_window: u32, send_credit: u64, read_waker: Option, write_waker: Option, } struct QueuedFrame { encoded: BytesMut, frame_type: FrameType, stream_id: u32, control: bool, cost: usize, } struct DownBatch { body: Bytes, lease: Arc, base_cursor: u64, next_cursor: u64, data_bytes: usize, data_items: usize, control_bytes: usize, control_items: usize, carrier_health_eligible: bool, } struct CarrierLane { instance: u64, pending_bytes: usize, pending_items: usize, resident: Arc, pending_frames: VecDeque, pending_windows: HashMap, unacked: Option, down_cursor: u64, down_epoch: u64, last_up_sequence: u64, last_up_digest: TokenHash, up_active: bool, notify: Arc, } impl CarrierLane { fn new(instance: u64) -> Self { Self { instance, pending_bytes: 0, pending_items: 0, resident: Arc::new(resident::ResidentCounters::default()), pending_frames: VecDeque::new(), pending_windows: HashMap::new(), unacked: None, down_cursor: 0, down_epoch: 0, last_up_sequence: 0, last_up_digest: [0; 32], up_active: false, notify: Arc::new(Notify::new()), } } } struct SessionState { streams: HashMap, closing_streams: HashMap, next_stream_instance: u64, active_peer_ports: HashSet, closed_streams: HashSet, closed_order: VecDeque, pending_frames: VecDeque, pending_windows: HashMap, unacked: Option, down_cursor: u64, down_epoch: u64, last_up_sequence: u64, last_up_digest: TokenHash, carrier_lanes: HashMap, lane_open_waits: usize, next_lane_instance: u64, websocket_lane_reservations: HashMap, pending_bytes: usize, pending_items: usize, pending_control_bytes: usize, pending_control_items: usize, last_activity: Instant, negotiation_phase: SessionNegotiationPhase, carrier_health_due_at: Option, carrier_health_activity_at: Option, carrier_health_uplink: bool, carrier_health_downlink: bool, carrier_commit_published: bool, carrier_health_reported: bool, websocket_carrier_active: bool, websocket_commit_ack_pending: bool, websocket_commit_ack_owner: Option, websocket_commit_ack_written: bool, websocket_probe_claimed: bool, close_requested: bool, closed: bool, } #[derive(Clone, Copy, PartialEq, Eq)] enum SessionNegotiationPhase { Uncommitted, Replacing, Committed, Superseded, } /// One bounded WEB carrier session containing logical MTProxy streams. pub(crate) struct WebSession { manager: std::sync::Weak, token_hash: TokenHash, client_ip: IpAddr, trace_session_id: u64, profile: Arc, profile_key: ProfileKey, selected_carrier: WebCarrier, carrier_attempt: u8, bootstrap_hash: TokenHash, carrier_deadline_at: Option, carrier_class: CarrierClientClass, learning_context: Option, automatic_carrier: bool, created_at: Instant, limits: WebLimitsConfig, timeouts: WebTimeoutsConfig, state: Mutex, down_notify: Arc, lane_open_notify: Arc, cancel: CancellationToken, tasks_live: AtomicUsize, tasks_done: Arc, resident: Arc, finished: AtomicBool, up_active: AtomicBool, } /// One successful downlink poll result. pub(crate) struct PollResult { /// Encoded downlink frame batch, or an empty long-poll result. pub(crate) body: Bytes, /// Cursor the client must present on its next downlink request. pub(crate) next_cursor: u64, /// Indicates that a drained non-zero lane no longer needs polling. pub(crate) lane_closed: bool, } impl WebSession { #[allow(clippy::too_many_arguments)] /// Creates one carrier session with immutable ownership and allocation policy. pub(crate) fn new( manager: std::sync::Weak, token_hash: TokenHash, client_ip: IpAddr, trace_session_id: u64, profile: Arc, profile_key: ProfileKey, selected_carrier: WebCarrier, carrier_attempt: u8, bootstrap_hash: TokenHash, carrier_deadline_at: Option, carrier_class: CarrierClientClass, learning_context: Option, automatic_carrier: bool, limits: WebLimitsConfig, timeouts: WebTimeoutsConfig, ) -> Arc { let mut carrier_lanes = HashMap::new(); let mut next_lane_instance = 1; if selected_carrier == WebCarrier::HttpsLanes { carrier_lanes.insert(0, CarrierLane::new(next_lane_instance)); next_lane_instance += 1; } Arc::new(Self { manager, token_hash, client_ip, trace_session_id, profile, profile_key, selected_carrier, carrier_attempt, bootstrap_hash, carrier_deadline_at, carrier_class, learning_context, automatic_carrier, created_at: Instant::now(), limits, timeouts, state: Mutex::new(SessionState { streams: HashMap::new(), closing_streams: HashMap::new(), next_stream_instance: 1, active_peer_ports: HashSet::new(), closed_streams: HashSet::new(), closed_order: VecDeque::new(), pending_frames: VecDeque::new(), pending_windows: HashMap::new(), unacked: None, down_cursor: 0, down_epoch: 0, last_up_sequence: 0, last_up_digest: [0; 32], carrier_lanes, lane_open_waits: 0, next_lane_instance, websocket_lane_reservations: HashMap::new(), pending_bytes: 0, pending_items: 0, pending_control_bytes: 0, pending_control_items: 0, last_activity: Instant::now(), negotiation_phase: SessionNegotiationPhase::Uncommitted, carrier_health_due_at: None, carrier_health_activity_at: None, carrier_health_uplink: false, carrier_health_downlink: false, carrier_commit_published: false, carrier_health_reported: false, websocket_carrier_active: false, websocket_commit_ack_pending: false, websocket_commit_ack_owner: None, websocket_commit_ack_written: false, websocket_probe_claimed: false, close_requested: false, closed: false, }), down_notify: Arc::new(Notify::new()), lane_open_notify: Arc::new(Notify::new()), cancel: CancellationToken::new(), tasks_live: AtomicUsize::new(0), tasks_done: Arc::new(Notify::new()), resident: Arc::new(resident::ResidentCounters::default()), finished: AtomicBool::new(false), up_active: AtomicBool::new(false), }) } /// Returns the stable hashed token identity without exposing the credential. pub(crate) fn token_hash(&self) -> TokenHash { self.token_hash } /// Checks the canonical virtual host that owns this bearer session. pub(crate) fn matches_host(&self, host: &str) -> bool { self.profile.host == host } /// Returns the immutable carrier selected when this session was created. pub(crate) fn carrier(&self) -> WebCarrier { self.selected_carrier } /// Returns the stable quota owner without exposing profile credentials. pub(crate) fn profile_key(&self) -> ProfileKey { self.profile_key } /// Returns the process-unique non-secret trace identifier. pub(crate) fn trace_session_id(&self) -> u64 { self.trace_session_id } /// Returns the immutable carrier-attempt incarnation number. pub(crate) fn carrier_attempt(&self) -> u8 { self.carrier_attempt } /// Creates a child cancellation boundary for one owned carrier task. pub(crate) fn carrier_cancellation(&self) -> CancellationToken { self.cancel.child_token() } /// Returns a cloned non-secret identity only for enabled debug capture. pub(crate) fn trace_identity(&self) -> crate::web::trace::TraceIdentity { crate::web::trace::TraceIdentity::from_profile(self.trace_session_id, &self.profile) } /// Records one typed lifecycle event without exposing session credentials. pub(super) fn trace_lifecycle( &self, event: crate::web::trace::TraceLifecycleEvent, stream_id: Option, reason: Option<&'static str>, ) { if let Some(manager) = self.manager.upgrade() { manager.trace().record_profile_lifecycle( self.client_ip, Some(self.trace_session_id), &self.profile, event, stream_id, reason, ); } } /// Returns the immutable limits frozen when this carrier chain was created. pub(crate) fn limits(&self) -> &WebLimitsConfig { &self.limits } /// Returns the immutable timeouts frozen when this carrier chain was created. pub(crate) fn timeouts(&self) -> &WebTimeoutsConfig { &self.timeouts } /// Polls client-to-server bytes and returns consumed flow-control credit. pub(super) fn poll_read( &self, stream: StreamIdentity, cx: &mut Context<'_>, output: &mut ReadBuf<'_>, ) -> Poll> { let mut state = self.state.lock(); let (count, finished) = { let Some(stream_state) = state .streams .get_mut(&stream.id) .filter(|state| state.instance == stream.instance) else { return Poll::Ready(Ok(())); }; let Some(chunk) = stream_state.inbound.front_mut() else { stream_state.read_waker = Some(cx.waker().clone()); return Poll::Pending; }; let available = &chunk.bytes[chunk.offset..]; let count = available.len().min(output.remaining()); output.put_slice(&available[..count]); chunk.offset += count; let finished = chunk.offset == chunk.bytes.len(); if finished { stream_state.inbound.pop_front(); } stream_state.receive_window = stream_state.receive_window.saturating_add(count as u32); (count, finished) }; let overhead = if finished { QUEUE_ITEM_COST } else { 0 }; self.release_locked(&mut state, count + overhead, usize::from(finished), false); if !self.queue_window_locked(&mut state, stream.id, count as u32) { drop(state); self.close(); return Poll::Ready(Err(io::Error::other( "WEB session control budget exhausted", ))); } Poll::Ready(Ok(())) } /// Polls server-to-client writes against stream credit and bounded queues. pub(super) fn poll_write( &self, stream: StreamIdentity, cx: &mut Context<'_>, input: &[u8], ) -> Poll> { if input.is_empty() { return Poll::Ready(Ok(0)); } let mut state = self.state.lock(); let Some(stream_state) = state .streams .get_mut(&stream.id) .filter(|state| state.instance == stream.instance) else { return Poll::Ready(Err(io::Error::new( io::ErrorKind::BrokenPipe, "WEB logical stream is closed", ))); }; let count = input .len() .min(frame::DATA_CHUNK_BYTES) .min(self.limits.max_frame_payload_bytes) .min(if self.carrier().uses_lanes() { self.limits .pending_bytes_per_lane .saturating_sub(frame::HEADER_BYTES + QUEUE_ITEM_COST) } else { usize::MAX }) .min(stream_state.send_credit as usize); if count == 0 { stream_state.write_waker = Some(cx.waker().clone()); return Poll::Pending; } if !self.queue_data_locked(&mut state, stream.id, &input[..count]) { if let Some(stream_state) = state .streams .get_mut(&stream.id) .filter(|state| state.instance == stream.instance) { stream_state.write_waker = Some(cx.waker().clone()); } return Poll::Pending; } let Some(stream_state) = state .streams .get_mut(&stream.id) .filter(|state| state.instance == stream.instance) else { return Poll::Ready(Err(io::Error::new( io::ErrorKind::BrokenPipe, "WEB logical stream is closed", ))); }; stream_state.send_credit -= count as u64; state.last_activity = Instant::now(); drop(state); if self.carrier().is_multiplexed() { self.down_notify.notify_waiters(); } Poll::Ready(Ok(count)) } /// Returns the process queue-capacity notification source while the manager lives. pub(super) fn budget_notify(&self) -> Option> { self.manager .upgrade() .map(|manager| manager.budget_notify()) } } fn inbound_queue_cost(queue: &VecDeque) -> (usize, usize) { let bytes = queue.iter().fold(0usize, |total, chunk| { total.saturating_add(chunk.bytes.len().saturating_sub(chunk.offset) + QUEUE_ITEM_COST) }); (bytes, queue.len()) } fn remember_closed(state: &mut SessionState, stream_id: u32, limit: usize) -> Option { if !state.closed_streams.insert(stream_id) { return None; } state.closed_order.push_back(stream_id); let mut evicted = None; while state.closed_order.len() > limit { if let Some(oldest) = state.closed_order.pop_front() { state.closed_streams.remove(&oldest); evicted = Some(oldest); } } evicted } fn insert_carrier_lane(state: &mut SessionState, lane_id: u32) -> Option { if state.carrier_lanes.contains_key(&lane_id) { return None; } let instance = state.next_lane_instance; state.next_lane_instance = instance.checked_add(1)?; state .carrier_lanes .insert(lane_id, CarrierLane::new(instance)); Some(CarrierLaneIdentity { lane_id, instance }) }