WEB

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
Co-Authored-By: John Preston <17900494+john-preston@users.noreply.github.com>

src/web/session/uplink.rs

@@ -0,0 +1,431 @@
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::time::Instant;
+
+use bytes::Bytes;
+use sha2::{Digest, Sha256};
+use subtle::ConstantTimeEq;
+
+use super::{
+    InboundChunk, PendingClass, QUEUE_ITEM_COST, SessionState, StreamState, WebSession,
+    inbound_queue_cost, remember_closed,
+};
+use crate::web::frame::{self, Frame, FrameType};
+use crate::web::manager::{ManagerError, TokenHash};
+
+impl WebSession {
+    /// Applies one exactly-once uplink batch.
+    pub(crate) fn process_up(
+        self: &Arc<Self>,
+        sequence: u64,
+        body: &[u8],
+    ) -> Result<u64, ManagerError> {
+        if self
+            .up_active
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_err()
+        {
+            return Err(ManagerError::Concurrent);
+        }
+        let _uplink = UplinkGuard(&self.up_active);
+        let frames = match frame::parse_all(body, &self.limits) {
+            Ok(frames) => frames,
+            Err(_) => {
+                self.close();
+                return Err(ManagerError::Protocol);
+            }
+        };
+        if frames
+            .iter()
+            .copied()
+            .any(|value| frame::validate_client_shape(value).is_err())
+        {
+            self.close();
+            return Err(ManagerError::Protocol);
+        }
+        let digest: TokenHash = Sha256::digest(body).into();
+        let mut opened = Vec::new();
+        let result = {
+            let mut state = self.state.lock();
+            if state.closed {
+                return Err(ManagerError::Closed);
+            }
+            state.last_activity = Instant::now();
+            if sequence == state.last_up_sequence && sequence != 0 {
+                return if bool::from(state.last_up_digest.ct_eq(&digest)) {
+                    Ok(sequence)
+                } else {
+                    drop(state);
+                    self.close();
+                    Err(ManagerError::Protocol)
+                };
+            }
+            if sequence == 0 || sequence != state.last_up_sequence.saturating_add(1) {
+                drop(state);
+                self.close();
+                return Err(ManagerError::Protocol);
+            }
+            if !validate_batch(&state, &frames) {
+                drop(state);
+                self.close();
+                return Err(ManagerError::Protocol);
+            }
+            let (reserve_bytes, reserve_items) = inbound_reservation(&state, &frames);
+            if !self.reserve_locked(
+                &mut state,
+                reserve_bytes,
+                reserve_items,
+                PendingClass::Uplink,
+            ) {
+                return Err(ManagerError::Backpressure);
+            }
+            let mut unused_bytes = reserve_bytes;
+            let mut unused_items = reserve_items;
+            let applied = self.apply_batch_locked(
+                &mut state,
+                &frames,
+                &mut opened,
+                &mut unused_bytes,
+                &mut unused_items,
+            );
+            self.release_locked(&mut state, unused_bytes, unused_items, false);
+            if !applied {
+                Err(ManagerError::Closed)
+            } else {
+                state.last_up_sequence = sequence;
+                state.last_up_digest = digest;
+                Ok(sequence)
+            }
+        };
+        if matches!(result, Err(ManagerError::Backpressure)) {
+            return result;
+        }
+        if result.is_err() {
+            self.close();
+            for (_, peer_port) in opened {
+                self.release_stream_reservation(peer_port);
+            }
+            return result;
+        }
+        for (stream_id, peer_port) in opened {
+            self.spawn_stream(stream_id, peer_port);
+        }
+        if let Some(manager) = self.manager.upgrade() {
+            manager.record_up(body.len());
+        }
+        result
+    }
+
+    fn apply_batch_locked(
+        &self,
+        state: &mut SessionState,
+        frames: &[Frame<'_>],
+        opened: &mut Vec<(u32, u16)>,
+        unused_bytes: &mut usize,
+        unused_items: &mut usize,
+    ) -> bool {
+        for value in frames {
+            if value.stream_id == 0 {
+                continue;
+            }
+            let was_closed = state.closed_streams.contains(&value.stream_id);
+            match value.frame_type {
+                FrameType::Open => {
+                    let Some(peer_port) = self.reserve_stream_locked(state) else {
+                        remember_closed(
+                            state,
+                            value.stream_id,
+                            self.limits.max_tombstones_per_session,
+                        );
+                        if !self.queue_control_locked(
+                            state,
+                            FrameType::Close,
+                            value.stream_id,
+                            &[],
+                        ) {
+                            return false;
+                        }
+                        continue;
+                    };
+                    state.streams.insert(
+                        value.stream_id,
+                        StreamState {
+                            inbound: VecDeque::new(),
+                            receive_window: frame::INITIAL_STREAM_WINDOW,
+                            send_credit: u64::from(frame::INITIAL_STREAM_WINDOW),
+                            read_waker: None,
+                            write_waker: None,
+                        },
+                    );
+                    opened.push((value.stream_id, peer_port));
+                }
+                FrameType::Data if !was_closed => {
+                    let Some(stream) = state.streams.get_mut(&value.stream_id) else {
+                        return false;
+                    };
+                    stream.receive_window -= value.payload.len() as u32;
+                    stream.inbound.push_back(InboundChunk {
+                        bytes: Bytes::copy_from_slice(value.payload),
+                        offset: 0,
+                    });
+                    *unused_bytes = unused_bytes
+                        .saturating_sub(value.payload.len() + QUEUE_ITEM_COST);
+                    *unused_items = unused_items.saturating_sub(1);
+                    if let Some(waker) = stream.read_waker.take() {
+                        waker.wake();
+                    }
+                }
+                FrameType::Window if !was_closed => {
+                    let Some(stream) = state.streams.get_mut(&value.stream_id) else {
+                        return false;
+                    };
+                    let amount = frame::window_amount(value.payload).unwrap_or(0);
+                    stream.send_credit = stream
+                        .send_credit
+                        .saturating_add(u64::from(amount))
+                        .min(u64::from(u32::MAX));
+                    if let Some(waker) = stream.write_waker.take() {
+                        waker.wake();
+                    }
+                }
+                FrameType::Close if !was_closed => {
+                    let Some(stream) = state.streams.remove(&value.stream_id) else {
+                        return false;
+                    };
+                    let (bytes, items) = inbound_queue_cost(&stream.inbound);
+                    self.release_locked(state, bytes, items, false);
+                    remember_closed(
+                        state,
+                        value.stream_id,
+                        self.limits.max_tombstones_per_session,
+                    );
+                    if let Some(waker) = stream.read_waker {
+                        waker.wake();
+                    }
+                    if let Some(waker) = stream.write_waker {
+                        waker.wake();
+                    }
+                }
+                FrameType::Data | FrameType::Window | FrameType::Close => {}
+                _ => return false,
+            }
+        }
+        true
+    }
+
+    fn reserve_stream_locked(&self, state: &mut SessionState) -> Option<u16> {
+        if state.active_peer_ports.len() >= self.profile.max_streams_per_session {
+            return None;
+        }
+        let manager = self.manager.upgrade()?;
+        let peer_port = manager.try_acquire_stream(
+            self.profile_key,
+            self.profile.max_streams,
+            self.client_ip,
+            self.profile.public_addr,
+        )?;
+        if state.active_peer_ports.insert(peer_port) {
+            return Some(peer_port);
+        }
+        manager.release_stream(
+            self.profile_key,
+            self.client_ip,
+            self.profile.public_addr,
+            peer_port,
+        );
+        None
+    }
+}
+
+struct UplinkGuard<'a>(&'a AtomicBool);
+
+impl Drop for UplinkGuard<'_> {
+    fn drop(&mut self) {
+        self.0.store(false, Ordering::Release);
+    }
+}
+
+fn validate_batch(state: &SessionState, frames: &[Frame<'_>]) -> bool {
+    let mut live = state
+        .streams
+        .iter()
+        .map(|(id, stream)| (*id, (stream.receive_window, stream.send_credit)))
+        .collect::<HashMap<_, _>>();
+    let mut closed = HashSet::new();
+    for value in frames {
+        if value.stream_id == 0 {
+            if value.frame_type != FrameType::Pong {
+                return false;
+            }
+            continue;
+        }
+        let was_closed = state.closed_streams.contains(&value.stream_id)
+            || closed.contains(&value.stream_id);
+        match value.frame_type {
+            FrameType::Open => {
+                if live.contains_key(&value.stream_id) || was_closed {
+                    return false;
+                }
+                live.insert(
+                    value.stream_id,
+                    (
+                        frame::INITIAL_STREAM_WINDOW,
+                        u64::from(frame::INITIAL_STREAM_WINDOW),
+                    ),
+                );
+            }
+            FrameType::Data if !was_closed => {
+                let Some((receive_window, send_credit)) = live.get_mut(&value.stream_id) else {
+                    return false;
+                };
+                let Ok(payload_len) = u32::try_from(value.payload.len()) else {
+                    return false;
+                };
+                if payload_len > *receive_window {
+                    return false;
+                }
+                *receive_window -= payload_len;
+                let _ = send_credit;
+            }
+            FrameType::Window if !was_closed => {
+                let Some((_, send_credit)) = live.get_mut(&value.stream_id) else {
+                    return false;
+                };
+                let Ok(amount) = frame::window_amount(value.payload) else {
+                    return false;
+                };
+                *send_credit = send_credit
+                    .saturating_add(u64::from(amount))
+                    .min(u64::from(u32::MAX));
+            }
+            FrameType::Close if !was_closed => {
+                if live.remove(&value.stream_id).is_none() {
+                    return false;
+                }
+                closed.insert(value.stream_id);
+            }
+            FrameType::Data | FrameType::Window | FrameType::Close => {}
+            _ => return false,
+        }
+    }
+    true
+}
+
+fn inbound_reservation(state: &SessionState, frames: &[Frame<'_>]) -> (usize, usize) {
+    let mut live = state.streams.keys().copied().collect::<HashSet<_>>();
+    let mut bytes = 0usize;
+    let mut items = 0usize;
+    for value in frames {
+        match value.frame_type {
+            FrameType::Open => {
+                live.insert(value.stream_id);
+            }
+            FrameType::Data if live.contains(&value.stream_id) => {
+                bytes = bytes.saturating_add(value.payload.len() + QUEUE_ITEM_COST);
+                items = items.saturating_add(1);
+            }
+            FrameType::Close => {
+                live.remove(&value.stream_id);
+            }
+            _ => {}
+        }
+    }
+    (bytes, items)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::SocketAddr;
+
+    use crate::config::{
+        WebLimitsConfig, WebRuntimeProfile, WebSecretMode, WebTimeoutsConfig,
+    };
+    use crate::web::manager::WebProcessRuntime;
+
+    fn session() -> Arc<WebSession> {
+        let profile = Arc::new(WebRuntimeProfile {
+            host: "proxy.example.com".to_string(),
+            public_addr: SocketAddr::from(([203, 0, 113, 10], 443)),
+            user: "alice".to_string(),
+            secret_mode: WebSecretMode::Plain,
+            capability: [0; 32],
+            max_sessions: 1,
+            max_streams: 1,
+            max_streams_per_session: 1,
+        });
+        WebSession::new(
+            std::sync::Weak::<WebProcessRuntime>::new(),
+            [1; 32],
+            "192.0.2.10".parse().unwrap(),
+            profile,
+            [2; 32],
+            WebLimitsConfig::default(),
+            WebTimeoutsConfig::default(),
+        )
+    }
+
+    #[test]
+    fn uplink_retry_commits_only_one_exact_body() {
+        let session = session();
+        let first = frame::encode(FrameType::Pong, 0, &[1, 2, 3]);
+        assert_eq!(session.process_up(1, &first), Ok(1));
+        assert_eq!(session.process_up(1, &first), Ok(1));
+
+        let changed = frame::encode(FrameType::Pong, 0, &[1, 2, 4]);
+        assert_eq!(session.process_up(1, &changed), Err(ManagerError::Protocol));
+        assert!(session.state.lock().closed);
+    }
+
+    #[test]
+    fn concurrent_uplink_does_not_commit_sequence() {
+        let session = session();
+        let body = frame::encode(FrameType::Pong, 0, &[]);
+        session.up_active.store(true, Ordering::Release);
+        assert_eq!(
+            session.process_up(1, &body),
+            Err(ManagerError::Concurrent)
+        );
+        assert_eq!(session.state.lock().last_up_sequence, 0);
+        session.up_active.store(false, Ordering::Release);
+        assert_eq!(session.process_up(1, &body), Ok(1));
+    }
+
+    #[test]
+    fn backpressured_uplink_does_not_commit_or_close() {
+        let session = session();
+        {
+            let mut state = session.state.lock();
+            state.streams.insert(
+                1,
+                StreamState {
+                    inbound: VecDeque::new(),
+                    receive_window: frame::INITIAL_STREAM_WINDOW,
+                    send_credit: u64::from(frame::INITIAL_STREAM_WINDOW),
+                    read_waker: None,
+                    write_waker: None,
+                },
+            );
+            state.pending_bytes = session.limits.pending_bytes_per_session;
+        }
+        let body = frame::encode(FrameType::Data, 1, &[1]);
+
+        assert_eq!(
+            session.process_up(1, &body),
+            Err(ManagerError::Backpressure)
+        );
+        let state = session.state.lock();
+        assert!(!state.closed);
+        assert_eq!(state.last_up_sequence, 0);
+        assert!(state.streams.get(&1).unwrap().inbound.is_empty());
+    }
+
+    #[test]
+    fn uplink_gap_is_fatal() {
+        let session = session();
+        let body = frame::encode(FrameType::Pong, 0, &[]);
+        assert_eq!(session.process_up(2, &body), Err(ManagerError::Protocol));
+        assert!(session.state.lock().closed);
+    }
+}