lesavka/server/src/main.rs

// lesavka-server - gadget cycle guarded by env
// server/src/main.rs
#[allow(clippy::useless_attribute)] #[forbid(unsafe_code)]
use anyhow::Context as _;
use futures_util::{Stream, StreamExt};
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use std::{backtrace::Backtrace, panic, pin::Pin, sync::Arc};
use tokio::sync::Mutex;
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::Server;
use tonic::{Request, Response, Status};
use tonic_reflection::server::Builder as ReflBuilder;
use tracing::{debug, error, info, warn};

use lesavka_common::lesavka::{
    AudioPacket, Empty, KeyboardReport, MonitorRequest, MouseReport, PasteReply, PasteRequest,
    ResetUsbReply, VideoPacket,
    relay_server::{Relay, RelayServer},
};

use lesavka_server::{
    audio, camera, camera_runtime::CameraRuntime, gadget::UsbGadget, handshake::HandshakeSvc,
    paste, runtime_support, runtime_support::init_tracing, uvc_runtime, video,
};

/*──────────────── constants ────────────────*/
const VERSION: &str = env!("CARGO_PKG_VERSION");
const PKG_NAME: &str = env!("CARGO_PKG_NAME");

/*──────────────── Handler ───────────────────*/
struct Handler {
    kb: Arc<Mutex<tokio::fs::File>>,
    ms: Arc<Mutex<tokio::fs::File>>,
    gadget: UsbGadget,
    did_cycle: Arc<AtomicBool>,
    camera_rt: Arc<CameraRuntime>,
}

impl Handler {
    async fn new(gadget: UsbGadget) -> anyhow::Result<Self> {
        if runtime_support::allow_gadget_cycle() {
            info!("🛠️  Initial USB reset…");
            let _ = gadget.cycle(); // ignore failure - may boot without host
        } else {
            info!(
                "🔒 gadget cycle disabled at startup (set LESAVKA_ALLOW_GADGET_CYCLE=1 to enable)"
            );
        }

        info!("🛠️  opening HID endpoints …");
        let kb = runtime_support::open_with_retry("/dev/hidg0").await?;
        let ms = runtime_support::open_with_retry("/dev/hidg1").await?;
        info!("✅ HID endpoints ready");

        Ok(Self {
            kb: Arc::new(Mutex::new(kb)),
            ms: Arc::new(Mutex::new(ms)),
            gadget,
            did_cycle: Arc::new(AtomicBool::new(false)),
            camera_rt: Arc::new(CameraRuntime::new()),
        })
    }

    async fn reopen_hid(&self) -> anyhow::Result<()> {
        let kb_new = runtime_support::open_with_retry("/dev/hidg0").await?;
        let ms_new = runtime_support::open_with_retry("/dev/hidg1").await?;
        *self.kb.lock().await = kb_new;
        *self.ms.lock().await = ms_new;
        Ok(())
    }
}

/*──────────────── gRPC service ─────────────*/
#[tonic::async_trait]
impl Relay for Handler {
    /* existing streams ─ unchanged, except: no more auto-reset */
    type StreamKeyboardStream = ReceiverStream<Result<KeyboardReport, Status>>;
    type StreamMouseStream = ReceiverStream<Result<MouseReport, Status>>;
    type CaptureVideoStream = Pin<Box<dyn Stream<Item = Result<VideoPacket, Status>> + Send>>;
    type CaptureAudioStream = Pin<Box<dyn Stream<Item = Result<AudioPacket, Status>> + Send>>;
    type StreamMicrophoneStream = ReceiverStream<Result<Empty, Status>>;
    type StreamCameraStream = ReceiverStream<Result<Empty, Status>>;

