view src/main.rs @ 23:1f33922b29f2

Use FileServer to serve static assets
author Lewin Bormann <lbo@spheniscida.de>
date Tue, 12 Jul 2022 12:02:47 -0700
parents 64eae1262165
children 758a4f6c160f
line wrap: on
line source

use anyhow::{self, Context, Error};
use either::Either;
use log::{debug, error, info, warn, Level};
use time::{Duration, OffsetDateTime};

use rocket::form::Form;
use rocket::fs::{relative, FileServer, NamedFile};
use rocket::futures::StreamExt;
use rocket::http::{Cookie, CookieJar, HeaderMap, Status};
use rocket::request::{self, FlashMessage, FromRequest, Outcome, Request};
use rocket::response::{self, Flash, Redirect, Responder};
use rocket::serde::json::{json, Value};
use rocket::serde::{self, Serialize};

use rocket_db_pools::sqlx::{self, Executor, Row, Sqlite, SqlitePool};
use rocket_db_pools::{Connection, Database, Pool};

use rocket_dyn_templates::{context, Template};

use std::collections::HashMap;
use std::net::IpAddr;
use std::path::Path;
use std::time::Instant;

#[cfg(feature = "sqlite")]
type DBType = Sqlite;
#[cfg(feature = "sqlite")]
type PoolType = SqlitePool;

// Current SQL queries don't work with postgres.
#[cfg(feature = "postgres")]
use rocket_db_pools::sqlx::{PgPool, Postgres};
#[cfg(feature = "postgres")]
type DBType = Postgres;
#[cfg(feature = "postgres")]
type PoolType = PgPool;

// TO DO: use other databases?
#[derive(Database)]
#[database("sqlite_main")]
struct ConfigDB(PoolType);

struct ConfigDBSession<'p, DB: sqlx::Database>(&'p mut sqlx::pool::PoolConnection<DB>);

impl<'p> ConfigDBSession<'p, Sqlite> {
    async fn check_user_password<S: AsRef<str>>(
        &mut self,
        user: S,
        password: S,
    ) -> Result<bool, Error> {
        // TODO: salt passwords.
        let salt: String = match sqlx::query("SELECT salt FROM users WHERE username = ? LIMIT 1;")
            .bind(user.as_ref())
            .fetch_one(&mut *self.0)
            .await
        {
            Ok(r) => r.get(0),
            Err(e) => {
                warn!("Error querying salt: {}", e);
                return Ok(false);
            }
        };
        let pwdhash = sha256::digest(format!("{}{}", salt, password.as_ref()));
        let q = sqlx::query("SELECT username FROM users WHERE username = ? AND password_hash = ?;")
            .bind(user.as_ref())
            .bind(pwdhash);
        let result = self.0.fetch_all(q).await?;
        Ok(result.len() == 1)
    }
}

#[derive(Database)]
#[database("sqlite_logs")]
struct LogsDB(PoolType);

struct LogsDBSession<'p, DB: sqlx::Database>(&'p mut sqlx::pool::PoolConnection<DB>);

impl<'p> LogsDBSession<'p, Sqlite> {
    async fn log_request<
        S1: AsRef<str>,
        S2: AsRef<str>,
        S3: AsRef<str>,
        S4: AsRef<str>,
        S5: AsRef<str>,
        S6: AsRef<str>,
    >(
        &mut self,
        session: Option<u32>,
        ip: S1,
        domain: S2,
        path: S3,
        status: u32,
        page: Option<S4>,
        refer: Option<S5>,
        ua: S6,
        ntags: u32,
    ) -> Result<u32, Error> {
        let q = sqlx::query::<DBType>("INSERT INTO RequestLog (session, ip, atime, domain, path, status, pagename, refer, ua, ntags) VALUES (?, ?, strftime('%s', 'now'), ?, ?, ?, ?, ?, ?, ?) RETURNING id");
        let q = q
            .bind(session)
            .bind(ip.as_ref())
            .bind(domain.as_ref())
            .bind(path.as_ref())
            .bind(status)
            .bind(page.map(|s| s.as_ref().to_string()))
            .bind(refer.map(|s| s.as_ref().to_string()))
            .bind(ua.as_ref())
            .bind(ntags);
        let row: u32 = q.fetch_one(&mut *self.0).await?.get(0);
        Ok(row)
    }

