WebSocket & Noise Protocol

src/
├── handshake.rs           # Noise protocol handshake
├── socket/
│   ├── noise_socket.rs    # Encrypted communication layer
│   ├── mod.rs            # Re-exports
│   └── error.rs          # Socket errors
├── transport/            # WebSocket transport abstraction
└── client.rs            # High-level client orchestration

// src/handshake.rs
enum HandshakePattern {
    /// Cold start / pairing / forced fallback after an earlier IK failure.
    Xx,
    /// Cached server static + valid cert chain available; attempt IK.
    Ik([u8; 32]),
}

// wacore/noise/src/handshake.rs
pub struct XxHandshakeState { /* ... */ }
pub struct IkHandshakeState { /* ... */ }
pub struct XxFallbackHandshakeState { /* ... */ }

pub enum IkServerHelloOutcome {
    /// Server accepted the cached static — IK can complete in one round trip.
    Continue(Box<IkHandshakeOutcome>),
    /// Server rejected the cached static; pivot the transcript into XXfallback.
    Fallback(Box<IkFallbackInputs>),
}

// src/handshake.rs
pub async fn do_handshake(
    runtime: Arc<dyn Runtime>,
    persistence_manager: &PersistenceManager,
    ik_handshake_failures: &AtomicU32,
    transport: Arc<dyn Transport>,
    transport_events: &mut async_channel::Receiver<TransportEvent>,
    observers: SendObservers,
) -> Result<Arc<NoiseSocket>>

// src/handshake.rs:run_ik_handshake
let mut ik = IkHandshakeState::new(
    device.noise_key.clone(),
    server_static_pub,
    client_payload,
    &WA_CONN_HEADER,
)?;

// IK ClientHello carries: ephemeral, encrypted client static, encrypted 0-RTT payload.
let client_hello_bytes = ik.build_client_hello()?;
send_first_handshake_message(&transport, device, &client_hello_bytes).await?;

let resp_frame = recv_frame(runtime, transport_events, &mut frame_decoder).await?;

match ik.read_server_hello(&resp_frame)? {
    IkServerHelloOutcome::Continue(out) => {
        // Single round trip — handshake is already complete.
        info!("Handshake complete (IK), switching to encrypted communication");
        // No client cert chain refresh — the cached entry stays authoritative.
    }
    IkServerHelloOutcome::Fallback(inputs) => {
        // Server rejected the cached static. Pivot to XXfallback below.
    }
}

// src/handshake.rs
let mut fb = XxFallbackHandshakeState::from_ik_failure(*inputs, &WA_CONN_HEADER)?;
let client_finish_bytes = fb.build_client_finish()?;

let framed = wacore::framing::encode_frame(&client_finish_bytes, None)?;
transport.send(bytes::Bytes::from(framed)).await?;

let outcome = fb.finish()?;
info!("Handshake complete (XXfallback), switching to encrypted communication");

// XXfallback re-derives the server cert chain, so persist the fresh value.
persistence_manager
    .process_command(DeviceCommand::SetServerCertChain(
        outcome.server_cert_chain.into(),
    ))
    .await;

impl HandshakeError {
    /// Transient — never invalidate the cache.
    pub fn is_transient(&self) -> bool { /* Transport, Timeout, Disconnected, StreamClosed */ }

    /// Crypto-fatal — the cached server static or cert chain is no longer trustworthy.
    pub fn is_crypto_fatal(&self) -> bool { /* AEAD/cert/parse failures */ }
}

pub fn build_handshake_header(
    edge_routing_info: Option<&[u8]>
) -> (Vec<u8>, bool) {
    let mut header = WA_CONN_HEADER.to_vec();  // [0x57, 0x41]
    
    if let Some(info) = edge_routing_info {
        if info.len() <= 8192 {  // Max size
            header.extend_from_slice(info);
            return (header, true);
        }
    }
    
    (header, false)
}

