Add regression and security tests for relay quota and TLS stream handling

- Introduced regression tests for relay quota wake liveness to ensure proper handling of contention and wake events.
- Added adversarial tests to validate the behavior of the quota system under stress and contention scenarios.
- Implemented security tests for the TLS stream to verify the preservation of pending plaintext during state transitions.
- Enhanced the pool writer tests to ensure proper quarantine behavior and validate the removal of writers from the registry.
- Included fuzz testing to assess the robustness of the quota and TLS handling mechanisms against unexpected inputs and states.

src/stream/tls_stream.rs

@@ -297,6 +297,11 @@ impl<R> FakeTlsReader<R> {
     pub fn into_inner_with_pending_plaintext(mut self) -> (R, Vec<u8>) {
         let pending = match std::mem::replace(&mut self.state, TlsReaderState::Idle) {
             TlsReaderState::Yielding { buffer } => buffer.as_slice().to_vec(),
+            TlsReaderState::ReadingBody { record_type, buffer, .. }
+                if record_type == TLS_RECORD_APPLICATION =>
+            {
+                buffer.to_vec()
+            }
             _ => Vec::new(),
         };
         (self.upstream, pending)
@@ -1293,3 +1298,7 @@ mod tests {
         assert_eq!(bytes, [0x17, 0x03, 0x03, 0x12, 0x34]);
     }
 }
+
+#[cfg(test)]
+#[path = "tls_stream_pending_plaintext_security_tests.rs"]
+mod pending_plaintext_security_tests;

src/stream/tls_stream_pending_plaintext_security_tests.rs

@@ -0,0 +1,143 @@
+use super::*;
+use bytes::{Bytes, BytesMut};
+
+#[test]
+fn reading_body_pending_application_plaintext_is_preserved_on_into_inner() {
+    let sample = b"coalesced-tail-after-mtproto";
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::ReadingBody {
+        record_type: TLS_RECORD_APPLICATION,
+        length: sample.len(),
+        buffer: BytesMut::from(&sample[..]),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert_eq!(
+        pending,
+        sample,
+        "partial application-data body must survive into fallback path"
+    );
+}
+
+#[test]
+fn yielding_pending_plaintext_is_preserved_on_into_inner() {
+    let sample = b"already-decoded-buffer";
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::Yielding {
+        buffer: YieldBuffer::new(Bytes::copy_from_slice(sample)),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert_eq!(pending, sample);
+}
+
+#[test]
+fn reading_body_non_application_record_does_not_produce_plaintext() {
+    let sample = b"unexpected-handshake-fragment";
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::ReadingBody {
+        record_type: TLS_RECORD_HANDSHAKE,
+        length: sample.len(),
+        buffer: BytesMut::from(&sample[..]),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert!(
+        pending.is_empty(),
+        "non-application partial body must not be surfaced as plaintext"
+    );
+}
+
+#[test]
+fn partial_header_state_does_not_produce_plaintext() {
+    let mut header = HeaderBuffer::<TLS_HEADER_SIZE>::new();
+    let unfilled = header.unfilled_mut();
+    unfilled[0] = TLS_RECORD_APPLICATION;
+    header.advance(1);
+
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::ReadingHeader { header };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert!(pending.is_empty(), "partial header bytes are not plaintext payload");
+}
+
+#[test]
+fn edge_zero_length_application_fragment_remains_empty_without_panics() {
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::ReadingBody {
+        record_type: TLS_RECORD_APPLICATION,
+        length: 0,
+        buffer: BytesMut::new(),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert!(pending.is_empty());
+}
+
+#[test]
+fn adversarial_poisoned_state_never_leaks_pending_bytes() {
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::Poisoned {
+        error: Some(std::io::Error::other("poisoned by adversarial input")),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert!(pending.is_empty(), "poisoned state must fail-closed for fallback payload");
+}
+
+#[test]
+fn stress_large_application_fragment_survives_state_extraction() {
+    let mut payload = vec![0u8; 96 * 1024];
+    for (i, b) in payload.iter_mut().enumerate() {
+        *b = (i as u8).wrapping_mul(17).wrapping_add(3);
+    }
+
+    let mut reader = FakeTlsReader::new(tokio::io::empty());
+    reader.state = TlsReaderState::ReadingBody {
+        record_type: TLS_RECORD_APPLICATION,
+        length: payload.len(),
+        buffer: BytesMut::from(&payload[..]),
+    };
+
+    let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+    assert_eq!(pending, payload, "large pending application plaintext must be preserved exactly");
+}
+
+#[test]
+fn light_fuzz_state_matrix_preserves_pending_contract() {
+    let mut seed = 0x9E37_79B9_7F4A_7C15u64;
+
+    for _ in 0..4096 {
+        seed ^= seed << 7;
+        seed ^= seed >> 9;
+        seed ^= seed << 8;
+
+        let len = (seed & 0x1ff) as usize;
+        let mut payload = vec![0u8; len];
+        for (idx, b) in payload.iter_mut().enumerate() {
+            *b = (seed as u8).wrapping_add(idx as u8);
+        }
+
+        let record_type = match seed & 0x3 {
+            0 => TLS_RECORD_APPLICATION,
+            1 => TLS_RECORD_HANDSHAKE,
+            2 => TLS_RECORD_ALERT,
+            _ => TLS_RECORD_CHANGE_CIPHER,
+        };
+
+        let mut reader = FakeTlsReader::new(tokio::io::empty());
+        reader.state = TlsReaderState::ReadingBody {
+            record_type,
+            length: payload.len(),
+            buffer: BytesMut::from(&payload[..]),
+        };
+
+        let (_inner, pending) = reader.into_inner_with_pending_plaintext();
+        if record_type == TLS_RECORD_APPLICATION {
+            assert_eq!(pending, payload);
+        } else {
+            assert!(pending.is_empty());
+        }
+    }
+}