Make the protocol transport-agnostic and use tokio-tungstenite

Replace the hand-rolled WebSocket implementation with a transport
boundary (Frame over any Sink/Stream) plus adapters for tokio-tungstenite
(server/client features) and axum (axum feature). The JSON wire format
is unchanged, so existing peers keep working.

Fixes:
- unbounded frame lengths were allocated up front; the server now
  enforces message/frame size limits (1 MiB default)
- responses arriving before `request` subscribed were lost
- an accept() error ended `session_loop`
- requests sent right after connecting could arrive before handlers
  were registered; the receiver now starts after `on_conn` returns
- panics in the receive loop skipped `on_close` and leaked sessions;
  `on_close` now runs exactly once and handler panics fail only their
  request
- unknown methods and invalid data got no reply; they now get an
  error response
- slow handlers blocked pongs and responses

Adds `on_notification`, `request_timeout`, `closed`, `is_closed`, `id`,
`ServerConfig`, `from_transport`, `from_tungstenite` and `from_axum`,
integration tests, and an axum example. Bumps to 0.2.0 since
`Session::connect` now takes a URL and the `ws` module is removed.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

src/server.rs

@@ -1,62 +1,124 @@
-use std::{net::SocketAddr, sync::Arc};
+use std::{net::SocketAddr, sync::Arc, time::Duration};
 
-use tokio::{net::TcpListener, time::timeout};
+use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
 
-use crate::{session::Session, ws::WebSocket};
+use crate::{Error, session::Session};
 
+/// Limits applied to every accepted connection.
+#[derive(Debug, Clone)]
+pub struct ServerConfig {
+    /// Largest message (after reassembling fragments) a peer may send.
+    pub max_message_size: usize,
+    /// Largest single frame a peer may send.
+    pub max_frame_size: usize,
+    /// How long a client gets to complete the WebSocket handshake.
+    pub handshake_timeout: Duration,
+}
+
+impl Default for ServerConfig {
+    fn default() -> Self {
+        Self {
+            max_message_size: 1 << 20,
+            max_frame_size: 1 << 20,
+            handshake_timeout: Duration::from_secs(5),
+        }
+    }
+}
+
+/// A standalone WebSocket server that hands each connection to a callback as a
+/// [`Session`].
 pub struct SessionServer {
     listener: TcpListener,
+    config: ServerConfig,
 }
 
 impl SessionServer {
-    pub async fn bind(addr: &str) -> crate::Result<Self> {
-        Ok(Self {
-            listener: TcpListener::bind(addr).await?,
-        })
+    pub async fn bind(addr: impl ToSocketAddrs) -> crate::Result<Self> {
+        Ok(Self::from_listener(TcpListener::bind(addr).await?))
     }
 
+    pub fn from_listener(listener: TcpListener) -> Self {
+        Self {
+            listener,
+            config: ServerConfig::default(),
+        }
+    }
+
+    pub fn with_config(mut self, config: ServerConfig) -> Self {
+        self.config = config;
+        self
+    }
+
+    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
+        self.listener.local_addr()
+    }
+
+    /// Accept one connection. The receiver is not started: register handlers,
+    /// then call [`Session::start_receiver`].
     pub async fn accept(&self) -> crate::Result<(Session, SocketAddr)> {
         let (stream, addr) = self.listener.accept().await?;
-
-        let ws = WebSocket::handshake(stream).await?;
-
-        Ok((Session::from_ws(ws), addr))
+        Ok((handshake(stream, &self.config).await?, addr))
     }
 
+    /// Accept connections forever, running `on_conn` for each one.
+    ///
+    /// `on_conn` should register handlers and return; the receiver starts once
+    /// it does, so no message is processed before its handler exists. If it
+    /// returns an error the connection is closed.
     pub async fn session_loop<F, Fut>(&self, on_conn: F) -> crate::Result<()>
     where
         F: Fn(Session, SocketAddr) -> Fut + Send + Sync + 'static,
         Fut: Future<Output = crate::Result<()>> + Send + 'static,
     {
-        let conn_handler = Arc::new(on_conn);
+        let on_conn = Arc::new(on_conn);
 
         loop {
-            let (stream, addr) = self.listener.accept().await?;
-            let conn_handler = conn_handler.clone();
+            let (stream, addr) = match self.listener.accept().await {
+                Ok(conn) => conn,
+                Err(e) => {
+                    // e.g. out of file descriptors: back off instead of exiting.
+                    eprintln!("Accept failed: {e}");
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    continue;
+                }
+            };
+
+            let on_conn = on_conn.clone();
+            let config = self.config.clone();
 
             tokio::spawn(async move {
-                match timeout(
-                    tokio::time::Duration::from_secs(5),
-                    WebSocket::handshake(stream),
-                )
-                .await
-                {
-                    Ok(Ok(ws)) => {
-                        let session = Session::from_ws(ws);
-                        session.start_receiver();
+                let session = match handshake(stream, &config).await {
+                    Ok(session) => session,
+                    Err(e) => {
+                        eprintln!("Handshake failed from {addr}: {e}");
+                        return;
+                    }
+                };
 
-                        if let Err(e) = conn_handler(session, addr).await {
-                            eprintln!("Connection error: {:?}", e);
-                        }
-                    }
-                    Ok(Err(e)) => {
-                        eprintln!("Handshake failed from {}: {:?}", addr, e);
-                    }
-                    Err(_) => {
-                        eprintln!("Handshake failed from {}: Handshake Timeout", addr);
-                    }
+                if let Err(e) = on_conn(session.clone(), addr).await {
+                    eprintln!("Connection error from {addr}: {e}");
+                    let _ = session.close().await;
+                    return;
                 }
+
+                session.start_receiver();
             });
         }
     }
 }
+
+async fn handshake(stream: TcpStream, config: &ServerConfig) -> crate::Result<Session> {
+    let ws_config = WebSocketConfig::default()
+        .max_message_size(Some(config.max_message_size))
+        .max_frame_size(Some(config.max_frame_size));
+
+    let ws = tokio::time::timeout(
+        config.handshake_timeout,
+        tokio_tungstenite::accept_async_with_config(stream, Some(ws_config)),
+    )
+    .await?
+    .map_err(|e| Error::Transport(Box::new(e)))?;
+
+    Ok(Session::from_tungstenite(ws))
+}