pub enum HandshakeError {
    Transport(#[from] anyhow::Error),
    Core(#[from] CoreHandshakeError),
    Timeout,
    /// Producer side of `transport_events` was dropped — distinct from a
    /// timeout because nothing more will ever arrive on the channel.
    StreamClosed,
    Disconnected,
    UnexpectedEvent(String),
}

pub struct NoiseSocket {
    read_key: Arc<NoiseCipher>,
    read_counter: Arc<AtomicU32>,
    /// Channel to send jobs to the dedicated sender task. A channel instead of
    /// a mutex lets callers enqueue their work and await the result without
    /// blocking on the send in progress.
    send_job_tx: async_channel::Sender<SendJob>,
    /// Aborts the sender task on drop (prevents resource leaks if the task is
    /// stuck on a slow/hanging network operation).
    _sender_task_handle: AbortHandle,
}

/// A job sent to the dedicated sender task.
struct SendJob {
    plaintext: bytes::Bytes,
    response_tx: oneshot::Sender<SendResult>,
}

/// What a socket reports its sends to. Both halves belong to the `Client`; a
/// VoIP relay socket and most tests pass `Default`, reporting to neither.
#[derive(Default, Clone)]
pub struct SendObservers {
    /// Wire-byte accounting, recorded after the transport write.
    stats: Option<Arc<wacore::stats::SessionStats>>,
    /// Publisher for the plaintext of each frame that reached the transport.
    sent_frames: Option<Arc<SentFrameTap>>,
}

impl SendObservers {
    /// Report wire bytes into `stats` and nothing else.
    pub fn with_stats(stats: Arc<wacore::stats::SessionStats>) -> Self { /* ... */ }

    /// Also publish each sent frame's plaintext through `tap`.
    pub(crate) fn with_sent_frames(mut self, tap: Arc<SentFrameTap>) -> Self { /* ... */ }
}

impl NoiseSocket {
    pub fn with_observers(
        runtime: Arc<dyn Runtime>,
        transport: Arc<dyn Transport>,
        write_key: NoiseCipher,
        read_key: NoiseCipher,
        observers: SendObservers,
    ) -> Self {
        let write_key = Arc::new(write_key);
        let read_key = Arc::new(read_key);

        // Small buffer matched to typical steady-state throughput; the sender
        // task is network-bound (awaits transport.send), so a transient
        // WebSocket stall backpressures producers here rather than queuing.
        let (send_job_tx, send_job_rx) = async_channel::bounded::<SendJob>(8);

        let sender_task_handle = runtime.spawn(Box::pin(Self::sender_task(
            runtime.clone(),
            transport.clone(),
            write_key.clone(),
            send_job_rx,
            observers,
        )));

        Self {
            read_key,
            read_counter: Arc::new(AtomicU32::new(0)),
            send_job_tx,
            _sender_task_handle: sender_task_handle,
        }
    }
}

async fn sender_task(
    runtime: Arc<dyn Runtime>,
    transport: Arc<dyn Transport>,
    write_key: Arc<NoiseCipher>,
    send_job_rx: async_channel::Receiver<SendJob>,
    observers: SendObservers,
) {
    let SendObservers { stats, sent_frames } = observers;
    let mut write_counter: u32 = 0;
    let mut enc_buf = Vec::with_capacity(4096);
    let mut out_buf = BytesMut::with_capacity(4096);
    let mut poisoned = false;
    let mut waiters: Vec<(oneshot::Sender<SendResult>, usize)> = Vec::new();
    let mut carry_over: Option<SendJob> = None;
    // Plaintexts of this batch's frames, held only while a consumer is
    // watching; empty (unallocated) otherwise. Kept until after the write so
    // what is published is what the transport actually accepted.
    let mut observed: Vec<bytes::Bytes> = Vec::new();

    loop {
        let job = match carry_over.take() {
            Some(job) => job,
            None => match send_job_rx.recv().await {
                Ok(job) => job,
                Err(_) => break,
            },
        };
        if poisoned {
            let _ = job.response_tx.send(Err(EncryptSendError::poisoned()));
            continue;
        }

        // Encrypt everything already queued into out_buf, cloning each
        // plaintext into `observed` first when `sent_frames` is enabled. Stop
        // at the MAX_BATCH_FRAMES / MAX_BATCH_WIRE_BYTES ceiling, or when
        // try_recv finds nothing more waiting.
        // ... encrypt_frame_into loop, then a single transport.send(out_buf) ...

        // After the write succeeds: stats.record_frame_sent for each wire
        // frame, then — re-checking sent_frames.enabled() rather than trusting
        // the read at capture time, so a batch that outlived its last lease
        // stays quiet — tap.publish(plaintext) for each entry in `observed`.
        // A write that fails clears `observed` instead: a frame the transport
        // refused is not reported as sent.
    }
}

if plaintext.len() <= INLINE_ENCRYPT_THRESHOLD {
    enc_buf.clear();
    enc_buf.extend_from_slice(&plaintext);
    write_key.encrypt_in_place_with_counter(counter, enc_buf)?;

    // append_frame_into, not encode_frame_into: it appends to out_buf
    // instead of clearing it, so several frames can be laid out back to
    // back in one buffer for a single write.
    wacore::framing::append_frame_into(enc_buf, None, out_buf)?;
}

else {
    let write_key = write_key.clone();
    let ciphertext = wacore::runtime::blocking(&**runtime, move || {
        write_key.encrypt_with_counter(counter, &plaintext)
    })
    .await??;

    wacore::framing::append_frame_into(&ciphertext, None, out_buf)?;
}

*write_counter = counter + 1;
Ok(out_buf.len() - before)

pub(crate) async fn enqueue_send(
    &self,
    plaintext: bytes::Bytes,
) -> std::result::Result<oneshot::Receiver<SendResult>, EncryptSendError> {
    let (response_tx, response_rx) = oneshot::channel();
    let job = SendJob { plaintext, response_tx };

    if let Err(_send_err) = self.send_job_tx.send(job).await {
        return Err(EncryptSendError::channel_closed());
    }

    Ok(response_rx)
}

pub(crate) async fn await_send(receiver: oneshot::Receiver<SendResult>) -> SendResult {
    match receiver.await {
        Ok(result) => result,
        Err(_) => Err(EncryptSendError::channel_closed()),
    }
}

pub async fn encrypt_and_send(&self, plaintext: bytes::Bytes) -> SendResult {
    let receiver = self.enqueue_send(plaintext).await?;
    Self::await_send(receiver).await
}

pub fn decrypt_frame(&self, mut ciphertext: BytesMut) -> Result<BytesMut> {
    let counter = self
        .read_counter
        .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |c| c.checked_add(1))
        .map_err(|_| SocketError::Cipher(NoiseError::CounterExhausted))?;
    self.read_key
        .decrypt_in_place_with_counter(counter, &mut ciphertext)
        .map_err(SocketError::Cipher)?;
    Ok(ciphertext)
}

pub fn encode_frame(
    payload: &[u8], 
    header: Option<&[u8]>
) -> Result<Vec<u8>> {
    let mut output = Vec::new();
    
    if let Some(header) = header {
        output.extend_from_slice(header);
    }
    
    // 3-byte big-endian length
    let len = payload.len() as u32;
    output.push(((len >> 16) & 0xFF) as u8);
    output.push(((len >> 8) & 0xFF) as u8);
    output.push((len & 0xFF) as u8);
    
    output.extend_from_slice(payload);
    Ok(output)
}

┌────────────────┬─────────────────┬────────────────────┐
│ Optional Header│  Length (3 bytes)│     Payload       │
│  (handshake)   │   (big-endian)  │   (encrypted)     │
└────────────────┴─────────────────┴────────────────────┘

use bytes::{Buf, BytesMut};

const CHUNK_SIZE: usize = 1024;
const MAX_IDLE_CAPACITY: usize = 64 * 1024;
const MAX_EAGER_RESERVE: usize = 64 * 1024;

pub struct FrameDecoder {
    buffer: BytesMut,
    /// Whether the buffer has ever grown past `MAX_IDLE_CAPACITY`.
    grew_oversized: bool,
}

impl FrameDecoder {
    pub fn new() -> Self {
        Self {
            buffer: BytesMut::with_capacity(CHUNK_SIZE),
            grew_oversized: false,
        }
    }