    async fn stream_keyboard(
        &self,
        req: Request<tonic::Streaming<KeyboardReport>>,
    ) -> Result<Response<Self::StreamKeyboardStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        info!(rpc_id, "⌨️ stream_keyboard opened");
        let (tx, rx) = tokio::sync::mpsc::channel(32);
        let kb = self.kb.clone();
        let ms = self.ms.clone();
        let gadget = self.gadget.clone();
        let did_cycle = self.did_cycle.clone();

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if let Err(e) = runtime_support::write_hid_report(&kb, &pkt.data).await {
                    if e.raw_os_error() == Some(libc::EAGAIN) {
                        debug!(rpc_id, "⌨️  write would block (dropped)");
                    } else {
                        warn!(rpc_id, "⌨️  write failed: {e} (dropped)");
                        runtime_support::recover_hid_if_needed(
                            &e,
                            gadget.clone(),
                            kb.clone(),
                            ms.clone(),
                            did_cycle.clone(),
                        )
                        .await;
                    }
                }
                tx.send(Ok(pkt)).await.ok();
            }
            info!(rpc_id, "⌨️ stream_keyboard closed");
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_mouse(
        &self,
        req: Request<tonic::Streaming<MouseReport>>,
    ) -> Result<Response<Self::StreamMouseStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        info!(rpc_id, "🖱️ stream_mouse opened");
        let (tx, rx) = tokio::sync::mpsc::channel(1024);
        let ms = self.ms.clone();
        let kb = self.kb.clone();
        let gadget = self.gadget.clone();
        let did_cycle = self.did_cycle.clone();

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if let Err(e) = runtime_support::write_hid_report(&ms, &pkt.data).await {
                    if e.raw_os_error() == Some(libc::EAGAIN) {
                        debug!(rpc_id, "🖱️  write would block (dropped)");
                    } else {
                        warn!(rpc_id, "🖱️  write failed: {e} (dropped)");
                        runtime_support::recover_hid_if_needed(
                            &e,
                            gadget.clone(),
                            kb.clone(),
                            ms.clone(),
                            did_cycle.clone(),
                        )
                        .await;
                    }
                }
                tx.send(Ok(pkt)).await.ok();
            }
            info!(rpc_id, "🖱️ stream_mouse closed");
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_microphone(
        &self,
        req: Request<tonic::Streaming<AudioPacket>>,
    ) -> Result<Response<Self::StreamMicrophoneStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        info!(rpc_id, "🎤 stream_microphone opened");
        // 1 ─ build once, early
        let uac_dev = std::env::var("LESAVKA_UAC_DEV").unwrap_or_else(|_| "hw:UAC2Gadget,0".into());
        info!(%uac_dev, "🎤 stream_microphone using UAC sink");
        let mut sink = runtime_support::open_voice_with_retry(&uac_dev)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;

        // 2 ─ dummy outbound stream (same trick as before)
        let (tx, rx) = tokio::sync::mpsc::channel(1);

        // 3 ─ drive the sink in a background task
        tokio::spawn(async move {
            let mut inbound = req.into_inner();
            static CNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

            while let Some(pkt) = inbound.next().await.transpose()? {
                let n = CNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                if n < 10 || n % 300 == 0 {
                    tracing::trace!(rpc_id, "🎤⬇ srv  pkt#{n}  {} bytes", pkt.data.len());
                }
                sink.push(&pkt);
            }
            sink.finish(); // flush on EOS
            let _ = tx.send(Ok(Empty {})).await;
            info!(rpc_id, "🎤 stream_microphone closed");
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_camera(
        &self,
        req: Request<tonic::Streaming<VideoPacket>>,
    ) -> Result<Response<Self::StreamCameraStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        let cfg = camera::current_camera_config();
        info!(
            rpc_id,
            output = cfg.output.as_str(),
            codec = cfg.codec.as_str(),
            width = cfg.width,
            height = cfg.height,
            fps = cfg.fps,
            hdmi = cfg.hdmi.as_ref().map(|h| h.name.as_str()).unwrap_or("none"),
            "🎥 stream_camera output selected"
        );

        let (session_id, relay) = self.camera_rt.activate(&cfg).await?;
        let camera_rt = self.camera_rt.clone();
        info!(rpc_id, session_id, "🎥 stream_camera opened");

        // dummy outbound (same pattern as other streams)
        let (tx, rx) = tokio::sync::mpsc::channel(1);

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if !camera_rt.is_active(session_id) {
                    info!(rpc_id, session_id, "🎥 stream_camera session superseded");
                    break;
                }
                relay.feed(pkt); // ← all logging inside video.rs
            }
            tx.send(Ok(Empty {})).await.ok();
            info!(rpc_id, session_id, "🎥 stream_camera closed");
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn capture_video(
        &self,
        req: Request<MonitorRequest>,
    ) -> Result<Response<Self::CaptureVideoStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        let req = req.into_inner();
        let id = req.id;
        let dev = match id {
            0 => "/dev/lesavka_l_eye",
            1 => "/dev/lesavka_r_eye",
            _ => return Err(Status::invalid_argument("monitor id must be 0 or 1")),
        };
        info!(
            rpc_id,
            id,
            max_bitrate = req.max_bitrate,
            "🎥 capture_video opened"
        );
        debug!(rpc_id, "🎥 streaming {dev}");
        let s = video::eye_ball(dev, id, req.max_bitrate)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;
        Ok(Response::new(Box::pin(s)))
    }

    async fn capture_audio(
        &self,
        req: Request<MonitorRequest>,
    ) -> Result<Response<Self::CaptureAudioStream>, Status> {
        let rpc_id = runtime_support::next_stream_id();
        // Only one speaker stream for now; both 0/1 → same ALSA dev.
        let _id = req.into_inner().id;
        // Allow override (`LESAVKA_ALSA_DEV=hw:2,0` for debugging).
        let dev = std::env::var("LESAVKA_ALSA_DEV").unwrap_or_else(|_| "hw:UAC2Gadget,0".into());
        info!(rpc_id, %dev, "🔊 capture_audio opened");

        let s = audio::ear(&dev, 0)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;

        Ok(Response::new(Box::pin(s)))
    }

    async fn paste_text(&self, req: Request<PasteRequest>) -> Result<Response<PasteReply>, Status> {
        let req = req.into_inner();
        let text = paste::decrypt(&req).map_err(|e| Status::unauthenticated(format!("{e}")))?;
        if let Err(e) = paste::type_text(self.kb.as_ref(), &text).await {
            return Ok(Response::new(PasteReply {
                ok: false,
                error: format!("{e}"),
            }));
        }
        Ok(Response::new(PasteReply {
            ok: true,
            error: String::new(),
        }))
    }

    /*────────────── USB-reset RPC ────────────*/
    async fn reset_usb(&self, _req: Request<Empty>) -> Result<Response<ResetUsbReply>, Status> {
        info!("🔴 explicit ResetUsb() called");
        match self.gadget.cycle() {
            Ok(_) => {
                if let Err(e) = self.reopen_hid().await {
                    error!("💥 reopen HID failed: {e:#}");
                    return Err(Status::internal(e.to_string()));
                }
                Ok(Response::new(ResetUsbReply { ok: true }))
            }
            Err(e) => {
                error!("💥 cycle failed: {e:#}");
                Err(Status::internal(e.to_string()))
            }
        }
    }
}

