maj/src/server/mod.rs

use crate::{Response, StatusCode};
use async_trait::async_trait;
use rustls::{Certificate, Session};
use std::{error::Error as StdError, net::SocketAddr, sync::Arc};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::{net::TcpListener, stream::StreamExt};
use tokio_rustls::TlsAcceptor;
use url::Url;

/// A Gemini request and its associated metadata.
#[allow(dead_code)]
pub struct Request {
    pub url: Url,
    pub certs: Option<Vec<Certificate>>,
}

pub type Error = Box<dyn StdError + Sync + Send>;

#[allow(dead_code, unused_assignments, unused_mut, unused_variables)]
mod routes;
pub use routes::*;

#[async_trait]
pub trait Handler {
    async fn handle(&self, r: Request) -> Result<Response, Error>;
}

pub async fn serve(
    h: &(dyn Handler + Sync),
    cfg: rustls::ServerConfig,
    host: String,
    port: u16,
) -> Result<(), Error>
where
{
    let cfg = Arc::new(cfg);
    let mut listener = TcpListener::bind(&format!("{}:{}", host, port)).await?;
    let mut incoming = listener.incoming();
    let acceptor = TlsAcceptor::from(cfg.clone());

    while let Some(stream) = incoming.next().await {
        let stream = stream?;
        let addr = stream.peer_addr().unwrap();
        let fut = async {
            let acceptor = acceptor.clone();
            let result = acceptor.accept(stream).await;

            if result.is_err() {
                return;
            }

            let mut stream = result.unwrap();

            let mut rd = BufReader::new(&mut stream);
            let mut u = String::new();
            if let Err(why) = rd.read_line(&mut u).await {
                log::error!("can't read request from {}: {:?}", addr, why);
                let _ = stream
                    .write(format!("{} Invalid URL", StatusCode::BadRequest as u8).as_bytes())
                    .await;
                return;
            }
            u = u.trim().to_string();
            if u.len() >= 1025 {
                let _ = stream
                    .write(format!("{} URL too long", StatusCode::BadRequest as u8).as_bytes())
                    .await;
                return;
            }

            if u.starts_with("//") {
                u = format!("gemini:{}", u);
            }

            match Url::parse(&u) {
                Err(why) => {
                    let _ = stream
                        .write(
                            format!("{} bad URL: {:?}", StatusCode::BadRequest as u8, why)
                                .as_bytes(),
                        )
                        .await;
                }
                Ok(u) => {
                    if u.scheme() != "gemini" {
                        let _ = stream
                            .write(
                                format!(
                                    "{} Cannot handle that kind of url",
                                    StatusCode::ProxyRequestRefused as u8
                                )
                                .as_bytes(),
                            )
                            .await;
                        return;
                    }

                    if let Some(u_port) = u.port() {
                        if port != u_port {
                            let _ = stream
                                .write(
                                    format!(
                                        "{} Cannot handle that kind of url",
                                        StatusCode::ProxyRequestRefused as u8
                                    )
                                    .as_bytes(),
                                )
                                .await;
                            return;
                        }
                    }

                    tokio::join!(handle(
                        h,
                        Request {
                            url: u.clone(),
                            certs: stream.get_ref().1.get_peer_certificates(),
                        },
                        &mut stream,
                        addr,
                    ));
                }
            }
        };

        tokio::join!(fut);
    }

    Ok(())
}

async fn handle<T>(h: &(dyn Handler + Sync), req: Request, stream: &mut T, addr: SocketAddr)
where
    T: AsyncWriteExt + Unpin,
{
    let u = req.url.clone();
    match h.handle(req).await {
        Ok(resp) => {
            stream
                .write(format!("{} {}\r\n", resp.status as u8, resp.meta).as_bytes())
                .await
                .unwrap();
            stream.write(&resp.body).await.unwrap();
            log::info!("{}: {} {:?}", addr, u, resp.status);
        }
        Err(why) => {
            stream
                .write(format!("{} {:?}\r\n", StatusCode::PermanentFailure as u8, why).as_bytes())
                .await
                .unwrap();
            log::error!("{}: {}: {:?}", addr, u, why);
        }
    };
}