creddy/src-tauri/src/server.rs

use core::time::Duration;
use std::io;
use std::net::{SocketAddr, SocketAddrV4};
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::oneshot;
use tokio::time::sleep;

use tauri::{AppHandle, Manager};

use crate::{clientinfo, clientinfo::Client};
use crate::errors::*;
use crate::ipc::{Request, Approval};
use crate::state::AppState;


struct Handler {
    request_id: u64,
    stream: TcpStream,
    receiver: Option<oneshot::Receiver<Approval>>,
    app: AppHandle,
}

impl Handler {
    fn new(stream: TcpStream, app: AppHandle) -> Self {
        let state = app.state::<AppState>();
        let (chan_send, chan_recv) = oneshot::channel();
        let request_id = state.register_request(chan_send);
        Handler { 
            request_id,
            stream,
            receiver: Some(chan_recv),
            app
        }
    }

    async fn handle(mut self) {
        if let Err(e) = self.try_handle().await {
            eprintln!("{e}");
        }
        let state = self.app.state::<AppState>();
        state.unregister_request(self.request_id);
    }

    async fn try_handle(&mut self) -> Result<(), RequestError> {
        let _ = self.recv_request().await?;
        let clients = self.get_clients()?;
        if self.includes_banned(&clients) {
            self.stream.write(b"HTTP/1.0 403 Access Denied\r\n\r\n").await?;
            return Ok(())
        }

        let req = Request {id: self.request_id, clients};
        self.app.emit_all("credentials-request", &req)?;
        let starting_visibility = self.show_window()?;

        match self.wait_for_response().await? {
            Approval::Approved => self.send_credentials().await?,
            Approval::Denied => {
                let state = self.app.state::<AppState>();
                for client in req.clients {
                    state.add_ban(client, self.app.clone());
                }
            }
        }

        // only hide the window if a) it was hidden to start with
        // and b) there are no other pending requests
        let state = self.app.state::<AppState>();
        if !starting_visibility && state.req_count() == 0 {
            let delay = {
                let config = state.config.read().unwrap();
                Duration::from_millis(config.rehide_ms)
            };
            sleep(delay).await;
            let window = self.app.get_window("main").ok_or(RequestError::NoMainWindow)?;
            window.hide()?;
        }

        Ok(())
    }

    async fn recv_request(&mut self) -> Result<Vec<u8>, RequestError> {
        let mut buf = vec![0; 8192]; // it's what tokio's BufReader uses
        let mut n = 0;
        loop {
            n += self.stream.read(&mut buf[n..]).await?;
            if n >= 4 && &buf[(n - 4)..n] == b"\r\n\r\n" {break;}
            if n == buf.len() {return Err(RequestError::RequestTooLarge);}
        }

        if cfg!(debug_assertions) {
            println!("{}", std::str::from_utf8(&buf).unwrap());
        }

        Ok(buf)
    }

    fn get_clients(&self) -> Result<Vec<Option<Client>>, RequestError> {
        let peer_addr = match self.stream.peer_addr()? {
            SocketAddr::V4(addr) => addr,
            _ => unreachable!(), // we only listen on IPv4
        };
        let clients = clientinfo::get_clients(peer_addr.port())?;
        Ok(clients)
    }

    fn includes_banned(&self, clients: &Vec<Option<Client>>) -> bool {
        let state = self.app.state::<AppState>();
        clients.iter().any(|c| state.is_banned(c))
    }

    fn show_window(&self) -> Result<bool, RequestError> {
        let window = self.app.get_window("main").ok_or(RequestError::NoMainWindow)?;
        let starting_visibility = window.is_visible()?;
        if !starting_visibility {
            window.unminimize()?;
            window.show()?;
        }
        window.set_focus()?;
        Ok(starting_visibility)
    }

    async fn wait_for_response(&mut self) -> Result<Approval, RequestError> {
        self.stream.write(b"HTTP/1.0 200 OK\r\n").await?;
        self.stream.write(b"Content-Type: application/json\r\n").await?;
        self.stream.write(b"X-Creddy-delaying-tactic: ").await?;

        #[allow(unreachable_code)] // seems necessary for type inference
        let stall = async {
            let delay = std::time::Duration::from_secs(1);
            loop {
                tokio::time::sleep(delay).await;
                self.stream.write(b"x").await?;
            }
            Ok(Approval::Denied)
        };

        // this is the only place we even read this field, so it's safe to unwrap
        let receiver = self.receiver.take().unwrap();
        tokio::select!{
            r = receiver => Ok(r.unwrap()), // only panics if the sender is dropped without sending, which shouldn't be possible
            e = stall => e,
        }
    }

    async fn send_credentials(&mut self) -> Result<(), RequestError> {
        let state = self.app.state::<AppState>();
        let creds = state.get_creds_serialized()?;

        self.stream.write(b"\r\nContent-Length: ").await?;
        self.stream.write(creds.as_bytes().len().to_string().as_bytes()).await?;
        self.stream.write(b"\r\n\r\n").await?;
        self.stream.write(creds.as_bytes()).await?;
        self.stream.write(b"\r\n\r\n").await?;
        Ok(())
    }
}


pub async fn serve(addr: SocketAddrV4, app_handle: AppHandle) -> io::Result<()> {
    let listener = TcpListener::bind(&addr).await?;
    println!("Listening on {addr}");
    loop {
        match listener.accept().await {
            Ok((stream, _)) => {
                let handler = Handler::new(stream, app_handle.app_handle());
                tauri::async_runtime::spawn(handler.handle());
            },
            Err(e) => {
                eprintln!("Error accepting connection: {e}");
            }
        }
    }
}