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

pub struct HandshakeState {
    noise_key: KeyPair,           // Client's static Curve25519 keypair
    client_payload: Vec<u8>,      // Client identification data
    pattern: &'static str,        // "Noise_XX_25519_AESGCM_SHA256"
    prologue: &'static [u8],      // "WA" header
}

// From src/handshake.rs:29-104
pub async fn do_handshake(
    device: &Device,
    transport: Arc<dyn Transport>,
    transport_events: &mut async_channel::Receiver<TransportEvent>,
) -> Result<Arc<NoiseSocket>>

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,
    UnexpectedEvent(String),
}

pub struct NoiseSocket {
    read_key: Arc<NoiseCipher>,
    read_counter: Arc<AtomicU32>,
    send_job_tx: mpsc::Sender<SendJob>,
    sender_task_handle: JoinHandle<()>,
}

struct SendJob {
    plaintext_buf: Vec<u8>,
    out_buf: Vec<u8>,
    response_tx: oneshot::Sender<SendResult>,
}

impl NoiseSocket {
    pub fn new(
        transport: Arc<dyn Transport>,
        write_key: NoiseCipher,
        read_key: NoiseCipher,
    ) -> Self {
        let (send_job_tx, send_job_rx) = mpsc::channel::<SendJob>(32);
        
        // Spawn dedicated sender task
        let sender_task_handle = tokio::spawn(
            Self::sender_task(transport, write_key, send_job_rx)
        );
        
        Self {
            read_key: Arc::new(read_key),
            read_counter: Arc::new(AtomicU32::new(0)),
            send_job_tx,
            sender_task_handle,
        }
    }
}

async fn sender_task(
    transport: Arc<dyn Transport>,
    write_key: Arc<NoiseCipher>,
    mut send_job_rx: mpsc::Receiver<SendJob>,
) {
    let mut write_counter: u32 = 0;
    
    while let Some(job) = send_job_rx.recv().await {
        let result = Self::process_send_job(
            &transport,
            &write_key,
            &mut write_counter,
            job.plaintext_buf,
            job.out_buf,
        ).await;
        
        // Return buffers to caller for reuse
        let _ = job.response_tx.send(result);
    }
}

if plaintext_buf.len() <= INLINE_ENCRYPT_THRESHOLD {
    // Copy to output buffer and encrypt in-place
    out_buf.clear();
    out_buf.extend_from_slice(&plaintext_buf);
    plaintext_buf.clear();
    
    write_key.encrypt_in_place_with_counter(counter, &mut out_buf)?;
    
    // Frame the ciphertext
    wacore::framing::encode_frame_into(&ciphertext, None, &mut out_buf)?;
}

else {
    let plaintext_arc = Arc::new(plaintext_buf);
    let plaintext_arc_for_task = plaintext_arc.clone();
    
    // Offload to blocking thread
    let spawn_result = tokio::task::spawn_blocking(move || {
        write_key.encrypt_with_counter(
            counter, 
            &plaintext_arc_for_task[..]
        )
    }).await;
    
    // Recover original buffer
    plaintext_buf = Arc::try_unwrap(plaintext_arc)
        .unwrap_or_else(|arc| (*arc).clone());
    
    let ciphertext = spawn_result??;
    
    // Frame and send
    wacore::framing::encode_frame_into(&ciphertext, None, &mut out_buf)?;
    transport.send(out_buf).await?;
}

pub async fn encrypt_and_send(
    &self, 
    plaintext_buf: Vec<u8>, 
    out_buf: Vec<u8>
) -> SendResult {
    let (response_tx, response_rx) = oneshot::channel();
    
    let job = SendJob {
        plaintext_buf,
        out_buf,
        response_tx,
    };
    
    // Send to dedicated task
    if let Err(send_err) = self.send_job_tx.send(job).await {
        let job = send_err.0;
        return Err(EncryptSendError::channel_closed(
            job.plaintext_buf,
            job.out_buf,
        ));
    }
    
    // Wait for result
    match response_rx.await {
        Ok(result) => result,
        Err(_) => Err(EncryptSendError::channel_closed(
            Vec::new(), Vec::new()
        )),
    }
}