    async fn log_tags<S: AsRef<str>, I: Iterator<Item = S>>(
        &mut self,
        requestid: u32,
        tags: I,
    ) -> Result<usize, Error> {
        let mut ntags = 0;
        for tag in tags {
            let (k, v) = tag.as_ref().split_once("=").unwrap_or((tag.as_ref(), ""));
            sqlx::query("INSERT INTO RequestTags (requestid, key, value) VALUES (?, ?, ?)")
                .bind(requestid)
                .bind(k)
                .bind(v)
                .execute(&mut *self.0)
                .await
                .map_err(|e| error!("Couldn't insert tag {}={}: {}", k, v, e))
                .unwrap();
            ntags += 1;
        }

        Ok(ntags)
    }

    async fn start_session(&mut self, domain: &str) -> Result<u32, Error> {
        Ok(sqlx::query("INSERT INTO Sessions (start, last, domain) VALUES (strftime('%s', 'now'), strftime('%s', 'now'), ?) RETURNING id")
            .bind(domain)
            .fetch_one(&mut *self.0)
            .await?.get(0))
    }

    async fn update_session_time(&mut self, id: u32) -> Result<(), Error> {
        sqlx::query("UPDATE Sessions SET last = strftime('%s', 'now') WHERE id = ?")
            .bind(id)
            .execute(&mut *self.0)
            .await?;
        Ok(())
    }

    async fn query_visits_sessions_counts<S: AsRef<str>>(
        &mut self,
        from: OffsetDateTime,
        to: OffsetDateTime,
        domainpattern: Option<S>,
    ) -> Result<(Vec<String>, Vec<u32>, Vec<u32>), Error> {
        let domain = domainpattern.as_ref().map(AsRef::as_ref).unwrap_or("%");
        let mut results = sqlx::query(
            r#"
SELECT DATE(atime, 'unixepoch') AS rqdate, COUNT(requestlog.id) AS rqcount, sesscount
FROM requestlog
JOIN (
    SELECT DATE(start, 'unixepoch') AS sessdate, COUNT(*) AS sesscount
    FROM sessions WHERE sessions.domain LIKE ? GROUP BY sessdate)
AS sc ON (rqdate = sessdate)
WHERE atime > ? AND atime < ? AND requestlog.domain LIKE ?
GROUP BY rqdate
ORDER BY rqdate ASC;"#,
        )
        .bind(domain)
        .bind(from.unix_timestamp())
        .bind(to.unix_timestamp())
        .bind(domain)
        .fetch(&mut *self.0);

        // Result table: date / visit count / session count
        let mut dates = vec![]; //Vec::<String>::with_capacity(results.len());
        let mut visits = vec![]; //Vec::<u32>::with_capacity(results.len());
        let mut sessions = vec![]; //Vec::<u32>::with_capacity(results.len());

        loop {
            match results.next().await {
                Some(Ok(row)) => {
                    dates.push(row.get(0));
                    visits.push(row.get(1));
                    sessions.push(row.get(2));
                }
                None => break,
                Some(Err(e)) => {
                    error!("Error querying visits/sessions: {}", e);
                    return Err(e).context("query visits/session counts");
                }
            }
        }

        Ok((dates, visits, sessions))
    }
}

const USER_ID_COOKIE_KEY: &str = "user_id";

struct LoggedInGuard(String);

#[rocket::async_trait]
impl<'r> FromRequest<'r> for LoggedInGuard {
    type Error = Error;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        let cookies = req.cookies();
        if let Some(uc) = cookies.get_private(USER_ID_COOKIE_KEY) {
            Outcome::Success(LoggedInGuard(uc.value().to_string()))
        } else {
            Outcome::Forward(())
        }
    }
}

struct HeadersGuard<'h>(HeaderMap<'h>);

#[rocket::async_trait]
impl<'r> FromRequest<'r> for HeadersGuard<'r> {
    type Error = Error;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        Outcome::Success(HeadersGuard(req.headers().clone()))
    }
}

#[derive(Responder)]
enum LoginResponse {
    // use templates later.
    #[response(status = 200, content_type = "html")]
    Ok { body: Template },
    #[response(status = 302, content_type = "html")]
    LoggedInAlready { redirect: Redirect },
}

#[derive(Responder)]
#[response(content_type = "text", status = 500)]
struct InternalServerError {
    body: String,
}

