Initial Commit

communicator/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "communicator"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+bytes = { workspace = true }
+h2 = { workspace = true }
+http = { workspace = true }
+log = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+tokio = { workspace = true, features = ["net", "rt-multi-thread", "io-util", "time"] }
+tokio-rustls = { workspace = true }
+rustls-pemfile = { workspace = true }
+webpki-roots = { workspace = true }
+anyhow = { workspace = true }
+tokio-util = { workspace = true }
+futures-util = { workspace = true }
+futures = "0.3.32"

communicator/src/config.rs

@@ -0,0 +1,48 @@
+use crate::stream::Identity;
+use crate::{ClientTunnelConfig, ServerTunnelConfig, SslConfig};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ConnectConfig {
+    #[serde(default)]
+    pub server: Vec<TcpServerTunnelConfig>,
+    #[serde(default)]
+    pub client: Vec<TcpClientTunnelConfig>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TcpServerTunnelConfig {
+    pub url: String,
+    #[serde(flatten)]
+    pub cert: Option<SslConfig>,
+    pub token: String,
+}
+
+impl TcpServerTunnelConfig {
+    pub fn into_tunnel_config(self, identity: Identity) -> anyhow::Result<ServerTunnelConfig> {
+        Ok(ServerTunnelConfig {
+            identity,
+            url: self.url,
+            cert: self.cert,
+            token: self.token,
+        })
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TcpClientTunnelConfig {
+    pub host: Option<String>,
+    pub url: String,
+    pub token: String,
+}
+
+impl TcpClientTunnelConfig {
+    pub fn into_tunnel_config(self, identity: Identity) -> ClientTunnelConfig {
+        ClientTunnelConfig {
+            identity,
+            host: self.host,
+            url: self.url,
+            token: self.token,
+        }
+    }
+}

communicator/src/http.rs

@@ -0,0 +1,216 @@
+use crate::stream::ServerManager;
+use bytes::{BufMut, Bytes, BytesMut};
+use futures_util::stream::{self};
+use h2::server::SendResponse;
+use h2::{RecvStream, client};
+use http::{Request, Response, StatusCode};
+use log::{debug, error};
+use serde::de::DeserializeOwned;
+use std::collections::HashMap;
+use std::io;
+use std::pin::Pin;
+use std::sync::Arc;
+use tokio_util::io::StreamReader;
+
+type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
+type Handler = Arc<
+    dyn Fn(Request<RecvStream>, SendResponse<Bytes>) -> BoxFuture<Result<(), h2::Error>>
+        + Send
+        + Sync,
+>;
+
+pub struct Router {
+    routes: HashMap<String, Handler>,
+}
+
+impl Default for Router {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Router {
+    pub fn new() -> Self {
+        Self {
+            routes: HashMap::new(),
+        }
+    }
+
+    pub fn add_route<F, Fut>(&mut self, path: &str, handler: F)
+    where
+        F: Fn(Request<RecvStream>, SendResponse<Bytes>) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = Result<(), h2::Error>> + Send + 'static,
+    {
+        self.routes.insert(
+            path.to_string(),
+            Arc::new(move |req, res| Box::pin(handler(req, res))),
+        );
+    }
+
+    pub async fn dispatch(
+        &self,
+        req: Request<RecvStream>,
+        mut res: SendResponse<Bytes>,
+    ) -> Result<(), h2::Error> {
+        let path = req.uri().path();
+        debug!("Received request for path: {}", path);
+        if let Some(handler) = self.routes.get(path) {
+            handler(req, res).await
+        } else {
+            // 404
+            let response = Response::builder().status(404).body(()).unwrap();
+            res.send_response(response, true)?;
+            Ok(())
+        }
+    }
+}
+
+pub struct Server {
+    router: Arc<Router>,
+}
+
+impl Server {
+    pub fn new(router: Arc<Router>) -> Self {
+        Self { router }
+    }
+
+    pub async fn on_conn(&self, server: Arc<ServerManager>) -> anyhow::Result<()> {
+        let router = self.router.clone();
+
+        while let Some(result) = server.accept().await {
+            let (request, respond) = result?;
+            let r = router.clone();
+            tokio::spawn(async move {
+                if let Err(e) = r.dispatch(request, respond).await {
+                    error!("Handler error: {:?}", e);
+                }
+            });
+        }
+        Ok(())
+    }
+}
+
+pub fn send(mut res: SendResponse<Bytes>, status: u16, content_type: &str, body: String) {
+    let response = Response::builder()
+        .status(status)
+        .header("content-type", content_type)
+        .body(())
+        .unwrap();
+
+    if let Ok(mut send_stream) = res.send_response(response, false) {
+        let _ = send_stream.send_data(Bytes::from(body), true);
+    }
+}
+
+pub fn send_error(mut res: SendResponse<Bytes>, error: anyhow::Error) {
+    let response = Response::builder()
+        .status(StatusCode::INTERNAL_SERVER_ERROR)
+        .header("content-type", "text/plain")
+        .body(())
+        .unwrap();
+
+    if let Ok(mut send_stream) = res.send_response(response, false) {
+        let error_msg = format!("Internal Server Error: {}", error);
+        let _ = send_stream.send_data(Bytes::from(error_msg), true);
+    }
+}
+
+pub async fn json_from_request<T>(req: &mut Request<RecvStream>) -> anyhow::Result<T>
+where
+    T: DeserializeOwned,
+{
+    let body_stream = req.body_mut();
+    let mut buf = BytesMut::new();
+
+    while let Some(chunk) = body_stream.data().await {
+        let data = chunk?;
+
+        let len = data.len();
+        buf.put(data);
+
+        body_stream.flow_control().release_capacity(len)?;
+    }
+
+    if buf.is_empty() {
+        return Err(anyhow::anyhow!("Request body is empty"));
+    }
+
+    let result = serde_json::from_slice(&buf)?;
+    Ok(result)
+}
+
+pub async fn request(
+    client: &mut client::SendRequest<Bytes>,
+    path: &str,
+    body: String,
+) -> anyhow::Result<Response<RecvStream>> {
+    let request = Request::builder()
+        .method("POST")
+        .uri("http://0.0.0.0".to_owned() + path)
+        .header("content-type", "application/json")
+        .body(())?;
+
+    let (response, mut send_stream) = client.send_request(request, false)?;
+    send_stream.send_data(Bytes::from(body), true)?;
+    Ok(response.await?)
+}
+
+async fn bytes_from_response(response: &mut Response<RecvStream>) -> anyhow::Result<Bytes> {
+    let body_stream = response.body_mut();
+    let mut buf = BytesMut::new();
+
+    while let Some(chunk) = body_stream.data().await {
+        let data = chunk?;
+
+        let len = data.len();
+        buf.put(data);
+
+        body_stream.flow_control().release_capacity(len)?;
+    }
+
+    Ok(buf.freeze())
+}
+
+pub async fn text_from_response(response: &mut Response<RecvStream>) -> anyhow::Result<String> {
+    let buf = bytes_from_response(response).await?;
+
+    if buf.is_empty() {
+        return Ok(String::new());
+    }
+
+    Ok(String::from_utf8_lossy(&buf).to_string())
+}
+
+pub async fn json_from_response<T>(response: &mut Response<RecvStream>) -> anyhow::Result<T>
+where
+    T: DeserializeOwned,
+{
+    let buf = bytes_from_response(response).await?;
+
+    if buf.is_empty() {
+        return Err(anyhow::anyhow!("Request body is empty"));
+    }
+
+    let result = serde_json::from_slice(&buf)?;
+    Ok(result)
+}
+
+pub fn response_to_async_read(res: Response<RecvStream>) -> impl tokio::io::AsyncRead {
+    let body = res.into_body();
+
+    let byte_stream = stream::unfold(body, |mut body| async move {
+        match body.data().await {
+            Some(Ok(bytes)) => {
+                let len = bytes.len();
+                if let Err(e) = body.flow_control().release_capacity(len) {
+                    return Some((Err(io::Error::other(e)), body));
+                }
+                Some((Ok::<_, io::Error>(bytes), body))
+            }
+            Some(Err(e)) => Some((Err(io::Error::other(e)), body)),
+            None => None,
+        }
+    });
+
+    StreamReader::new(byte_stream)
+}

communicator/src/lib.rs

@@ -0,0 +1,9 @@
+mod config;
+pub mod http;
+mod stream;
+
+pub use config::{ConnectConfig, TcpClientTunnelConfig, TcpServerTunnelConfig};
+pub use stream::{
+    ClientManager, ClientTunnelConfig, Identity, ServerManager, ServerTunnelConfig, SslConfig,
+    TunnelEndpoint, TunnelListener, connect_tunnel,
+};

communicator/src/stream.rs

@@ -0,0 +1,611 @@
+use anyhow::anyhow;
+use bytes::Bytes;
+use h2::{RecvStream, client, server};
+use http::Request;
+use log::{debug, error, info};
+use serde::{Deserialize, Serialize};
+use std::collections::{HashMap, HashSet};
+use std::error::Error;
+use std::fs::File;
+use std::io::BufReader;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::sync::{Mutex, Notify};
+use tokio::time::{Duration, sleep, timeout};
+use tokio_rustls::rustls;
+use tokio_rustls::rustls::pki_types::{CertificateDer, ServerName};
+use tokio_rustls::{TlsAcceptor, TlsConnector};
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum Identity {
+    Server,
+    Client,
+}
+
+impl Identity {
+    pub fn as_u8(&self) -> u8 {
+        match self {
+            Identity::Server => b'S',
+            Identity::Client => b'C',
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ClientTunnelConfig {
+    pub host: Option<String>,
+    pub url: String,
+    pub token: String,
+    pub identity: Identity,
+}
+
+impl Default for ClientTunnelConfig {
+    fn default() -> Self {
+        Self {
+            host: None,
+            url: "127.0.0.1:3333".to_string(),
+            token: "super_secret_magic_token".to_string(),
+            identity: Identity::Client,
+        }
+    }
+}
+
+unsafe impl Send for ClientTunnelConfig {}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct SslConfig {
+    pub cert: String,
+    pub key: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ServerTunnelConfig {
+    pub url: String,
+    pub cert: Option<SslConfig>,
+    pub token: String,
+    pub identity: Identity,
+}
+
+impl Default for ServerTunnelConfig {
+    fn default() -> Self {
+        Self {
+            url: "127.0.0.1:3333".to_string(),
+            cert: None,
+            token: "super_secret_magic_token".to_string(),
+            identity: Identity::Client,
+        }
+    }
+}
+
+static CLIENT_HELLO: &[u8] = b"Sekai unpacker client hello";
+static SERVER_HELLO: &[u8] = b"Sekai unpacker server hello";
+static TLS_BOOTSTRAP_MAGIC: &[u8] = b"SPTLSP1";
+static RESUME_MAGIC: &[u8] = b"SPRESM1";
+
+pub enum TunnelEndpointRaw {
+    Client(Box<client::SendRequest<Bytes>>),
+    Server(Box<server::Connection<TcpStream, Bytes>>),
+}
+
+#[derive(Clone)]
+pub enum TunnelEndpoint {
+    Client(Arc<ClientManager>),
+    Server(Arc<ServerManager>),
+}
+
+pub struct ClientManager {
+    pub session_id: AtomicU64,
+    pub current_client: Mutex<Option<client::SendRequest<Bytes>>>,
+    pub config: Option<ClientTunnelConfig>,
+    pub notify: Arc<Notify>,
+}
+
+impl ClientManager {
+    pub async fn get_client(&self) -> anyhow::Result<client::SendRequest<Bytes>> {
+        loop {
+            let cached = { self.current_client.lock().await.clone() };
+            if let Some(mut c) = cached {
+                if futures::future::poll_fn(|cx| c.poll_ready(cx))
+                    .await
+                    .is_ok()
+                {
+                    return Ok(c);
+                }
+                debug!("Client physical connection dead, preparing to recover...");
+            }
+
+            let mut lock = self.current_client.lock().await;
+            if lock.is_some()
+                && futures::future::poll_fn(|cx| lock.as_mut().unwrap().poll_ready(cx))
+                    .await
+                    .is_ok()
+            {
+                continue;
+            }
+
+            if let Some(config) = &self.config {
+                let mut retry_delay = Duration::from_secs(1);
+                let mut sid = self.session_id.load(Ordering::Relaxed);
+
+                let new_client = loop {
+                    match do_client_reconnect(config, &mut sid).await {
+                        Ok(TunnelEndpointRaw::Client(c)) => {
+                            self.session_id.store(sid, Ordering::Relaxed);
+                            break c;
+                        }
+                        Ok(_) => return Err(anyhow::anyhow!("Identity mismatch on reconnect")),
+                        Err(e) => {
+                            error!("Reconnect failed: {}. Retrying in {:?}...", e, retry_delay);
+                            drop(e);
+                            sleep(retry_delay).await;
+                            retry_delay = std::cmp::min(retry_delay * 2, Duration::from_secs(30));
+                        }
+                    }
+                };
+                *lock = Some(*new_client.clone());
+                return Ok(*new_client);
+            } else {
+                *lock = None;
+                drop(lock);
+                match timeout(Duration::from_secs(15), self.notify.notified()).await {
+                    Ok(_) => continue,
+                    Err(_) => {
+                        anyhow::bail!("Session did not reconnect within 15s. Throwing error!")
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub struct ServerManager {
+    pub session_id: AtomicU64,
+    pub current_server: Mutex<Option<server::Connection<TcpStream, Bytes>>>,
+    pub config: Option<ClientTunnelConfig>,
+    pub notify: Arc<Notify>,
+}
+
+impl ServerManager {
+    pub async fn accept(
+        &self,
+    ) -> Option<Result<(Request<RecvStream>, server::SendResponse<Bytes>), h2::Error>> {
+        loop {
+            let mut conn_guard = self.current_server.lock().await;
+            if let Some(conn) = conn_guard.as_mut() {
+                if let Some(res) = conn.accept().await {
+                    return Some(res);
+                }
+                *conn_guard = None;
+                log::warn!("Server physical connection dropped, waiting for recovery...");
+            }
+            drop(conn_guard);
+
+            if let Some(config) = &self.config {
+                let mut retry_delay = Duration::from_secs(1);
+                let mut sid = self.session_id.load(Ordering::Relaxed);
+
+                let new_server = loop {
+                    match do_client_reconnect(config, &mut sid).await {
+                        Ok(TunnelEndpointRaw::Server(s)) => {
+                            self.session_id.store(sid, Ordering::Relaxed);
+                            break s;
+                        }
+                        Ok(_) => {
+                            error!("Identity mismatch on reconnect");
+                            return None;
+                        }
+                        Err(e) => {
+                            error!("Reconnect failed: {}. Retrying in {:?}...", e, retry_delay);
+                            drop(e);
+                            sleep(retry_delay).await;
+                            retry_delay = std::cmp::min(retry_delay * 2, Duration::from_secs(30));
+                        }
+                    }
+                };
+                *self.current_server.lock().await = Some(*new_server);
+                continue;
+            } else {
+                match timeout(Duration::from_secs(15), self.notify.notified()).await {
+                    Ok(_) => continue,
+                    Err(_) => {
+                        error!(
+                            "Session did not reconnect within 15s. Throwing error (returning None)!"
+                        );
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+}
+
+enum ResumeResult {
+    NotResume(TcpStream),
+    NewSession(TunnelEndpointRaw, u64),
+    ResumedExisting,
+    Invalid,
+}
+
+pub struct TunnelListener {
+    listener: TcpListener,
+    config: ServerTunnelConfig,
+    pending_plain_sessions: Mutex<HashSet<u64>>,
+    next_session_id: AtomicU64,
+    active_sessions: Mutex<HashMap<u64, TunnelEndpoint>>,
+}
+
+impl TunnelListener {
+    pub async fn bind(config: ServerTunnelConfig) -> anyhow::Result<Self> {
+        let listener = TcpListener::bind(&config.url).await?;
+        info!("TCP tunnel listener bound to {}", &config.url);
+        Ok(Self {
+            listener,
+            config,
+            pending_plain_sessions: Mutex::new(HashSet::new()),
+            next_session_id: AtomicU64::new(1),
+            active_sessions: Mutex::new(HashMap::new()),
+        })
+    }
+
+    pub async fn accept(&self) -> anyhow::Result<TunnelEndpoint> {
+        loop {
+            let (stream, peer_addr) = self.listener.accept().await?;
+            debug!("[{}] Connected on tcp", peer_addr);
+
+            if is_tls_client_hello(&stream).await {
+                if let Err(e) = self.handle_tls_bootstrap(stream, peer_addr).await {
+                    error!("[{}] TLS bootstrap failed: {}", peer_addr, e);
+                }
+                continue;
+            }
+
+            let mut stream = match self.try_resume_plain_session(stream, peer_addr).await? {
+                ResumeResult::NewSession(ep_raw, sid) => {
+                    let ep = wrap_raw_endpoint(sid, ep_raw, None);
+                    self.active_sessions.lock().await.insert(sid, ep.clone());
+                    return Ok(ep);
+                }
+                ResumeResult::ResumedExisting => continue,
+                ResumeResult::Invalid => continue,
+                ResumeResult::NotResume(s) => s,
+            };
+
+            if perform_server_handshake(&mut stream, &self.config.token, self.config.identity)
+                .await
+                .is_err()
+            {
+                debug!("[{}] Plain handshake failed", peer_addr);
+                continue;
+            }
+
+            debug!(
+                "[{}] Plain handshake completed, upgrading to H2...",
+                peer_addr
+            );
+            let ep_raw = upgrade_to_h2_raw(stream, self.config.identity)
+                .await
+                .map_err(|e| anyhow::anyhow!("{}", e))?;
+
+            let sid = self.next_session_id.fetch_add(1, Ordering::Relaxed);
+            let ep = wrap_raw_endpoint(sid, ep_raw, None);
+            self.active_sessions.lock().await.insert(sid, ep.clone());
+            return Ok(ep);
+        }
+    }
+
+    async fn handle_tls_bootstrap(
+        &self,
+        stream: TcpStream,
+        peer_addr: std::net::SocketAddr,
+    ) -> Result<(), Box<dyn Error>> {
+        let cert_cfg = self
+            .config
+            .cert
+            .as_ref()
+            .ok_or("TLS is not enabled on server")?;
+
+        let acceptor = TlsAcceptor::from(build_server_tls_config(cert_cfg)?);
+        let mut tls_stream = acceptor.accept(stream).await?;
+
+        perform_server_handshake(&mut tls_stream, &self.config.token, self.config.identity).await?;
+
+        let mut magic = vec![0u8; TLS_BOOTSTRAP_MAGIC.len()];
+        tls_stream.read_exact(&mut magic).await?;
+        if magic != TLS_BOOTSTRAP_MAGIC {
+            return Err("TLS bootstrap marker mismatch".into());
+        }
+
+        let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed);
+        {
+            let mut pending = self.pending_plain_sessions.lock().await;
+            pending.insert(session_id);
+        }
+
+        tls_stream.write_all(TLS_BOOTSTRAP_MAGIC).await?;
+        tls_stream.write_all(&session_id.to_be_bytes()).await?;
+        debug!(
+            "[{}] TLS bootstrap done, issued plain-H2 session {}",
+            peer_addr, session_id
+        );
+        Ok(())
+    }
+
+    async fn try_resume_plain_session(
+        &self,
+        mut stream: TcpStream,
+        peer_addr: std::net::SocketAddr,
+    ) -> anyhow::Result<ResumeResult> {
+        let mut peek_buf = vec![0u8; RESUME_MAGIC.len()];
+        let n = stream.peek(&mut peek_buf).await?;
+        if n < RESUME_MAGIC.len() || peek_buf != RESUME_MAGIC {
+            return Ok(ResumeResult::NotResume(stream));
+        }
+
+        let mut magic = vec![0u8; RESUME_MAGIC.len()];
+        stream.read_exact(&mut magic).await?;
+        let mut sid_buf = [0u8; 8];
+        stream.read_exact(&mut sid_buf).await?;
+        let session_id = u64::from_be_bytes(sid_buf);
+
+        let is_pending = {
+            let mut pending = self.pending_plain_sessions.lock().await;
+            pending.remove(&session_id)
+        };
+
+        if is_pending {
+            let ep_raw = upgrade_to_h2_raw(stream, self.config.identity)
+                .await
+                .map_err(|e| anyhow::anyhow!("{}", e))?;
+            return Ok(ResumeResult::NewSession(ep_raw, session_id));
+        }
+
+        let active = self.active_sessions.lock().await.get(&session_id).cloned();
+        if let Some(ep) = active {
+            let ep_raw = upgrade_to_h2_raw(stream, self.config.identity)
+                .await
+                .map_err(|e| anyhow::anyhow!("{}", e))?;
+            update_endpoint(&ep, ep_raw).await;
+            info!(
+                "[{}] Successfully resumed existing session {}",
+                peer_addr, session_id
+            );
+            return Ok(ResumeResult::ResumedExisting);
+        }
+
+        error!("[{}] Invalid plain-H2 session {}", peer_addr, session_id);
+        Ok(ResumeResult::Invalid)
+    }
+}
+
+async fn is_tls_client_hello(stream: &TcpStream) -> bool {
+    let mut header = [0u8; 3];
+    match stream.peek(&mut header).await {
+        Ok(n) if n >= 3 => header[0] == 0x16 && header[1] == 0x03 && (1..=4).contains(&header[2]),
+        _ => false,
+    }
+}
+
+async fn perform_server_handshake<S>(
+    stream: &mut S,
+    token: &str,
+    identity: Identity,
+) -> Result<(), Box<dyn Error>>
+where
+    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
+{
+    let mut client_hello_buf = vec![0u8; CLIENT_HELLO.len()];
+    stream.read_exact(&mut client_hello_buf).await?;
+    if client_hello_buf != CLIENT_HELLO {
+        return Err("Client Hello mismatch".into());
+    }
+
+    let mut token_buf = vec![0u8; token.len()];
+    stream.read_exact(&mut token_buf).await?;
+    if String::from_utf8_lossy(&token_buf) != token {
+        return Err("Wrong token".into());
+    }
+
+    stream.write_all(SERVER_HELLO).await?;
+
+    let mut peer_id_buf = [0u8; 1];
+    stream.read_exact(&mut peer_id_buf).await?;
+    stream.write_all(&[identity.as_u8()]).await?;
+
+    if peer_id_buf[0] == identity.as_u8() {
+        return Err("Identity collision with peer".into());
+    }
+
+    Ok(())
+}
+
+async fn perform_client_handshake<S>(
+    stream: &mut S,
+    token: &str,
+    identity: Identity,
+) -> anyhow::Result<()>
+where
+    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
+{
+    stream.write_all(CLIENT_HELLO).await?;
+    stream.write_all(token.as_bytes()).await?;
+
+    let mut server_hello_buf = vec![0u8; SERVER_HELLO.len()];
+    stream.read_exact(&mut server_hello_buf).await?;
+    if server_hello_buf != SERVER_HELLO {
+        return Err(anyhow!("Server Hello mismatch"));
+    }
+
+    stream.write_all(&[identity.as_u8()]).await?;
+    let mut peer_id_buf = [0u8; 1];
+    stream.read_exact(&mut peer_id_buf).await?;
+
+    if peer_id_buf[0] == identity.as_u8() {
+        return Err(anyhow!("Identity collision with server"));
+    }
+
+    Ok(())
+}
+
+fn build_server_tls_config(
+    cert_cfg: &SslConfig,
+) -> Result<Arc<rustls::ServerConfig>, Box<dyn Error>> {
+    let cert_file = File::open(&cert_cfg.cert)?;
+    let mut cert_reader = BufReader::new(cert_file);
+    let cert_chain: Vec<CertificateDer<'static>> =
+        rustls_pemfile::certs(&mut cert_reader).collect::<Result<Vec<_>, _>>()?;
+    if cert_chain.is_empty() {
+        return Err("No certificate found in PEM".into());
+    }
+
+    let key_file = File::open(&cert_cfg.key)?;
+    let mut key_reader = BufReader::new(key_file);
+    let private_key =
+        rustls_pemfile::private_key(&mut key_reader)?.ok_or("No private key found in PEM")?;
+
+    let server_config = rustls::ServerConfig::builder()
+        .with_no_client_auth()
+        .with_single_cert(cert_chain, private_key)?;
+    Ok(Arc::new(server_config))
+}
+
+fn build_client_tls_connector() -> TlsConnector {
+    let root_store =
+        rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
+    let config = rustls::ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+    TlsConnector::from(Arc::new(config))
+}
+
+async fn upgrade_to_h2_raw(
+    stream: TcpStream,
+    identity: Identity,
+) -> anyhow::Result<TunnelEndpointRaw> {
+    match identity {
+        Identity::Client => {
+            let (h2_client, h2_conn) = client::handshake(stream).await?;
+            tokio::spawn(async move {
+                if let Err(e) = h2_conn.await {
+                    debug!("H2 connection driver finished/error: {:?}", e);
+                }
+            });
+            Ok(TunnelEndpointRaw::Client(Box::new(h2_client)))
+        }
+        Identity::Server => {
+            let h2_conn = server::handshake(stream).await?;
+            Ok(TunnelEndpointRaw::Server(Box::new(h2_conn)))
+        }
+    }
+}
+
+fn wrap_raw_endpoint(
+    sid: u64,
+    raw: TunnelEndpointRaw,
+    config: Option<ClientTunnelConfig>,
+) -> TunnelEndpoint {
+    match raw {
+        TunnelEndpointRaw::Client(c) => TunnelEndpoint::Client(Arc::new(ClientManager {
+            session_id: AtomicU64::new(sid),
+            current_client: Mutex::new(Some(*c)),
+            config,
+            notify: Arc::new(Notify::new()),
+        })),
+        TunnelEndpointRaw::Server(s) => TunnelEndpoint::Server(Arc::new(ServerManager {
+            session_id: AtomicU64::new(sid),
+            current_server: Mutex::new(Some(*s)),
+            config,
+            notify: Arc::new(Notify::new()),
+        })),
+    }
+}
+
+async fn update_endpoint(ep: &TunnelEndpoint, ep_raw: TunnelEndpointRaw) {
+    match (ep, ep_raw) {
+        (TunnelEndpoint::Client(mgr), TunnelEndpointRaw::Client(c)) => {
+            *mgr.current_client.lock().await = Some(*c);
+            mgr.notify.notify_waiters();
+        }
+        (TunnelEndpoint::Server(mgr), TunnelEndpointRaw::Server(s)) => {
+            *mgr.current_server.lock().await = Some(*s);
+            mgr.notify.notify_waiters();
+        }
+        _ => error!("Identity mismatch during session resume!"),
+    }
+}
+
+pub async fn connect_tunnel(config: ClientTunnelConfig) -> Result<TunnelEndpoint, Box<dyn Error>> {
+    info!("Connecting to tunnel at {}", &config.url);
+    let mut sid = 0;
+    let ep_raw = do_client_reconnect(&config, &mut sid).await?;
+    Ok(wrap_raw_endpoint(sid, ep_raw, Some(config)))
+}
+
+async fn do_client_reconnect(
+    config: &ClientTunnelConfig,
+    current_sid: &mut u64,
+) -> anyhow::Result<TunnelEndpointRaw> {
+    if *current_sid != 0 {
+        match resume_tunnel_client(config, *current_sid).await {
+            Ok(ep) => {
+                info!("Resumed existing session {}", current_sid);
+                return Ok(ep);
+            }
+            Err(e) => {
+                log::warn!("Resume failed: {}. Falling back to full bootstrap.", e);
+            }
+        }
+    }
+
+    if let Some(host) = config.host.clone() {
+        let (sid, _) = bootstrap_tls_and_get_sid(config, &host).await?;
+        *current_sid = sid;
+        resume_tunnel_client(config, sid).await
+    } else {
+        let mut stream = TcpStream::connect(&config.url).await?;
+        perform_client_handshake(&mut stream, &config.token, config.identity).await?;
+        let raw = upgrade_to_h2_raw(stream, config.identity).await?;
+        *current_sid = 0;
+        Ok(raw)
+    }
+}
+
+async fn bootstrap_tls_and_get_sid(
+    config: &ClientTunnelConfig,
+    host: &str,
+) -> anyhow::Result<(u64, ())> {
+    let connector = build_client_tls_connector();
+    let tcp = TcpStream::connect(&config.url).await?;
+    let server_name = ServerName::try_from(host.to_string())
+        .map_err(|_| anyhow!("Invalid TLS host: {}", host))?
+        .to_owned();
+
+    let mut tls_stream = connector
+        .connect(server_name, tcp)
+        .await
+        .map_err(|e| anyhow!("TLS handshake failed (server may not support TLS): {}", e))?;
+
+    perform_client_handshake(&mut tls_stream, &config.token, config.identity).await?;
+
+    tls_stream.write_all(TLS_BOOTSTRAP_MAGIC).await?;
+    let mut magic = vec![0u8; TLS_BOOTSTRAP_MAGIC.len()];
+    tls_stream.read_exact(&mut magic).await?;
+    if magic != TLS_BOOTSTRAP_MAGIC {
+        return Err(anyhow!("TLS bootstrap response mismatch"));
+    }
+
+    let mut sid_buf = [0u8; 8];
+    tls_stream.read_exact(&mut sid_buf).await?;
+    Ok((u64::from_be_bytes(sid_buf), ()))
+}
+
+async fn resume_tunnel_client(
+    config: &ClientTunnelConfig,
+    session_id: u64,
+) -> anyhow::Result<TunnelEndpointRaw> {
+    let mut plain_stream = TcpStream::connect(&config.url).await?;
+    plain_stream.write_all(RESUME_MAGIC).await?;
+    plain_stream.write_all(&session_id.to_be_bytes()).await?;
+    upgrade_to_h2_raw(plain_stream, config.identity).await
+}