    pub fn feed(&mut self, data: &[u8]) {
        self.reserve_for(data.len());
        self.buffer.extend_from_slice(data);
    }

    fn reserve_for(&mut self, incoming: usize) {
        if self.buffer.capacity() - self.buffer.len() < incoming {
            self.buffer.reserve(incoming.max(CHUNK_SIZE));
            self.grew_oversized |= self.buffer.capacity() > MAX_IDLE_CAPACITY;
        }
    }

    pub fn decode_frame(&mut self) -> Option<BytesMut> {
        if self.buffer.len() < FRAME_LENGTH_SIZE {
            return None;
        }

        let frame_len = ((self.buffer[0] as usize) << 16)
            | ((self.buffer[1] as usize) << 8)
            | (self.buffer[2] as usize);

        if self.buffer.len() < FRAME_LENGTH_SIZE + frame_len {
            // Size the buffer ahead for the rest of a large frame instead of
            // growing geometrically across the reads that carry it.
            let missing = FRAME_LENGTH_SIZE + frame_len - self.buffer.len();
            self.reserve_for(missing.min(MAX_EAGER_RESERVE));
            return None;
        }

        self.buffer.advance(FRAME_LENGTH_SIZE);
        let frame_data = self.buffer.split_to(frame_len);

        // Copy a short residue into a fresh chunk so an oversized allocation
        // doesn't get pinned by a few leftover bytes of the next frame.
        if (self.grew_oversized || self.buffer.capacity() > MAX_IDLE_CAPACITY)
            && self.buffer.len() <= CHUNK_SIZE
        {
            let mut fresh = BytesMut::with_capacity(CHUNK_SIZE);
            fresh.extend_from_slice(&self.buffer);
            self.buffer = fresh;
            self.grew_oversized = false;
        }

        Some(frame_data)
    }
}

use bytes::Bytes;

#[async_trait]
pub trait Transport: Send + Sync {
    async fn send(&self, data: Vec<u8>) -> Result<(), anyhow::Error>;
    async fn disconnect(&self);
}

pub enum TransportEvent {
    Connected,
    Disconnected,
    /// Inbound payload, fed into the frame decoder's accumulation buffer
    /// and dropped immediately after — see [Frame Decoder](#frame-decoder).
    DataReceived(Bytes),
}

pub fn from_websocket<S>(
    ws: WebSocketStream<S>,
) -> (Arc<dyn Transport>, async_channel::Receiver<TransportEvent>)
where
    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,

const TRANSPORT_CONNECT_TIMEOUT: Duration = Duration::from_secs(20);

let version_future = tokio::time::timeout(
    TRANSPORT_CONNECT_TIMEOUT,
    resolve_and_update_version(&persistence_manager, &http_client, override_version),
);
let transport_future = tokio::time::timeout(
    TRANSPORT_CONNECT_TIMEOUT,
    transport_factory.create_transport(),
);

let (version_result, transport_result) = tokio::join!(version_future, transport_future);

impl Client {
    pub async fn connect(&self) -> Result<()> {
        // Version fetch + transport connection run in parallel, both under 20s timeout
        let (transport, mut events) = tokio::time::timeout(
            TRANSPORT_CONNECT_TIMEOUT,
            self.transport_factory.create_transport(),
        ).await??;
        
        // Perform Noise handshake
        let noise_socket = do_handshake(
            self.runtime.clone(),
            &self.persistence_manager,
            &self.ik_handshake_failures,
            transport,
            &mut events,
            SendObservers::with_stats(self.stats.clone())
                .with_sent_frames(self.sent_frame_tap.clone()),
        ).await?;
        
        // Store socket and start receivers
        *self.noise_socket.lock().unwrap_or_else(|p| p.into_inner()) = Some(Arc::clone(&noise_socket));
        self.start_frame_receiver(events, noise_socket).await;
        
        Ok(())
    }
}

async fn read_messages_loop(self: &Arc<Self>) -> Result<(), anyhow::Error> {
    let transport_events = self.transport_events.lock().await.take()
        .ok_or_else(|| anyhow!("Cannot start message loop: not connected"))?;
    // Noise socket is stable for the lifetime of this loop; resolve once
    // instead of locking the mutex on every inbound frame.
    let noise_socket = self.get_noise_socket()
        .map_err(|_| anyhow!("Cannot start message loop: no noise socket"))?;
    let mut frame_decoder = FrameDecoder::new();

    loop {
        futures::select_biased! {
            _ = self.shutdown_notifier.listen().fuse() => {
                return Ok(());
            },
            event_result = transport_events.recv().fuse() => {
                match event_result {
                    Ok(TransportEvent::DataReceived(data)) => {
                        // Update dead-socket timer on arrival
                        self.last_data_received_ms.store(now_millis(), Ordering::Relaxed);
                        // Dropped before any await below: `data` is a view into the
                        // websocket's shared read buffer, so holding it while a node
                        // is processed would keep that allocation alive alongside
                        // the decoder's own copy of the same bytes.
                        frame_decoder.feed(&data);
                        drop(data);

                        let mut frames_in_batch: u32 = 0;
                        while let Some(encrypted_frame) = frame_decoder.decode_frame() {
                            if let Some(node) = self.decrypt_frame(&noise_socket, encrypted_frame) {
                                // Inline vs spawned processing (see below)
                                if is_critical(&node) {
                                    self.process_decrypted_node(node).await;
                                } else {
                                    self.runtime.spawn_detached(/* ... */);
                                }
                            }

                            // Cooperative yield every N frames
                            frames_in_batch += 1;
                            if frames_in_batch.is_multiple_of(self.runtime.yield_frequency()) {
                                if let Some(yield_fut) = self.runtime.yield_now() {
                                    yield_fut.await;
                                }
                            }
                        }

                        // Refresh timestamp after batch so keepalive sees
                        // batch completion time, not just arrival time
                        if frames_in_batch > 1 {
                            self.last_data_received_ms.store(now_millis(), Ordering::Relaxed);
                        }
                    },
                    Ok(TransportEvent::Disconnected) | Err(_) => { /* handle disconnect */ }
                    _ => {}
                }
            }
        }
    }
}

### 3. Send message

```rust
impl Client {
    pub async fn send_node(&self, node: &Node) -> Result<()> {
        let noise_socket = self.get_noise_socket()?;
        
        // Marshal node to binary with auto-sized buffer
        let plaintext_buf = marshal_auto(node)?;
        
        // Encrypt and send
        self.send_raw_bytes(plaintext_buf).await
    }
}

### 4. Disconnect

On disconnect, `cleanup_connection_state()` runs exactly once, after the message loop exits. As of PR #1258 that call site is the shared `drive_connection()` helper, used by both `run()`'s loop body and [`Connection::read_until_disconnected()`](/api/client#connection) — so the reset below happens the same way whether `run()` or a direct `connect()` drove the connection:

```rust
impl Client {
    async fn cleanup_connection_state(&self) {
        self.shutdown_notifier.notify(usize::MAX);
        *self.transport.lock().await = None;
        *self.noise_socket.lock().unwrap_or_else(|p| p.into_inner()) = None;
        self.is_connected.store(false, Ordering::Release);

        // Drop per-chat lane senders so workers exit via channel close.
        // Without this, stale workers from the old connection survive reconnects
        // holding outdated signal/crypto state.
        self.chat_lanes.invalidate_all();

        // Clear signal cache, pending retries, IQ waiters, offline sync state...
    }
}

// Lock-free connection check — never affected by mutex contention
pub fn is_connected(&self) -> bool {
    self.is_connected.load(Ordering::Acquire)
}

pub enum SocketError {
    SocketClosed,
    Io(#[from] std::io::Error),
    Cipher(#[from] NoiseError),
    Marshal(#[source] BinaryError),
}

pub enum EncryptSendErrorKind {
    Crypto,
    Framing,
    Transport,
    Join,
    ChannelClosed,
    /// A previous frame failed at the transport; this connection's write
    /// keystream can no longer be extended safely. See "Send poisoning"
    /// below.
    Poisoned,
}

pub struct EncryptSendError {
    pub kind: EncryptSendErrorKind,
    pub source: anyhow::Error,
}

Server sends <stream:error code="..."/>
  → read_messages_loop (inline, not spawned)
  → StanzaRouter → StreamErrorHandler
  → Client::handle_stream_error()
  → is_logged_in = false
  → Code-specific handling (see table above)
  → shutdown_notifier fires → keepalive exits
  → run() loop checks enable_auto_reconnect

Sequence: 1s, 1s, 2s, 3s, 5s, 8s, 13s, 21s, 34s, 55s, 89s, 144s, ...
Maximum:  900s (15 minutes)
Jitter:   ±10%

// Simplified reconnection logic in run()
loop {
    match self.connect().await {
        Ok(connection) => {
            // read_until_disconnected() drives the same connection lifecycle
            // (read loop, teardown, Disconnected dispatch) that run() has
            // always used — see the `Connection` type below.
            connection.read_until_disconnected().await;
        }
        Err(ConnectError::Shutdown) => break, // shutdown is sticky and final; never retry
        Err(e) => { /* log and fall through to the retry delay */ }
    }

    if !self.enable_auto_reconnect.load(Ordering::Relaxed) {
        break; // 401, 409, 516 — stop permanently
    }

    if self.expected_disconnect.load(Ordering::Relaxed) {
        continue; // 515 — reconnect immediately
    }

    let errors = self.auto_reconnect_errors.fetch_add(1, Ordering::Relaxed);
    let delay = fibonacci_backoff(errors + 1);

    // This races the wait against the client's terminal shutdown listener.
    // If you call `disconnect()`, `logout()`, or `signal_shutdown_sync()`,
    // the wait ends immediately. A bare `sleep(delay).await` would otherwise
    // hold `run()` for up to the 900s cap after you'd already asked it to stop.
    select! {
        _ = sleep(delay) => {}
        _ = self.shutdown_signal() => break,
    }
}

pub struct NoiseCipher {
    cipher: Aes256Gcm,  // Uses AES-NI on x86_64
}

// Bad: Allocates new buffer
let data = node.to_bytes();
socket.send(data).await?;

// Good: Reuses buffer
let mut buf = Vec::with_capacity(1024);
marshal_to_vec(&node, &mut buf)?;
socket.send(buf).await?;

frame_decoder.feed(&data);
drop(data); // release the transport's view before any await below

// Each frame is split out of the shared accumulation buffer, so the
// downstream freeze() is a pointer move rather than a fresh allocation.
while let Some(frame) = frame_decoder.decode_frame() { /* ... */ }

pub struct MockTransport;

#[async_trait]
impl Transport for MockTransport {
    async fn send(&self, data: Vec<u8>) -> Result<()> {
        // Record for assertions
        Ok(())
    }
    async fn disconnect(&self) {}
}

#[tokio::test]
async fn test_concurrent_sends_maintain_order()

#[tokio::test]
async fn test_encrypted_buffer_sizing_is_sufficient()

#[tokio::test]
async fn test_handshake_with_edge_routing()

// Write batching (frame coalescing)
#[tokio::test]
async fn queued_frames_leave_in_one_write_in_counter_order()

#[tokio::test]
async fn a_batch_never_overshoots_the_byte_ceiling()

#[tokio::test]
async fn every_encrypted_frame_burns_its_own_counter()

#[tokio::test]
async fn the_transport_cause_reaches_the_caller()