joecalsend/src/lib.rs

pub mod discovery;
pub mod error;
pub mod http_server;
pub mod models;
pub mod transfer;

use std::{
    collections::HashMap,
    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
    sync::{Arc, OnceLock},
};

use julid::Julid;
use log::error;
use models::Device;
use serde::{Deserialize, Serialize};
use tokio::{
    net::UdpSocket,
    sync::{
        Mutex,
        mpsc::{self, UnboundedSender},
    },
    task::JoinSet,
};
use transfer::Session;

pub const DEFAULT_PORT: u16 = 53317;
pub const MULTICAST_IP: Ipv4Addr = Ipv4Addr::new(224, 0, 0, 167);
pub const LISTENING_SOCKET_ADDR: SocketAddrV4 =
    SocketAddrV4::new(Ipv4Addr::from_bits(0), DEFAULT_PORT);

pub type ShutdownSender = mpsc::Sender<()>;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Listeners {
    Udp,
    Http,
    Multicast,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UploadDialog {
    UploadDeny,
    UploadConfirm,
}

pub enum TransferEvent {
    Sent,
    Received,
    Failed,
    UploadRequest { alias: String, id: Julid },
}

/// Contains the main network and backend state for an application session.
#[derive(Clone)]
pub struct JoecalState {
    pub device: Device,
    pub peers: Arc<Mutex<HashMap<String, (SocketAddr, Device)>>>,
    pub sessions: Arc<Mutex<HashMap<String, Session>>>, // Session ID to Session
    pub running_state: Arc<Mutex<RunningState>>,
    pub socket: Arc<UdpSocket>,
    pub client: reqwest::Client,
    upload_requests: Arc<Mutex<HashMap<Julid, UnboundedSender<UploadDialog>>>>,
    shutdown_sender: OnceLock<ShutdownSender>,
    // the receiving end will be held by the application so it can update the UI based on backend
    // events
    transfer_event_tx: UnboundedSender<TransferEvent>,
}

impl JoecalState {
    pub async fn new(
        device: Device,
        transfer_event_tx: UnboundedSender<TransferEvent>,
    ) -> crate::error::Result<Self> {
        let socket = UdpSocket::bind(LISTENING_SOCKET_ADDR).await?;
        socket.set_multicast_loop_v4(true)?;
        socket.set_multicast_ttl_v4(2)?; // one hop out from localnet
        socket.join_multicast_v4(MULTICAST_IP, Ipv4Addr::from_bits(0))?;

        Ok(Self {
            device,
            client: reqwest::Client::new(),
            socket: socket.into(),
            peers: Default::default(),
            sessions: Default::default(),
            running_state: Default::default(),
            shutdown_sender: Default::default(),
            upload_requests: Default::default(),
            transfer_event_tx,
        })
    }

    pub async fn start(&self, config: &Config, handles: &mut JoinSet<Listeners>) {
        let state = self.clone();
        let konfig = config.clone();
        handles.spawn({
            let (tx, shutdown_rx) = mpsc::channel(1);
            let _ = self.shutdown_sender.set(tx);
            async move {
                if let Err(e) = state.start_http_server(shutdown_rx, &konfig).await {
                    error!("HTTP server error: {e}");
                }
                Listeners::Http
            }
        });
        let state = self.clone();
        let konfig = config.clone();
        handles.spawn({
            async move {
                if let Err(e) = state.listen_multicast(&konfig).await {
                    error!("UDP listener error: {e}");
                }
                Listeners::Multicast
            }
        });

        let state = self.clone();
        let config = config.clone();
        handles.spawn({
            async move {
                loop {
                    let rstate = state.running_state.lock().await;
                    if *rstate == RunningState::Stopping {
                        break;
                    }
                    if let Err(e) = state.announce(None, &config).await {
                        error!("Announcement error: {e}");
                    }
                    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                }
                Listeners::Udp
            }
        });

        // TODO: add a task that periodically clears out the upload requests if
        // they're too old; the keys are julids so they have the time in them
    }

    pub async fn stop(&self) {
        let mut rstate = self.running_state.lock().await;
        *rstate = RunningState::Stopping;
        let _ = self
            .shutdown_sender
            .get()
            .expect("Could not get stop signal transmitter")
            .send(())
            .await;
    }

    pub async fn refresh_peers(&self) {
        let mut peers = self.peers.lock().await;
        peers.clear();
    }

    pub async fn get_upload_request(&self, id: Julid) -> Option<UnboundedSender<UploadDialog>> {
        self.upload_requests.lock().await.get(&id).cloned()
    }

    pub async fn clear_upload_request(&self, id: Julid) {
        let _ = self.upload_requests.lock().await.remove(&id);
    }

    /// Add a transmitter for an upload request confirmation dialog that the
    /// application frontend can use to tell the Axum handler whether or not to
    /// accept the upload.
    ///
    /// IMPORTANT! Be sure to call `clear_upload_request(id)` when you're done
    /// getting an answer back/before you exit!
    pub async fn add_upload_request(&self, id: Julid, tx: UnboundedSender<UploadDialog>) {
        self.upload_requests.lock().await.entry(id).insert_entry(tx);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RunningState {
    Running,
    Stopping,
}

impl Default for RunningState {
    fn default() -> Self {
        Self::Running
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
    pub multicast_addr: SocketAddrV4,
    pub port: u16,
    pub download_dir: String,
}

impl Default for Config {
    fn default() -> Self {
        let home = std::env::home_dir().unwrap_or("/tmp".into());
        let dd = home.join("joecalsend-downloads");
        Self {
            multicast_addr: SocketAddrV4::new(MULTICAST_IP, DEFAULT_PORT),
            port: DEFAULT_PORT,
            download_dir: dd.to_string_lossy().into(),
        }
    }
}