type SendResult = Result<(Vec<u8>, Vec<u8>), EncryptSendError>;
//                        ^^^^^^^^^^^^^^^^^^^^^                  
//                        (plaintext_buf, out_buf) for reuse

let mut plaintext_buf = Vec::with_capacity(1024);
let mut out_buf = Vec::with_capacity(1056);  // plaintext + 32 overhead

loop {
    plaintext_buf.clear();
    marshal_to_vec(&message, &mut plaintext_buf)?;
    
    (plaintext_buf, out_buf) = 
        noise_socket.encrypt_and_send(plaintext_buf, out_buf).await?;
}

pub fn decrypt_frame(&self, ciphertext: &[u8]) -> Result<Vec<u8>> {
    let counter = self.read_counter.fetch_add(1, Ordering::SeqCst);
    self.read_key
        .decrypt_with_counter(counter, ciphertext)
        .map_err(|e| SocketError::Crypto(e.to_string()))
}

impl Drop for NoiseSocket {
    fn drop(&mut self) {
        // Abort sender task to prevent resource leaks
        self.sender_task_handle.abort();
    }
}

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)     │
└────────────────┴─────────────────┴────────────────────┘

pub struct FrameDecoder {
    buffer: Vec<u8>,
}

impl FrameDecoder {
    pub fn feed(&mut self, data: &[u8]) {
        self.buffer.extend_from_slice(data);
    }
    
    pub fn decode_frame(&mut self) -> Option<Vec<u8>> {
        if self.buffer.len() < 3 {
            return None;  // Need length header
        }
        
        let len = u32::from_be_bytes([
            0,
            self.buffer[0],
            self.buffer[1],
            self.buffer[2],
        ]) as usize;
        
        if self.buffer.len() < 3 + len {
            return None;  // Incomplete frame
        }
        
        let frame = self.buffer[3..3 + len].to_vec();
        self.buffer.drain(..3 + len);
        Some(frame)
    }
}

#[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,
    DataReceived(Vec<u8>),
}

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 device = self.persistence_manager.get_device_snapshot().await;
        let noise_socket = do_handshake(
            &device,
            transport,
            &mut events
        ).await?;
        
        // Store socket and start receivers
        self.noise_socket.store(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"))?;
    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);
                        frame_decoder.feed(&data);

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

                            // 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.noise_socket.load()
            .ok_or_else(|| anyhow!("not connected"))?;
        
        // 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 — in the `run()` method after the message loop exits — resetting all connection-scoped state:

```rust
impl Client {
    async fn cleanup_connection_state(&self) {
        self.shutdown_notifier.notify(usize::MAX);
        *self.transport.lock().await = None;
        *self.noise_socket.lock().await = 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,
    NoiseHandshake(String),
    Io(String),
    Crypto(String),
}

pub enum EncryptSendError {
    Crypto { error: anyhow::Error, ... },
    Transport { error: anyhow::Error, ... },
    Framing { error: FramingError, ... },
    Join { error: JoinError, ... },
    ChannelClosed { ... },
}

impl EncryptSendError {
    pub fn into_buffers(self) -> (Vec<u8>, Vec<u8>) {
        match self {
            Self::Crypto { plaintext_buf, out_buf, .. } => 
                (plaintext_buf, out_buf),
            Self::Transport { plaintext_buf, out_buf, .. } => 
                (plaintext_buf, out_buf),
            // ...
        }
    }
}

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 {
    self.connect().await;
    self.read_messages_loop().await;

    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);
    sleep(delay).await;
}

// Encrypted size = plaintext + 16 (AES-GCM tag) + 3 (frame header)
let buffer_capacity = plaintext.len() + 32;  // Extra headroom
let out_buf = Vec::with_capacity(buffer_capacity);

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?;

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_encrypt_and_send_returns_both_buffers()

#[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()