pukram-server/src/main.rs

use std::{
    collections::{HashMap, hash_map::Entry},
    env,
    fs::File,
    io::{BufReader, Error, ErrorKind, Result, prelude::*},
    net::{TcpListener, TcpStream},
    path::{Path, PathBuf},
    sync::{
        Arc, Mutex,
        atomic::{AtomicUsize, Ordering},
    },
};

use data_stream::{
    FromStream, ToStream, collections::SizeSettings, default_settings::PortableSettings,
    from_stream, to_stream,
};
use maud::html;
use percent_encoding::percent_decode_str;
use pukram2html::{Settings, convert, convert_extended, convert_subheader};
use rayon::prelude::*;

mod request;
use request::Request;

fn main() {
    let Ok(path) = env::current_dir() else {
        eprintln!("Current directory does not exist!");
        return;
    };

    let mut args = env::args();
    args.next();
    let address = args.next().unwrap_or_else(|| "127.0.0.1:8080".to_string());
    let password = args.next();
    let partial_password = args.next();
    start_server(path, &address, password, partial_password);
}

fn start_server(
    path: PathBuf,
    address: &str,
    password: Option<String>,
    partial_password: Option<String>,
) {
    let listener = TcpListener::bind(address).expect("Invalid bind address!");
    eprintln!("Strated server!");

    let context: Arc<Mutex<Context>> = Arc::default();

    listener.incoming().par_bridge().for_each(|stream| {
        eprintln!("New connection!");

        let Ok(stream) = stream else {
            eprintln!("Connection failed!");
            return;
        };

        let context = context.clone();
        let path = path.clone();
        let password = password.clone();
        let hidden_password = partial_password.clone();

        handle_connection(
            context,
            path,
            stream,
            password.as_deref(),
            hidden_password.as_deref(),
        )
    })
}

fn fail(mut stream: TcpStream) {
    let _ = write!(stream, "HTTP/1.1 404 Not Found\r\n\r\n");
    let _ = writeln!(stream, "Page not found!");
}

#[derive(Clone)]
struct Comment {
    name: Box<str>,
    text: Box<str>,
}

#[derive(Default)]
struct SiteInfo {
    comments: Mutex<Vec<Comment>>,
    visits: AtomicUsize,
    up: AtomicUsize,
    down: AtomicUsize,
}

impl<S: SizeSettings> ToStream<S> for SiteInfo {
    fn to_stream<W: Write>(&self, stream: &mut W) -> Result<()> {
        let Ok(comments) = self.comments.lock() else {
            return Err(std::io::ErrorKind::ResourceBusy.into());
        };
        S::size_to_stream(comments.len(), stream)?;
        for Comment { name, text } in comments.iter() {
            let name_bytes = name.as_bytes();
            S::size_to_stream(name_bytes.len(), stream)?;
            stream.write_all(name_bytes)?;

            let text_bytes = text.as_bytes();
            S::size_to_stream(text_bytes.len(), stream)?;
            stream.write_all(text_bytes)?;
        }

        S::size_to_stream(self.visits.load(Ordering::Acquire), stream)?;
        S::size_to_stream(self.up.load(Ordering::Acquire), stream)?;
        S::size_to_stream(self.down.load(Ordering::Acquire), stream)?;

        Ok(())
    }
}

impl<S: SizeSettings> FromStream<S> for SiteInfo {
    fn from_stream<R: Read>(stream: &mut R) -> Result<Self> {
        let size = S::size_from_stream(stream)?;
        let comments = Mutex::new(
            (0..size)
                .map(|_| {
                    let name_bytes = <Vec<_> as FromStream<S>>::from_stream(stream)?;
                    let name = std::str::from_utf8(&name_bytes)
                        .map_err(|e| Error::new(ErrorKind::InvalidData, e))?
                        .into();

                    let text_bytes = <Vec<_> as FromStream<S>>::from_stream(stream)?;
                    let text = std::str::from_utf8(&text_bytes)
                        .map_err(|e| Error::new(ErrorKind::InvalidData, e))?
                        .into();

                    Ok(Comment { name, text })
                })
                .collect::<Result<Vec<_>>>()?,
        );

        let visits = S::size_from_stream(stream)?.into();
        let up = S::size_from_stream(stream)?.into();
        let down = S::size_from_stream(stream)?.into();

        Ok(Self {
            comments,
            visits,
            up,
            down,
        })
    }
}

