mirror of https://github.com/ospab/ostp.git
603 lines
24 KiB
Rust
603 lines
24 KiB
Rust
use anyhow::Result;
|
|
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
|
use hmac::{Hmac, Mac};
|
|
use sha2::Sha256;
|
|
use std::collections::HashMap;
|
|
use std::net::SocketAddr;
|
|
use std::sync::{Arc, RwLock as StdRwLock};
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::sync::{mpsc, RwLock};
|
|
use tracing::info;
|
|
use tokio::net::TcpStream;
|
|
use base64::Engine;
|
|
use std::pin::Pin;
|
|
use std::task::{Context as TaskContext, Poll};
|
|
use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, Nonce};
|
|
use x25519_dalek::StaticSecret;
|
|
|
|
use ostp_core::framing::wss::{encode_wss_frame, decode_wss_frame, WssFrameResult};
|
|
use ostp_core::crypto::reality::{parse_client_hello, derive_keys, verify_session_id, REALITY_SERVER_HANDSHAKE_RECORDS};
|
|
use crate::RealityServerConfig;
|
|
|
|
pub async fn handle_tcp_connection<S>(
|
|
mut stream: S,
|
|
peer_addr: SocketAddr,
|
|
shared_keys: Arc<StdRwLock<HashMap<String, crate::api::UserMeta>>>,
|
|
udp_tx: mpsc::Sender<(Bytes, SocketAddr)>,
|
|
tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>>,
|
|
reality_config: Option<Arc<RealityServerConfig>>,
|
|
fb_target: Option<String>,
|
|
) -> Result<()>
|
|
where
|
|
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
|
{
|
|
let mut initial_buf = vec![0u8; 16384];
|
|
let mut header_len = 0;
|
|
|
|
// Read the first chunk to determine if it's TLS or HTTP
|
|
let n = stream.read(&mut initial_buf).await?;
|
|
if n == 0 {
|
|
anyhow::bail!("connection closed before data received");
|
|
}
|
|
header_len += n;
|
|
|
|
// Check if it's a TLS record (0x16 0x03 0x01 or 0x16 0x03 0x03)
|
|
if initial_buf[0] == 0x16 && initial_buf[1] == 0x03 {
|
|
// It's a TLS record. We need to ensure we read the entire record.
|
|
if header_len >= 5 {
|
|
let record_len = 5 + u16::from_be_bytes([initial_buf[3], initial_buf[4]]) as usize;
|
|
if record_len > initial_buf.len() {
|
|
anyhow::bail!("TLS record too large");
|
|
}
|
|
while header_len < record_len {
|
|
let n = stream.read(&mut initial_buf[header_len..record_len]).await?;
|
|
if n == 0 {
|
|
anyhow::bail!("connection closed while reading TLS record");
|
|
}
|
|
header_len += n;
|
|
}
|
|
}
|
|
|
|
if let Some(rc) = reality_config {
|
|
return handle_reality_connection(stream, initial_buf[..header_len].to_vec(), peer_addr, shared_keys, udp_tx, tcp_map, rc).await;
|
|
} else {
|
|
// Received TLS but Reality is not enabled
|
|
if let Some(target) = fb_target {
|
|
tracing::info!("Fallback triggered for {} -> {}", peer_addr, target);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&target).await?;
|
|
dest_stream.write_all(&initial_buf[..header_len]).await?;
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
} else {
|
|
anyhow::bail!("received TLS but Reality is not configured and no fallback target");
|
|
}
|
|
}
|
|
}
|
|
|
|
// Otherwise, assume it's HTTP (Standard xhttp/wss)
|
|
loop {
|
|
if initial_buf[..header_len].windows(4).any(|w| w == b"\r\n\r\n") {
|
|
break;
|
|
}
|
|
if header_len == initial_buf.len() {
|
|
anyhow::bail!("handshake headers too large");
|
|
}
|
|
let n = stream.read(&mut initial_buf[header_len..]).await?;
|
|
if n == 0 {
|
|
anyhow::bail!("connection closed before HTTP handshake complete");
|
|
}
|
|
header_len += n;
|
|
}
|
|
|
|
let headers_str = String::from_utf8_lossy(&initial_buf[..header_len]);
|
|
|
|
let wss = if headers_str.starts_with("GET /wss HTTP/1.1\r\n") {
|
|
true
|
|
} else if headers_str.starts_with("GET /stream HTTP/1.1\r\n") {
|
|
false
|
|
} else {
|
|
if let Some(target) = fb_target {
|
|
tracing::info!("Fallback triggered for {} -> {}", peer_addr, target);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&target).await?;
|
|
dest_stream.write_all(&initial_buf[..header_len]).await?;
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
} else {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("invalid request line");
|
|
}
|
|
};
|
|
|
|
// Extract Authorization
|
|
let mut signature_base64 = None;
|
|
for line in headers_str.lines() {
|
|
let lower = line.to_ascii_lowercase();
|
|
if lower.starts_with("authorization: bearer ") {
|
|
signature_base64 = Some(line[22..].trim().to_string());
|
|
}
|
|
}
|
|
|
|
let sig_b64 = match signature_base64 {
|
|
Some(s) => s,
|
|
None => {
|
|
if let Some(target) = fb_target {
|
|
tracing::info!("Fallback triggered for {} -> {}", peer_addr, target);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&target).await?;
|
|
dest_stream.write_all(&initial_buf[..header_len]).await?;
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
} else {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("missing authorization");
|
|
}
|
|
}
|
|
};
|
|
|
|
let sig_bytes = match base64::Engine::decode(&base64::engine::general_purpose::STANDARD_NO_PAD, &sig_b64) {
|
|
Ok(b) => b,
|
|
Err(_) => {
|
|
if let Some(target) = fb_target {
|
|
tracing::info!("Fallback triggered for {} -> {}", peer_addr, target);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&target).await?;
|
|
dest_stream.write_all(&initial_buf[..header_len]).await?;
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
} else {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("invalid base64 signature");
|
|
}
|
|
}
|
|
};
|
|
|
|
if sig_bytes.len() < 8 {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("signature too short");
|
|
}
|
|
|
|
let ts_bytes: [u8; 8] = sig_bytes[0..8].try_into().unwrap();
|
|
let client_ts = u64::from_be_bytes(ts_bytes);
|
|
let provided_mac = &sig_bytes[8..];
|
|
|
|
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
|
|
if client_ts > now + 30 || client_ts < now.saturating_sub(60) {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("timestamp out of bounds (replay protection)");
|
|
}
|
|
|
|
// Verify HMAC against known keys
|
|
let keys = {
|
|
let guard = shared_keys.read().unwrap();
|
|
guard.keys().cloned().collect::<Vec<_>>()
|
|
};
|
|
|
|
let mut authenticated = false;
|
|
for key in keys {
|
|
let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(key.as_bytes())
|
|
.unwrap_or_else(|_| <Hmac<Sha256> as Mac>::new_from_slice(b"default").unwrap());
|
|
mac.update(&ts_bytes);
|
|
if mac.verify_slice(provided_mac).is_ok() {
|
|
authenticated = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if !authenticated {
|
|
if let Some(target) = fb_target {
|
|
tracing::info!("Fallback triggered for {} -> {}", peer_addr, target);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&target).await?;
|
|
dest_stream.write_all(&initial_buf[..header_len]).await?;
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
} else {
|
|
let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\n\r\n").await;
|
|
anyhow::bail!("unauthorized (invalid HMAC)");
|
|
}
|
|
}
|
|
|
|
if wss {
|
|
let response = "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\nX-Ostp-Server: 1\r\n\r\n";
|
|
stream.write_all(response.as_bytes()).await?;
|
|
} else {
|
|
let response = "HTTP/1.1 200 OK\r\nX-Ostp-Server: 1\r\nContent-Type: application/octet-stream\r\n\r\n";
|
|
stream.write_all(response.as_bytes()).await?;
|
|
}
|
|
|
|
info!("UoT client authenticated from {} (xhttp)", peer_addr);
|
|
|
|
start_uot_loops(stream, peer_addr, wss, tcp_map, udp_tx).await
|
|
}
|
|
|
|
async fn handle_reality_connection<S>(
|
|
mut stream: S,
|
|
initial_buf: Vec<u8>,
|
|
peer_addr: SocketAddr,
|
|
_shared_keys: Arc<StdRwLock<HashMap<String, crate::api::UserMeta>>>, // Note: Reality uses its own keys (sid)
|
|
udp_tx: mpsc::Sender<(Bytes, SocketAddr)>,
|
|
tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>>,
|
|
reality_config: Arc<RealityServerConfig>,
|
|
) -> Result<()>
|
|
where
|
|
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
|
{
|
|
// Try to parse ClientHello
|
|
let parsed_ch = parse_client_hello(&initial_buf);
|
|
|
|
let mut authenticated = false;
|
|
let mut data_key_opt = None;
|
|
|
|
if let Some(ch) = parsed_ch {
|
|
// Validate SNI
|
|
if reality_config.sni_list.contains(&ch.sni) {
|
|
// Decode Server Private Key
|
|
if let Ok(priv_bytes) = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(&reality_config.private_key) {
|
|
if priv_bytes.len() == 32 {
|
|
let mut secret_bytes = [0u8; 32];
|
|
secret_bytes.copy_from_slice(&priv_bytes);
|
|
let server_priv = StaticSecret::from(secret_bytes);
|
|
|
|
let shared_secret = server_priv.diffie_hellman(&ch.c_pub);
|
|
let (auth_key, data_key) = derive_keys(shared_secret.as_bytes());
|
|
|
|
// Attempt to decrypt Session ID
|
|
if let Some((sid, _ts)) = verify_session_id(&auth_key, &ch.session_id) {
|
|
// Check if sid is in config
|
|
let sid_hex = hex::encode(sid);
|
|
if reality_config.sid == sid_hex {
|
|
authenticated = true;
|
|
data_key_opt = Some(data_key);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if authenticated {
|
|
let data_key = data_key_opt.unwrap();
|
|
info!("Reality client authenticated from {} (sid matched)", peer_addr);
|
|
|
|
// Build a fake TLS 1.3 server flight that matches what a real server sends.
|
|
// Must be exactly REALITY_SERVER_HANDSHAKE_RECORDS (5) TLS records:
|
|
// 1. ServerHello (0x16) - static blob with fake key share
|
|
// 2. ChangeCipherSpec (0x14) - RFC 8446 §D.4 middlebox compat
|
|
// 3. Fake EE (0x17) - simulates EncryptedExtensions
|
|
// 4. Fake Certificate (0x17) - simulates Certificate (big, DPI-realistic)
|
|
// 5. Fake Finished (0x17) - simulates CertificateVerify + Finished
|
|
let _ = REALITY_SERVER_HANDSHAKE_RECORDS; // assert constant is imported (= 5)
|
|
|
|
// Record 1: ServerHello (0x16), same static blob as before (valid structure)
|
|
let server_hello_rec = hex::decode(
|
|
"160303007a0200007603030000000000000000000000000000000000000000000000\
|
|
000000000000000000000000200000000000000000000000000000000000000000\
|
|
0000000000000000000000000000130100002e002b0002030400330024001d0020\
|
|
e29b191a62d0572e9a30d0fb9d08e50bc78d591dfc1dbafbfa533411db1c8e11"
|
|
).unwrap();
|
|
|
|
// Record 2: ChangeCipherSpec (0x14)
|
|
let ccs_rec: &[u8] = &[0x14, 0x03, 0x03, 0x00, 0x01, 0x01];
|
|
|
|
// Record 3: Fake EncryptedExtensions (0x17), 108 zero bytes payload
|
|
let mut fake_ee = vec![0x17u8, 0x03, 0x03, 0x00, 108];
|
|
fake_ee.extend_from_slice(&[0u8; 108]);
|
|
|
|
// Record 4: Fake Certificate (0x17), 812 zero bytes (realistic cert size for DPI)
|
|
let cert_payload_len: u16 = 812;
|
|
let mut fake_cert = vec![0x17u8, 0x03, 0x03,
|
|
(cert_payload_len >> 8) as u8, (cert_payload_len & 0xff) as u8];
|
|
fake_cert.extend_from_slice(&vec![0u8; cert_payload_len as usize]);
|
|
|
|
// Record 5: Fake Finished (0x17), 52 zero bytes (CertificateVerify + Finished)
|
|
let mut fake_fin = vec![0x17u8, 0x03, 0x03, 0x00, 52];
|
|
fake_fin.extend_from_slice(&[0u8; 52]);
|
|
|
|
let mut server_flight = Vec::with_capacity(
|
|
server_hello_rec.len() + ccs_rec.len() +
|
|
fake_ee.len() + fake_cert.len() + fake_fin.len()
|
|
);
|
|
server_flight.extend_from_slice(&server_hello_rec);
|
|
server_flight.extend_from_slice(ccs_rec);
|
|
server_flight.extend_from_slice(&fake_ee);
|
|
server_flight.extend_from_slice(&fake_cert);
|
|
server_flight.extend_from_slice(&fake_fin);
|
|
|
|
stream.write_all(&server_flight).await?;
|
|
|
|
// The client now sends ClientHello + CCS (6 bytes) as two separate TLS records.
|
|
// The ClientHello was already consumed into initial_buf above.
|
|
// The CCS may arrive as a separate TCP segment - drain it from the raw stream
|
|
// before wrapping in RealityStream so RealityStream only ever sees 0x17 records.
|
|
{
|
|
let mut ccs_head = [0u8; 5];
|
|
if stream.read_exact(&mut ccs_head).await.is_ok() {
|
|
// Expected: CCS record 0x14 0x03 0x03 0x00 0x01
|
|
// If it's something else (unlikely), we still drain its payload to stay in sync.
|
|
let ccs_payload_len = u16::from_be_bytes([ccs_head[3], ccs_head[4]]) as usize;
|
|
if ccs_payload_len <= 64 {
|
|
let mut _discard = vec![0u8; ccs_payload_len];
|
|
let _ = stream.read_exact(&mut _discard).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
let reality_stream = RealityStream::new(stream, data_key);
|
|
return process_inner_reality_stream(reality_stream, peer_addr, tcp_map, udp_tx).await;
|
|
|
|
} else {
|
|
// Fallback: act as a transparent proxy to `reality_config.dest`
|
|
info!("Reality fallback triggered for {} -> {}", peer_addr, reality_config.dest);
|
|
let mut dest_stream: TcpStream = TcpStream::connect(&reality_config.dest).await?;
|
|
dest_stream.write_all(&initial_buf).await?;
|
|
|
|
tokio::io::copy_bidirectional(&mut stream, &mut dest_stream).await?;
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
async fn process_inner_reality_stream<S>(
|
|
mut stream: S,
|
|
peer_addr: SocketAddr,
|
|
tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>>,
|
|
udp_tx: mpsc::Sender<(Bytes, SocketAddr)>,
|
|
) -> Result<()>
|
|
where
|
|
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
|
{
|
|
// 1. Read the inner HTTP Handshake
|
|
let mut buf = [0u8; 4096];
|
|
let mut header_len = 0;
|
|
loop {
|
|
let n = stream.read(&mut buf[header_len..]).await?;
|
|
if n == 0 {
|
|
anyhow::bail!("inner connection closed before handshake complete");
|
|
}
|
|
header_len += n;
|
|
if buf[..header_len].windows(4).any(|w| w == b"\r\n\r\n") {
|
|
break;
|
|
}
|
|
if header_len == buf.len() {
|
|
anyhow::bail!("inner handshake headers too large");
|
|
}
|
|
}
|
|
|
|
let headers_str = String::from_utf8_lossy(&buf[..header_len]);
|
|
|
|
let wss = if headers_str.starts_with("GET /wss HTTP/1.1\r\n") {
|
|
true
|
|
} else if headers_str.starts_with("GET /stream HTTP/1.1\r\n") {
|
|
false
|
|
} else {
|
|
anyhow::bail!("invalid inner request line");
|
|
};
|
|
|
|
// We skip signature validation because Reality already authenticated the user via Session ID!
|
|
|
|
if wss {
|
|
let response = "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\nX-Ostp-Server: 1\r\n\r\n";
|
|
stream.write_all(response.as_bytes()).await?;
|
|
} else {
|
|
let response = "HTTP/1.1 200 OK\r\nX-Ostp-Server: 1\r\nContent-Type: application/octet-stream\r\n\r\n";
|
|
stream.write_all(response.as_bytes()).await?;
|
|
}
|
|
|
|
start_uot_loops(stream, peer_addr, wss, tcp_map, udp_tx).await
|
|
}
|
|
|
|
async fn start_uot_loops<S>(
|
|
stream: S,
|
|
peer_addr: SocketAddr,
|
|
wss: bool,
|
|
tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>>,
|
|
udp_tx: mpsc::Sender<(Bytes, SocketAddr)>,
|
|
) -> Result<()>
|
|
where
|
|
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
|
{
|
|
// Register this connection in the map
|
|
let (tx, mut rx) = mpsc::channel::<Bytes>(16384);
|
|
{
|
|
tcp_map.write().await.insert(peer_addr, tx);
|
|
}
|
|
|
|
// Process streams
|
|
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
|
|
|
// Spawn writer task
|
|
let peer_clone = peer_addr;
|
|
let tcp_map_clone = tcp_map.clone();
|
|
let writer_task = tokio::spawn(async move {
|
|
while let Some(packet) = rx.recv().await {
|
|
if wss {
|
|
let header = encode_wss_frame(&packet, false); // Server sends unmasked WSS frames
|
|
if write_half.write_all(&header).await.is_err() { break; }
|
|
} else {
|
|
let mut out = BytesMut::with_capacity(2 + packet.len());
|
|
out.put_u16(packet.len() as u16);
|
|
out.put_slice(&packet);
|
|
if write_half.write_all(&out).await.is_err() { break; }
|
|
}
|
|
}
|
|
let _ = tcp_map_clone.write().await.remove(&peer_clone);
|
|
});
|
|
|
|
// Spawn reader task
|
|
let tcp_map_clone2 = tcp_map.clone();
|
|
let reader_task = tokio::spawn(async move {
|
|
if wss {
|
|
let mut read_buf = BytesMut::with_capacity(65536);
|
|
let mut tmp = [0u8; 8192];
|
|
loop {
|
|
match read_half.read(&mut tmp).await {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
read_buf.put_slice(&tmp[..n]);
|
|
loop {
|
|
match decode_wss_frame(&mut read_buf) {
|
|
WssFrameResult::Frame { payload, total_len } => {
|
|
if udp_tx.send((Bytes::from(payload), peer_clone)).await.is_err() { return; }
|
|
read_buf.advance(total_len);
|
|
}
|
|
WssFrameResult::Incomplete => break,
|
|
}
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
} else {
|
|
let mut len_buf = [0u8; 2];
|
|
loop {
|
|
if read_half.read_exact(&mut len_buf).await.is_err() { break; }
|
|
let len = u16::from_be_bytes(len_buf) as usize;
|
|
if len > 65535 { break; }
|
|
let mut data = vec![0u8; len];
|
|
if read_half.read_exact(&mut data).await.is_err() { break; }
|
|
if udp_tx.send((Bytes::from(data), peer_clone)).await.is_err() { break; }
|
|
}
|
|
}
|
|
let _ = tcp_map_clone2.write().await.remove(&peer_clone);
|
|
});
|
|
|
|
let _ = tokio::join!(writer_task, reader_task);
|
|
Ok(())
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// RealityStream: Wraps a TCP stream in fake TLS Application Data Records
|
|
// -----------------------------------------------------------------------
|
|
struct RealityStream<S> {
|
|
inner: S,
|
|
data_key: ChaCha20Poly1305,
|
|
rx_nonce: u64,
|
|
tx_nonce: u64,
|
|
rx_buf: BytesMut,
|
|
plaintext_buf: BytesMut,
|
|
tx_buf: BytesMut,
|
|
}
|
|
|
|
impl<S> RealityStream<S> {
|
|
fn new(inner: S, data_key: ChaCha20Poly1305) -> Self {
|
|
Self {
|
|
inner,
|
|
data_key,
|
|
rx_nonce: 0,
|
|
tx_nonce: 0,
|
|
rx_buf: BytesMut::with_capacity(16384),
|
|
plaintext_buf: BytesMut::new(),
|
|
tx_buf: BytesMut::new(),
|
|
}
|
|
}
|
|
|
|
fn make_nonce(seq: u64) -> [u8; 12] {
|
|
let mut nonce = [0u8; 12];
|
|
nonce[4..12].copy_from_slice(&seq.to_le_bytes());
|
|
nonce
|
|
}
|
|
}
|
|
|
|
impl<S: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for RealityStream<S> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>, buf: &mut tokio::io::ReadBuf<'_>) -> Poll<std::io::Result<()>> {
|
|
loop {
|
|
if !self.plaintext_buf.is_empty() {
|
|
let out_len = std::cmp::min(buf.remaining(), self.plaintext_buf.len());
|
|
buf.put_slice(&self.plaintext_buf[..out_len]);
|
|
self.plaintext_buf.advance(out_len);
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
if self.rx_buf.len() >= 5 {
|
|
let len = u16::from_be_bytes([self.rx_buf[3], self.rx_buf[4]]) as usize;
|
|
if self.rx_buf.len() >= 5 + len {
|
|
if self.rx_buf[0] != 0x17 {
|
|
return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "expected application data record")));
|
|
}
|
|
|
|
let ciphertext = &self.rx_buf[5..5+len];
|
|
let nonce_bytes = Self::make_nonce(self.rx_nonce);
|
|
let nonce = Nonce::from_slice(&nonce_bytes);
|
|
|
|
match self.data_key.decrypt(nonce, ciphertext) {
|
|
Ok(plaintext) => {
|
|
self.rx_nonce += 1;
|
|
self.plaintext_buf.put_slice(&plaintext);
|
|
self.rx_buf.advance(5 + len);
|
|
continue;
|
|
}
|
|
Err(_) => return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "reality decrypt failed"))),
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut read_buf = [0u8; 8192];
|
|
let mut tokio_buf = tokio::io::ReadBuf::new(&mut read_buf);
|
|
match Pin::new(&mut self.inner).poll_read(cx, &mut tokio_buf) {
|
|
Poll::Ready(Ok(())) => {
|
|
if tokio_buf.filled().is_empty() { return Poll::Ready(Ok(())); }
|
|
self.rx_buf.put_slice(tokio_buf.filled());
|
|
}
|
|
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
|
Poll::Pending => return Poll::Pending,
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for RealityStream<S> {
|
|
fn poll_write(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
|
|
let this = self.get_mut();
|
|
while !this.tx_buf.is_empty() {
|
|
match Pin::new(&mut this.inner).poll_write(cx, &this.tx_buf) {
|
|
Poll::Ready(Ok(n)) => this.tx_buf.advance(n),
|
|
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
|
Poll::Pending => return Poll::Pending,
|
|
}
|
|
}
|
|
|
|
let nonce_bytes = Self::make_nonce(this.tx_nonce);
|
|
let nonce = Nonce::from_slice(&nonce_bytes);
|
|
|
|
match this.data_key.encrypt(nonce, buf) {
|
|
Ok(ciphertext) => {
|
|
this.tx_nonce += 1;
|
|
this.tx_buf.reserve(5 + ciphertext.len());
|
|
this.tx_buf.put_u8(0x17);
|
|
this.tx_buf.put_u16(0x0303);
|
|
this.tx_buf.put_u16(ciphertext.len() as u16);
|
|
this.tx_buf.put_slice(&ciphertext);
|
|
|
|
match Pin::new(&mut this.inner).poll_write(cx, &this.tx_buf) {
|
|
Poll::Ready(Ok(n)) => this.tx_buf.advance(n),
|
|
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
|
Poll::Pending => {}
|
|
}
|
|
Poll::Ready(Ok(buf.len()))
|
|
}
|
|
Err(_) => Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::Other, "reality encrypt failed"))),
|
|
}
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<std::io::Result<()>> {
|
|
let this = self.get_mut();
|
|
while !this.tx_buf.is_empty() {
|
|
match Pin::new(&mut this.inner).poll_write(cx, &this.tx_buf) {
|
|
Poll::Ready(Ok(n)) => this.tx_buf.advance(n),
|
|
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
|
Poll::Pending => return Poll::Pending,
|
|
}
|
|
}
|
|
Pin::new(&mut this.inner).poll_flush(cx)
|
|
}
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<std::io::Result<()>> {
|
|
let this = self.get_mut();
|
|
while !this.tx_buf.is_empty() {
|
|
match Pin::new(&mut this.inner).poll_write(cx, &this.tx_buf) {
|
|
Poll::Ready(Ok(n)) => this.tx_buf.advance(n),
|
|
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
|
Poll::Pending => return Poll::Pending,
|
|
}
|
|
}
|
|
Pin::new(&mut this.inner).poll_shutdown(cx)
|
|
}
|
|
}
|