diff --git a/core/src/db.rs b/core/src/db.rs index ee979c5..0dde04f 100644 --- a/core/src/db.rs +++ b/core/src/db.rs @@ -281,6 +281,20 @@ impl UsefulFunctions for sqlite::SqlitePooledConnection { } impl DbConn { + /// Tries to fetch the user with the given id. + pub fn fetch_user_id(&self, id: i32) -> Result { + use crate::schema::users::dsl; + let query = dsl::users.filter(dsl::id.eq(id)); + match self { + #[cfg(feature = "mysql")] + DbConn::Mysql(conn) => query.first(conn), + #[cfg(feature = "postgres")] + DbConn::Postgres(conn) => query.first(conn), + #[cfg(feature = "sqlite")] + DbConn::Sqlite(conn) => query.first(conn), + } + .map_err(Error::from) + } /// Tries to fetch the user with the given email. (TODO: lookup by username) pub fn fetch_user(&self, email: impl AsRef) -> Result { use crate::schema::users::dsl; diff --git a/frontend/src/base.rs b/frontend/src/base.rs index 3482b35..247345b 100644 --- a/frontend/src/base.rs +++ b/frontend/src/base.rs @@ -1,8 +1,11 @@ -use tera::Context; use warp::Filter; +use crate::extractors::{get_context, navbar, Context}; use crate::render::render_template; pub fn get_index() -> Resp!() { - warp::path::end().and_then(|| render_template("base/index.html", Context::new())) + warp::path::end() + .and(navbar()) + .and(get_context()) + .and_then(|ctx: Context| render_template("base/index.html", ctx.into())) } diff --git a/frontend/src/extractors.rs b/frontend/src/extractors.rs index 499ff85..1938c5a 100644 --- a/frontend/src/extractors.rs +++ b/frontend/src/extractors.rs @@ -1,20 +1,60 @@ -use core::State; +use core::{Error, State}; +use std::ops::{Deref, DerefMut}; +use tera::Context as TeraContext; use warp::{Filter, Rejection}; use crate::session::Session; -#[derive(Default)] -pub struct Context(::tera::Context); +#[derive(Clone, Default)] +pub struct Context(pub TeraContext); -fn get_context() -> impl Clone + Filter { - warp::ext::get::() - .and(warp::ext::get::()) - .map(|state, session| Context::default()) +impl Deref for Context { + type Target = TeraContext; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for Context { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl Into for Context { + fn into(self) -> TeraContext { + self.0 + } +} + +pub fn get_context() -> impl Clone + Filter { + warp::ext::get::() .or(warp::any().map(|| Context::default())) .unify() } +fn render_navbar(state: &State, session: &Session, ctx: &mut Context) -> Result<(), Error> { + let conn = state.get_connection()?; + + // retrieve the user data + if let Some(user_id) = session.user_id { + conn.fetch_user_id(user_id) + .map(|user| ctx.insert("user", &user))?; + } + + Ok(()) +} + /// Generates all of the information needed to populate the navbar. pub fn navbar() -> impl Clone + Filter { - get_context().map(|ctx| {}).untuple_one() + warp::ext::get::() + .and(warp::ext::get::()) + .and(get_context()) + .map(|state: State, session: Session, mut ctx: Context| { + render_navbar(&state, &session, &mut ctx).unwrap_or_else(|err| { + println!("err: {:?}", err); + }); + warp::ext::set::(ctx); + }) + .untuple_one() } diff --git a/frontend/src/session.rs b/frontend/src/session.rs index 003efc7..2ab544d 100644 --- a/frontend/src/session.rs +++ b/frontend/src/session.rs @@ -6,6 +6,7 @@ use warp::{Filter, Rejection}; #[derive(Debug, Copy, Clone, Default, Serialize, Deserialize)] pub struct Session { pub user_id: Option, + pub team_id: Option, } pub fn extract() -> impl Clone + Filter { @@ -48,8 +49,9 @@ pub fn apply() -> impl Clone + Filter,), Error = Rejec .map(|data| { jar.private(&key).add( Cookie::build("session", data) - .secure(true) + // .secure(true) // TODO: enable this based on a config .http_only(true) + .path("/") .same_site(SameSite::Strict) .finish(), ); diff --git a/frontend/src/users.rs b/frontend/src/users.rs index cb70548..4831fcb 100644 --- a/frontend/src/users.rs +++ b/frontend/src/users.rs @@ -4,16 +4,18 @@ use core::{ Error, State, }; use http::uri::Uri; -use tera::Context; use warp::Filter; use wtforms::Form; -use crate::extractors::navbar; +use crate::extractors::{get_context, navbar, Context}; use crate::render::render_template; use crate::session::Session; pub fn get_login() -> Resp!() { - navbar().and_then(|| render_template("users/login.html", Context::new())) + warp::path::end() + .and(navbar()) + .and(get_context()) + .and_then(|ctx: Context| render_template("users/login.html", ctx.into())) } pub fn post_login() -> Resp!() { @@ -39,11 +41,17 @@ pub fn post_login() -> Resp!() { } pub fn get_profile() -> Resp!() { - navbar().and_then(|| render_template("users/profile.html", Context::new())) + warp::path::end() + .and(navbar()) + .and(get_context()) + .and_then(|ctx: Context| render_template("users/profile.html", ctx.into())) } pub fn get_register() -> Resp!() { - navbar().and_then(|| render_template("users/register.html", Context::new())) + warp::path::end() + .and(navbar()) + .and(get_context()) + .and_then(|ctx: Context| render_template("users/register.html", ctx.into())) } pub fn post_register() -> Resp!() { diff --git a/frontend/templates/layout.html b/frontend/templates/layout.html index 4e3c47c..0164812 100644 --- a/frontend/templates/layout.html +++ b/frontend/templates/layout.html @@ -9,7 +9,7 @@ {% block title %}{% endblock %} - LibreCTF -