/*──────────────── main ───────────────────────*/
#[tokio::main(worker_threads = 4)]
async fn main() -> anyhow::Result<()> {
    let _guard = init_tracing()?;
    info!("🚀 {} v{} starting up", PKG_NAME, VERSION);

    panic::set_hook(Box::new(|p| {
        let bt = Backtrace::force_capture();
        error!("💥 panic: {p}\n{bt}");
    }));

    let gadget = UsbGadget::new("lesavka");
    if std::env::var("LESAVKA_DISABLE_UVC").is_err() {
        if std::env::var("LESAVKA_UVC_EXTERNAL").is_ok() {
            info!("📷 UVC control helper external; not spawning");
        } else {
            let bin = uvc_runtime::uvc_ctrl_bin();
            tokio::spawn(uvc_runtime::supervise_uvc_control(bin));
        }
    } else {
        info!("📷 UVC disabled (LESAVKA_DISABLE_UVC set)");
    }
    let handler = Handler::new(gadget.clone()).await?;

    info!("🌐 lesavka-server listening on 0.0.0.0:50051");
    Server::builder()
        .tcp_nodelay(true)
        .max_frame_size(Some(2 * 1024 * 1024))
        .add_service(RelayServer::new(handler))
        .add_service(HandshakeSvc::server())
        .add_service(ReflBuilder::configure().build_v1().unwrap())
        .serve(([0, 0, 0, 0], 50051).into())
        .await?;
    Ok(())
}