#[derive(Default)]
struct Context {
    infos: HashMap<Box<str>, Arc<SiteInfo>>,
}

fn handle_connection(
    context: Arc<Mutex<Context>>,
    path: PathBuf,
    mut stream: TcpStream,
    password: Option<&str>,
    partial_password: Option<&str>,
) {
    let Some(request) = Request::from(&stream) else {
        eprintln!("Invalid request!");
        return;
    };

    eprintln!();
    eprintln!("Request:");
    eprintln!("- Method: {}", request.method);
    eprintln!("- Path: {}", request.path);
    eprintln!("- Version: {}", request.version);
    eprintln!("- Headers:");
    for header in &request.headers {
        eprintln!("  - {header}");
    }
    eprintln!("- Body: {}", request.body);
    eprintln!();

    #[derive(PartialEq, Eq)]
    enum Access {
        None,
        Partial,
        Full,
    }

    let mut access = Access::None;

    let mut cookie = None;

    if let Some(password) = password {
        for entry in request.body.split('&') {
            let Some((key, input)) = entry.split_once('=') else {
                continue;
            };

            if key == "password" {
                if input == password {
                    access = Access::Full;
                    cookie = Some(password);
                } else if Some(input) == partial_password {
                    access = Access::Partial;
                    cookie = partial_password;
                }
                break;
            }
        }

        if access == Access::None {
            for header in request.headers {
                let Some((key, values)) = header.split_once(':') else {
                    continue;
                };
                let key = key.trim();

                if key.to_lowercase() != "cookie" {
                    continue;
                }

                for cookie in values.split(',') {
                    let cookie = cookie.trim();
                    let Some((name, state)) = cookie.split_once('=') else {
                        continue;
                    };
                    if name != "password" {
                        continue;
                    }

                    if state == password {
                        access = Access::Full;
                        break;
                    } else if Some(state) == partial_password {
                        access = Access::Partial;
                        break;
                    }
                }
            }
        }

        if access == Access::None {
            let _ = write!(stream, "HTTP/1.1 200 OK\r\n");
            let _ = write!(stream, "Content-Type: text/html; charset=\"utf-8\"\r\n");
            let _ = write!(stream, "\r\n");

            let html = html! {
                h1 { "Authentification required" }
                form method="POST" {
                    input type="password" name="password" placeholder="Password";
                    br;
                    input type="submit" value="Authenticate!";
                }
            };

            let _ = stream.write_all(html.into_string().as_bytes());

            return;
        }
    } else {
        access = Access::Full;
    }

    let (mut relative_path, _) = request
        .path
        .split_once('?')
        .unwrap_or_else(|| (&request.path, Default::default()));

    if relative_path.contains("//") {
        fail(stream);
        return;
    }

    if let Some(path) = relative_path.strip_prefix('/') {
        relative_path = path;
    }

    let mut pk_path = path.clone();
    let mut data_path = path.clone();

    let (pki_path, audio_path, start_level, relative_path) = if relative_path.is_empty() {
        if access == Access::Full {
            pk_path.push("index.pk");
            data_path.push("index.dat");
        } else {
            pk_path.push("partial.pk");
            data_path.push("partial.dat");
        }

        (None, None, 0, String::new())
    } else {
        let mut pki_path = path.clone();
        let mut audio_path = path.clone();

        let path = percent_decode_str(relative_path).decode_utf8_lossy();
        if path.contains('_') {
            let path = path.replace('_', " ");
            let _ = write!(stream, "HTTP/1.1 308 Permanent Redirect\r\n");
            let _ = write!(stream, "Location: /{path}\r\n\r\n");
            return;
        }

        pk_path.push(format!("{path}.pk"));
        if access == Access::Full {
            pki_path.push(format!("{path}.pki"));
        } else {
            pki_path.push(format!("{path}.pkc"));
        }
        audio_path.push(format!("{path}.mp3"));
        data_path.push(format!("{path}.dat"));

        (Some(pki_path), Some(audio_path), 1, path.to_string())
    };

    let file_paths = DocumentPaths {
        pk: &pk_path,
        pki: pki_path.as_ref().map(PathBuf::as_ref),
        audio: audio_path.as_ref().map(PathBuf::as_ref),
        data: &data_path,
    };

    if let Some((_, ending)) = relative_path.rsplit_once('.') {
        let ending = ending.to_lowercase();
        match ending.as_ref() {
            "jpg" => reply_binary(stream, &relative_path, "image", "jpeg", path),
            "png" | "tiff" | "gif" | "jpeg" => {
                reply_binary(stream, &relative_path, "image", &ending, path)
            }
            "mp3" | "wav" => reply_binary(stream, &relative_path, "audio", &ending, path),
            _ => fail(stream),
        }
        return;
    }

    if !Path::is_file(&pk_path) {
        fail(stream);
        return;
    }

    if let Some(pki_path) = &pki_path {
        if !Path::is_file(pki_path) {
            fail(stream);
            return;
        }
    }

    let info = if let Ok(mut context) = context.lock() {
        use Entry::*;
        match context.infos.entry(relative_path.clone().into_boxed_str()) {
            Occupied(o) => o.get().clone(),
            Vacant(v) => v
                .insert(Arc::new(
                    File::open(&data_path)
                        .map(|mut file| {
                            from_stream::<PortableSettings, _, _>(&mut file).unwrap_or_default()
                        })
                        .unwrap_or_default(),
                ))
                .clone(),
        }
    } else {
        return;
    };

    handle_relative_connection(
        info,
        stream,
        &request.body,
        &relative_path,
        &path,
        file_paths,
        start_level,
        cookie,
        access != Access::Full,
    )
}