#[rocket::get("/login")]
async fn route_login_form<'r>(
    cc: &rocket::State<CustomConfig>,
    flash: Option<FlashMessage<'_>>,
    cookies: &CookieJar<'_>,
) -> Result<LoginResponse, InternalServerError> {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }
    if let Some(cookie) = cookies.get_private(USER_ID_COOKIE_KEY) {
        Ok(LoginResponse::LoggedInAlready {
            redirect: Redirect::to(cc.deploy_path.clone()),
        })
    } else {
        Ok(LoginResponse::Ok {
            body: Template::render("login", context![flash: f]),
        })
    }
}

#[rocket::post("/logout")]
fn route_logout(cc: &rocket::State<CustomConfig>, cookies: &CookieJar<'_>) -> Flash<Redirect> {
    if let Some(cookie) = cookies.get_private(USER_ID_COOKIE_KEY) {
        cookies.remove_private(cookie);
    }
    Flash::success(
        Redirect::to(cc.deploy_path.clone()),
        format!("Logged out. Goodbye!"),
    )
}

#[derive(rocket::FromForm)]
struct LoginForm {
    username: String,
    password: String,
}

#[rocket::post("/login", data = "<login>")]
async fn route_login_post(
    mut db: Connection<ConfigDB>,
    cc: &rocket::State<CustomConfig>,
    cookies: &CookieJar<'_>,
    login: Form<LoginForm>,
) -> Flash<Redirect> {
    // TO DO: look up user in database.
    let mut conn = ConfigDBSession(&mut db);
    match conn
        .check_user_password(&login.username, &login.password)
        .await
    {
        Ok(true) => {
            let c = Cookie::new(USER_ID_COOKIE_KEY, login.username.clone());
            cookies.add_private(c);
            Flash::success(
                Redirect::to(cc.deploy_path.clone()),
                "Successfully logged in.",
            )
        }
        Ok(false) => Flash::error(
            Redirect::to(format!("{}/login", cc.deploy_path)),
            "User/password not found",
        ),
        Err(e) => Flash::error(
            Redirect::to(format!("{}/login", cc.deploy_path)),
            format!("User/password lookup failed: {}", e),
        ),
    }
}

fn create_chart<S1: Serialize, S2: Serialize>(
    labels: Vec<S1>,
    values: Vec<(String, Vec<S2>)>,
    typ: String,
) -> Value {
    let colors = vec![
        "red", "blue", "orange", "green", "gray", "purple", "black", "brown",
    ];
    if values.len() > colors.len() {
        error!("Not enough colors for line chart!");
    }
    let datasets: Vec<Value> = values
        .iter()
        .zip(colors)
        .map(|((name, val), col)| json!({"label": name, "data": val, "borderColor": col}))
        .collect();
    let inner = json!({
        "type": typ,
        "data": {
            "labels": labels,
            "datasets": datasets,
        },
        "options": {
            "scales": { "y": { "beginAtZero": true }},
        },
    });
    inner
}

#[rocket::get("/?<domain>", rank = 1)]
async fn route_index_loggedin(
    mut conn: Connection<LogsDB>,
    lig: LoggedInGuard,
    flash: Option<FlashMessage<'_>>,
    domain: Option<&str>,
) -> Template {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }

    let mut charts = HashMap::<String, String>::new();

    // TODO: Make configurable
    let begin = OffsetDateTime::now_utc() - Duration::days(30);
    let end = OffsetDateTime::now_utc();
    match LogsDBSession(&mut conn)
        .query_visits_sessions_counts(begin, end, domain)
        .await
    {
        Ok((dates, visits, sessions)) => {
            info!(
                "Successfully queried visits/sessions: {:?}, {:?}",
                visits, sessions
            );
            let vissess = create_chart(
                dates,
                vec![("Visits".into(), visits), ("Sessions".into(), sessions)],
                "line".into(),
            )
            .to_string();
            Template::render(
                "index",
                context![
                loggedin: true,
                domain: domain,
                username: lig.0,
                flash: f,
                chartconfig: context![ visitsAndSessions: vissess ]],
            )
        }
        Err(e) => Template::render(
            "index",
            context![loggedin: true, domain: domain, username: lig.0, flash: f, chartconfig: context![], error: format!("{:?}", e)],
        ),
    }
}

#[rocket::get("/", rank = 2)]
async fn route_index_loggedout(flash: Option<FlashMessage<'_>>) -> Template {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }
    Template::render("index", context![loggedin: false, flash: f])
}

