One-KVM/src/web/audio_ws.rs

//! Audio WebSocket handler for MJPEG mode
//!
//! Provides a dedicated WebSocket endpoint (`/api/ws/audio`) for streaming
//! Opus-encoded audio data in binary format.
//!
//! ## Binary Protocol
//!
//! Each audio packet is sent as a binary WebSocket message with the following format:
//!
//! ```text
//! Byte 0:      Type (0x02 = audio)
//! Bytes 1-4:   Timestamp (u32 LE, milliseconds since stream start)
//! Bytes 5-6:   Duration (u16 LE, milliseconds)
//! Bytes 7-10:  Sequence (u32 LE)
//! Bytes 11-14: Data length (u32 LE)
//! Bytes 15+:   Opus encoded data
//! ```

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
};
use futures::{SinkExt, StreamExt};
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::broadcast;
use tracing::{debug, info, warn};

use crate::audio::OpusFrame;
use crate::state::AppState;

/// Audio packet type identifier
const AUDIO_PACKET_TYPE: u8 = 0x02;

/// Audio WebSocket upgrade handler
///
/// Upgrades HTTP connection to WebSocket for audio streaming.
pub async fn audio_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
) -> Response {
    ws.on_upgrade(move |socket| handle_audio_socket(socket, state))
}

/// Handle audio WebSocket connection
async fn handle_audio_socket(socket: WebSocket, state: Arc<AppState>) {
    let (mut sender, mut receiver) = socket.split();

    // Try to get Opus frame subscription
    let opus_rx = match state.audio.subscribe_opus_async().await {
        Some(rx) => rx,
        None => {
            warn!("Audio not streaming, rejecting WebSocket connection");
            // Send error message before closing
            let _ = sender
                .send(Message::Text(
                    r#"{"error": "Audio not streaming"}"#.to_string(),
                ))
                .await;
            return;
        }
    };

    let mut opus_rx = opus_rx;
    let stream_start = Instant::now();

    info!("Audio WebSocket client connected");

    // Track connection for cleanup
    let mut closed = false;

    // Use interval instead of sleep for more efficient keepalive
    let mut ping_interval = tokio::time::interval(std::time::Duration::from_secs(30));
    ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);

    loop {
        tokio::select! {
            // Receive Opus frames and send to client
            opus_result = opus_rx.recv() => {
                match opus_result {
                    Ok(frame) => {
                        let binary = encode_audio_packet(&frame, stream_start);
                        if sender.send(Message::Binary(binary)).await.is_err() {
                            debug!("Failed to send audio frame, client disconnected");
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Lagged(n)) => {
                        warn!("Audio WebSocket client lagged by {} frames", n);
                        // Continue - just skip the missed frames
                    }
                    Err(broadcast::error::RecvError::Closed) => {
                        info!("Audio stream closed");
                        break;
                    }
                }
            }

            // Handle client messages (ping/close)
            msg = receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) => {
                        debug!("Audio WebSocket client requested close");
                        closed = true;
                        break;
                    }
                    Some(Ok(Message::Ping(data))) => {
                        if sender.send(Message::Pong(data)).await.is_err() {
                            break;
                        }
                    }
                    Some(Ok(Message::Pong(_))) => {
                        // Pong received, connection is alive
                    }
                    Some(Ok(Message::Text(text))) => {
                        // Handle potential control messages
                        debug!("Received text message on audio WS: {}", text);
                    }
                    Some(Err(e)) => {
                        warn!("Audio WebSocket receive error: {}", e);
                        break;
                    }
                    None => {
                        // Connection closed
                        break;
                    }
                    _ => {}
                }
            }

            // Periodic ping to keep connection alive (using interval)
            _ = ping_interval.tick() => {
                if sender.send(Message::Ping(vec![])).await.is_err() {
                    warn!("Failed to send ping, disconnecting");
                    break;
                }
            }
        }
    }

    if !closed {
        // Try to send close message
        let _ = sender.send(Message::Close(None)).await;
    }

    info!("Audio WebSocket client disconnected");
}

/// Encode Opus frame to binary packet format
///
/// ## Format
///
/// | Offset | Size | Description |
/// |--------|------|-------------|
/// | 0 | 1 | Packet type (0x02 for audio) |
/// | 1 | 4 | Timestamp (u32 LE, ms since start) |
/// | 5 | 2 | Duration (u16 LE, ms) |
/// | 7 | 4 | Sequence number (u32 LE) |
/// | 11 | 4 | Data length (u32 LE) |
/// | 15 | N | Opus encoded data |
fn encode_audio_packet(frame: &OpusFrame, stream_start: Instant) -> Vec<u8> {
    let timestamp_ms = stream_start.elapsed().as_millis() as u32;
    let data_len = frame.data.len() as u32;

    let mut buf = Vec::with_capacity(15 + frame.data.len());

    // Header
    buf.push(AUDIO_PACKET_TYPE);
    buf.extend_from_slice(&timestamp_ms.to_le_bytes());
    buf.extend_from_slice(&(frame.duration_ms as u16).to_le_bytes());
    buf.extend_from_slice(&(frame.sequence as u32).to_le_bytes());
    buf.extend_from_slice(&data_len.to_le_bytes());

    // Opus data
    buf.extend_from_slice(&frame.data);

    buf
}

/// Decode audio packet from binary format (for testing/debugging)
#[allow(dead_code)]
pub fn decode_audio_packet(data: &[u8]) -> Option<AudioPacketHeader> {
    if data.len() < 15 {
        return None;
    }

    if data[0] != AUDIO_PACKET_TYPE {
        return None;
    }

    let timestamp = u32::from_le_bytes([data[1], data[2], data[3], data[4]]);
    let duration_ms = u16::from_le_bytes([data[5], data[6]]);
    let sequence = u32::from_le_bytes([data[7], data[8], data[9], data[10]]);
    let data_length = u32::from_le_bytes([data[11], data[12], data[13], data[14]]);

    Some(AudioPacketHeader {
        packet_type: data[0],
        timestamp,
        duration_ms,
        sequence,
        data_length,
    })
}

/// Audio packet header (for decoding/testing)
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct AudioPacketHeader {
    pub packet_type: u8,
    pub timestamp: u32,
    pub duration_ms: u16,
    pub sequence: u32,
    pub data_length: u32,
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;

    #[test]
    fn test_encode_decode_packet() {
        let frame = OpusFrame {
            data: Bytes::from(vec![1, 2, 3, 4, 5]),
            duration_ms: 20,
            sequence: 42,
            timestamp: Instant::now(),
            rtp_timestamp: 0,
        };

        let stream_start = Instant::now();
        let encoded = encode_audio_packet(&frame, stream_start);

        assert!(encoded.len() >= 15);
        assert_eq!(encoded[0], AUDIO_PACKET_TYPE);

        let header = decode_audio_packet(&encoded).unwrap();
        assert_eq!(header.packet_type, AUDIO_PACKET_TYPE);
        assert_eq!(header.duration_ms, 20);
        assert_eq!(header.sequence, 42);
        assert_eq!(header.data_length, 5);
    }

    #[test]
    fn test_decode_invalid_packet() {
        // Too short
        assert!(decode_audio_packet(&[]).is_none());
        assert!(decode_audio_packet(&[0x02; 10]).is_none());

        // Wrong type
        let mut bad = vec![0x01; 20];
        assert!(decode_audio_packet(&bad).is_none());
    }
}