fn reply_binary(
    mut stream: TcpStream,
    relative_path: &str,
    data_type: &str,
    file_type: &str,
    mut path: PathBuf,
) {
    path.push(relative_path);
    let Ok(mut file) = File::open(path) else {
        fail(stream);
        return;
    };

    let _ = write!(stream, "HTTP/1.1 200 OK\r\n");
    let _ = write!(stream, "Content-Type: {data_type}/{file_type}\r\n");
    let _ = write!(stream, "\r\n");

    let mut buf = [0; 0x400];

    while let Ok(len) = file.read(&mut buf) {
        if len == 0 {
            break;
        }
        let _ = stream.write_all(&buf[0..len]);
    }
}

#[derive(Copy, Clone)]
struct DocumentPaths<'a> {
    pk: &'a Path,
    pki: Option<&'a Path>,
    audio: Option<&'a Path>,
    data: &'a Path,
}

fn handle_relative_connection(
    info: Arc<SiteInfo>,
    mut stream: TcpStream,
    body: &str,
    relative_path: &str,
    path: &Path,
    file_paths: DocumentPaths,
    start_level: usize,
    cookie: Option<&str>,
    censored: bool,
) {
    let mut name = None;
    let mut text = None;

    let mut up = false;
    let mut down = false;

    for entry in body.split('&') {
        let Some((key, input)) = entry.split_once('=') else {
            continue;
        };

        match key {
            "name" => {
                let name_value = input.replace('+', " ");
                let decoded_name = percent_decode_str(&name_value).decode_utf8_lossy();
                eprintln!("Received comment by \"{decoded_name}\":");
                name = Some(
                    decoded_name
                        .replace('&', "&amp;")
                        .replace('<', "&lt;")
                        .replace('>', "&gt;")
                        .into(),
                );
            }
            "text" => {
                let text_value = input.replace('+', " ");
                let decoded_text = percent_decode_str(&text_value).decode_utf8_lossy();
                for line in decoded_text.lines() {
                    eprintln!("  {line}");
                }
                eprintln!();
                let mut text_buf = Vec::new();
                convert(decoded_text.lines(), &mut text_buf);
                text = Some(
                    std::str::from_utf8(text_buf.as_slice())
                        .unwrap_or_default()
                        .into(),
                );
            }
            "up" => up = true,
            "down" => down = true,
            _ => (),
        }
    }

    let comments = {
        let Ok(mut comments) = info.comments.lock() else {
            return;
        };
        if let (Some(name), Some(text)) = (name, text) {
            comments.push(Comment { name, text });
        }
        comments.clone()
    };
    let up = if up {
        info.up.fetch_add(1, Ordering::Relaxed)
    } else {
        info.up.load(Ordering::Relaxed)
    };
    let down = if down {
        info.down.fetch_add(1, Ordering::Relaxed)
    } else {
        info.down.load(Ordering::Relaxed)
    };

    let visits = info.visits.fetch_add(1, Ordering::Relaxed);

    if let Ok(mut file) = File::create(file_paths.data) {
        if to_stream::<PortableSettings, _, _>(&*info, &mut file).is_err() {
            eprintln!("Error saving data!");
            eprintln!();
        }
    }

    let _ = write!(stream, "HTTP/1.1 200 OK\r\n");
    let _ = write!(stream, "Content-Type: text/html; charset=\"utf-8\"\r\n");
    if let Some(password) = cookie {
        let _ = write!(stream, "Set-Cookie: password={password}\r\n");
    }

    let _ = write!(stream, "\r\n");

    let parent_path = relative_path
        .rsplit_once('/')
        .map(|(path, _)| path)
        .unwrap_or_default();

    let section = |stream: &mut TcpStream| {
        let _ = writeln!(stream, "<hr>");
        let _ = writeln!(stream, "<a href=\"/{parent_path}\">&lt;&lt; Back</a>");
    };

    let _ = writeln!(stream, "<p>👁️{visits} 💖️{up} 💔️{down}</p>");

    section(&mut stream);

    let title = relative_path
        .rsplit_once('/')
        .map_or(relative_path, |(_, title)| title);

    fn entry_handler<W: Write>(
        path: &Path,
        relative_path: &str,
        censored: bool,
    ) -> impl Fn(&str, &mut W, usize) {
        move |mut entry, output, level| {
            let level = level + 1;
            let mut pki_path = path.to_path_buf();
            let mut audio_path = path.to_path_buf();
            if let Some((real_entry, _)) = entry.split_once(':') {
                entry = real_entry
            }
            let pki_extension = if censored { "pkc" } else { "pki" };
            if relative_path.is_empty() {
                pki_path.push(format!("{entry}.{pki_extension}"));
                audio_path.push(format!("{entry}.mp3"));
            } else {
                pki_path.push(format!("{relative_path}/{entry}.{pki_extension}"));
                audio_path.push(format!("{relative_path}/{entry}.mp3"));
            }

            let Ok(file) = File::open(pki_path) else {
                return;
            };

            let _ = writeln!(
                output,
                "<h{level}><a href=\"{relative_path}/{entry}\">{entry}</a></h{level}>"
            );

            if Path::is_file(&audio_path) {
                let _ = writeln!(
                    output,
                    "<p><audio controls src=\"/{relative_path}/{entry}.mp3\"/></p>"
                );
            }

            convert_subheader(
                BufReader::new(file).lines().map(Result::unwrap_or_default),
                output,
                level,
            );
        }
    }

    let Ok(pk_file) = File::open(file_paths.pk) else {
        unreachable!();
    };

    let lines = BufReader::new(pk_file).lines();

    if !title.is_empty() {
        let _ = writeln!(stream, "<h1>{title}</h1>");
    }

    let check_path: &Path = relative_path.as_ref();
    let interactive = check_path.parent().is_some_and(|parent| {
        std::fs::metadata(parent.with_extension("vng")).is_ok_and(|file| file.is_file())
    });

    if interactive {
        let mut sections = Vec::new();
        let mut current_section = Vec::new();

        let mut add_section = |current_section: Vec<String>| {
            if current_section.is_empty() {
                return;
            }

            let mut section_html = Vec::new();
            convert_extended(
                current_section,
                &mut section_html,
                Settings::default()
                    .with_handler(entry_handler(path, relative_path, censored))
                    .with_start_level(start_level)
                    .with_use_textboxes(true),
            );
            sections.push(String::from_utf8_lossy(&section_html).into_owned());
        };

        for line in lines {
            let line = line.unwrap_or_default();
            if !line.is_empty() {
                current_section.push(line);
                continue;
            }

            add_section(std::mem::take(&mut current_section));
        }

        add_section(current_section);

        let html = html! {
            div id="story-container" {
                div {
                    span id="section-counter" { "1/" (sections.len()) }
                }

                div {
                    button onclick="prev()" { "Prev" }
                    button onclick="next()" { "Next" }
                }

                @for (index, section) in sections.iter().enumerate() {
                    div class="story-section" data-section-index=(index)
                        style=(format!("display: {};", if index == 0 { "block" } else { "none" })) {
                        (maud::PreEscaped(section))
                    }
                }

                script {
                    (maud::PreEscaped(r"
                        let currentSection = 0;
                        const totalSections = ".to_owned() + &sections.len().to_string() + r";

                        function updateSection() {
                            document.querySelectorAll('.story-section').forEach((el, index) => {
                                el.style.display = index === currentSection ? 'block' : 'none';
                            });
                            document.getElementById('section-counter').textContent = 
                                `${currentSection + 1}/${totalSections}`;
                        }

                        function prev() {
                            if (currentSection > 0) {
                                --currentSection;
                                updateSection();
                            }
                        }

                        function next() {
                            if (currentSection < totalSections - 1) {
                                ++currentSection;
                                updateSection();
                            }
                        }
                    "))
                }
            }
        };

        let _ = write!(stream, "{}", html.into_string());
    } else {
        convert_extended(
            lines.map(Result::unwrap_or_default),
            &mut stream,
            Settings::default()
                .with_handler(entry_handler(path, relative_path, censored))
                .with_start_level(start_level)
                .with_use_textboxes(true),
        );
    }

    section(&mut stream);

    if let Some(pki_path) = file_paths.pki {
        let Ok(pki_file) = File::open(pki_path) else {
            unreachable!();
        };

        let _ = writeln!(stream, "<h1>Description</h1>");

        if let Some(audio_path) = &file_paths.audio {
            if Path::is_file(audio_path) {
                let _ = writeln!(
                    stream,
                    "<p><audio controls src=\"/{relative_path}.mp3\"/></p>"
                );
            }
        }

        let lines = BufReader::new(pki_file).lines();
        convert_subheader(lines.map(Result::unwrap_or_default), &mut stream, 1);

        section(&mut stream);
    }

    let html = html! {
        h1 { "Comments" }
        form method="POST" {
            input type="text" name="name" value="anon" placeholder="Name";
            br;
            textarea rows="5" cols="60" name="text" placeholder="Enter comment..." {}
            br;
            input type="submit" value="Send!";
        }
        form method="POST" {
            input type="submit" value="💖️" name="up";
            input type="submit" value="💔️" name="down";
        }
    };
    let _ = stream.write_all(html.into_string().as_bytes());

    for Comment { name, text } in comments {
        let html = html! {
            fieldset {
                legend { (name) }
                (maud::PreEscaped(text))
            }
        };
        let _ = stream.write_all(html.into_string().as_bytes());
    }

    section(&mut stream);
}