// TODO: ignore requests when logged in.
#[rocket::get("/log?<host>&<status>&<path>&<pagename>&<referer>&<tags>")]
async fn route_log(
    mut conn: Connection<LogsDB>,
    cookies: &CookieJar<'_>,
    host: Option<String>,
    status: Option<u32>,
    path: Option<String>,
    pagename: Option<String>,
    referer: Option<String>,
    tags: Vec<String>,
    config: &rocket::State<CustomConfig>,
    ip: IpAddr,
    headers: HeadersGuard<'_>,
) -> (Status, Either<NamedFile, &'static str>) {
    let mut conn = LogsDBSession(&mut conn);
    let mut session_id = None;
    // Only log if this is not an analytics user.
    if cookies.get_private(USER_ID_COOKIE_KEY).is_none() {
        // Get session ID from cookie, or start new session.
        if let Some(sessioncookie) = cookies.get_private("analyrics_session") {
            if let Ok(id) = u32::from_str_radix(sessioncookie.value(), 10) {
                session_id = Some(id);
                match conn.update_session_time(id).await {
                    Ok(()) => {}
                    Err(e) => error!("Couldn't update session time: {}", e),
                }
            }
        }
        if session_id.is_none() {
            match conn
                .start_session(host.as_ref().map(String::as_str).unwrap_or(""))
                .await
            {
                Ok(id) => {
                    session_id = Some(id);
                    let c = Cookie::build("analyrics_session", id.to_string())
                        .max_age(time::Duration::hours(12))
                        .finish();
                    cookies.add_private(c);
                }
                Err(e) => error!("Couldn't start session: {}", e),
            }
        }

        let ntags = tags.len() as u32;
        let ua = headers.0.get_one("user-agent").unwrap_or("");
        let ip = ip.to_string();
        let ip = headers.0.get_one("x-real-ip").unwrap_or(&ip);
        let host: String = host.unwrap_or(
            headers
                .0
                .get("host")
                .take(1)
                .map(|s| s.to_string())
                .collect::<Vec<String>>()
                .pop()
                .unwrap_or(String::new()),
        );
        match conn
            .log_request(
                session_id,
                ip,
                host,
                path.unwrap_or(String::new()),
                status.unwrap_or(200),
                pagename,
                referer,
                ua,
                ntags,
            )
            .await
        {
            Err(e) => error!("Couldn't log request: {}", e),
            Ok(id) => {
                conn.log_tags(id, tags.iter()).await.ok();
            }
        }
    }

    if let Ok(f) = NamedFile::open(Path::new(config.asset_path.as_str()).join("pixel.png")).await {
        (Status::Ok, Either::Left(f))
    } else {
        (Status::Ok, Either::Right(""))
    }
}

#[rocket::get("/testchart")]
async fn route_testchart() -> Template {
    Template::render(
        "testchart",
        context![ chartconfig: create_chart(vec!["1", "2", "3", "4"], vec![("Series A".into(), vec![1,4,3,8]), ("Series B".into(), vec![5,2,3,9])], "line".into()).to_string() ],
    )
}

#[rocket::get("/static/<path>")]
async fn route_static(
    cc: &rocket::State<CustomConfig>,
    path: &str,
) -> Either<NamedFile, (Status, String)> {
    match NamedFile::open(Path::new(&cc.asset_path).join(path)).await {
        Ok(f) => Either::Left(f),
        Err(e) => {
            warn!("Static file not found: {}", path);
            Either::Right((
                Status::NotFound,
                format!("Error loading file at {}: {}", path, e),
            ))
        }
    }
}
#[derive(rocket::serde::Deserialize)]
#[serde(crate = "rocket::serde")]
struct CustomConfig {
    asset_path: String,
    deploy_path: String,
}

#[rocket::launch]
fn rocketmain() -> _ {
    env_logger::init();

    let r = rocket::build()
        .attach(ConfigDB::init())
        .attach(LogsDB::init())
        .attach(Template::fairing())
        .attach(rocket::fairing::AdHoc::config::<CustomConfig>())
        .mount("/static", FileServer::from(relative!("assets/static")))
        .mount(
            "/",
            rocket::routes![
                route_index_loggedin,
                route_index_loggedout,
                route_logout,
                route_login_form,
                route_login_post,
                route_log,
                route_testchart,
            ],
        );
    r
}