uidc/src/server/session.rs

use crate::{schema, user::UserExt, UIDCError};
use microrm::{prelude::*, schema::Stored};
use serde::Deserialize;
use tide::http::Cookie;

pub(super) struct SessionHelper<'l> {
    db: &'l schema::UIDCDatabase,
    tmpl: &'l handlebars::Handlebars<'l>,
    realm_str: &'l str,
}

type Request = tide::Request<super::ServerStateWrapper>;

const SESSION_COOKIE_NAME: &str = "uidc_session";

impl<'l> SessionHelper<'l> {
    pub fn new(req: &'l Request) -> Self {
        Self {
            db: &req.state().core.db,
            tmpl: &req.state().core.templates,
            realm_str: req.param("realm").expect("no realm param?"),
        }
    }

    pub fn get_realm(&self) -> tide::Result<Stored<schema::Realm>> {
        self.db
            .realms
            .keyed(self.realm_str)
            .get()?
            .ok_or(tide::Error::from_str(404, "No such realm"))
    }

    fn build_session(
        &self,
    ) -> tide::Result<(schema::Session, Option<tide::http::Cookie<'static>>)> {
        let rng = ring::rand::SystemRandom::new();
        let session_id: [u8; 32] = ring::rand::generate(&rng)
            .map_err(|_| tide::Error::from_str(500, "Failed to generate session ID"))?
            .expose();
        let session_id = base64::encode_config(session_id, base64::URL_SAFE_NO_PAD);

        // XXX: replace with in-place insertion once support for that is added to microrm
        let session = self.db.sessions.insert_and_return(schema::Session {
            session_id: session_id.clone(),
            auth: Default::default(),
            expiry: time::OffsetDateTime::now_utc() + time::Duration::minutes(10),
        })?;
        let session_cookie = Cookie::build(SESSION_COOKIE_NAME, session_id)
            .path("/")
            .finish();
        Ok((session.wrapped(), Some(session_cookie)))
    }

    pub fn verify_session(&self, req: &Request) -> Option<(Stored<schema::Realm>, schema::UserID)> {
        self.get_or_build_session(req)
            .ok()
            .zip(self.get_realm().ok())
            .and_then(|((sid, _cookie), realm)| {
                self.get_auth_for_session(realm.id(), &sid)
                    .and_then(|auth| auth.user.map(|user| (realm, user)))
            })
    }

    pub fn get_session(&self, req: &Request) -> Option<schema::Session> {
        req.cookie(SESSION_COOKIE_NAME).and_then(|sid| {
            self.db
                .sessions
                .keyed(sid.value())
                .get()
                .ok()
                .flatten()
                .map(|v| v.wrapped())
        })
    }

    pub fn get_or_build_session(
        &self,
        req: &Request,
    ) -> tide::Result<(schema::Session, Option<tide::http::Cookie<'static>>)> {
        match self.get_session(req) {
            Some(s) => Ok((s, None)),
            None => self.build_session(),
        }
    }

    pub fn get_auth_for_session(
        &self,
        realm: schema::RealmID,
        session: &schema::Session,
    ) -> Option<Stored<schema::SessionAuth>> {
        session
            .auth
            .with(schema::SessionAuth::Realm, realm)
            .first()
            .get()
            .ok()?
    }

    pub fn destroy_auth(
        &self,
        realm: schema::RealmID,
        session: &schema::Session,
    ) -> Result<(), UIDCError> {
        session
            .auth
            .with(schema::SessionAuth::Realm, realm)
            .first()
            .delete()?;
        Ok(())
    }
}

impl<'l> SessionHelper<'l> {
    fn render_login_from_auth(
        &self,
        mut response: tide::Response,
        redirect: String,
        auth: Option<schema::SessionAuth>,
        error_msg: Option<String>,
    ) -> tide::Response {
        let to_present: Option<schema::AuthChallengeType> = match auth {
            None => Some(schema::AuthChallengeType::Username),
            Some(auth) => auth.pending_challenges.as_ref().first().copied(),
        };

        if let Some(to_present) = to_present {
            self.render_login_page(response, redirect, to_present, error_msg)
        } else {
            response.set_status(302);
            tide::Redirect::new(redirect).into()
        }
    }

    fn render_login_page(
        &self,
        mut response: tide::Response,
        redirect: String,
        to_present: schema::AuthChallengeType,
        error_msg: Option<String>,
    ) -> tide::Response {
        let do_challenge = |ty, ch| {
            self.tmpl
                .render(
                    "id_v1_login",
                    &serde_json::json!(
                        {
                            "challenge":
                                format!(r#"
                            <td class="challenge-type">
                                <input type="hidden" name="challenge_type" value="{:?}" />
                                {}
                            </td>
                            <td class="challenge-content">{}</td>
                            "#,
                                    to_present, ty, ch),
                            "redirect": redirect,
                            "error_msg": error_msg.iter().collect::<Vec<_>>()
                        }
                    ),
                )
                .unwrap()
        };

        response.set_content_type("text/html");

        match to_present {
            schema::AuthChallengeType::Username => {
                response.set_body(do_challenge(
                    "Username",
                    r#"<input name="challenge" type="text" autofocus />"#,
                ));
            }
            schema::AuthChallengeType::Password => {
                response.set_body(do_challenge(
                    "Password",
                    r#"<input name="challenge" type="password" autofocus />"#,
                ));
            }
            schema::AuthChallengeType::Totp => {
                response.set_body(do_challenge(
                    "Authenticator code",
                    r#"<input name="challenge" type="text" autofocus />"#,
                ));
            }
            _ => todo!(),
        }

        response
    }
}

async fn v1_login(req: Request) -> tide::Result<tide::Response> {
    log::info!("in v1_login");
    let mut response = tide::Response::builder(200).build();

    let shelper = SessionHelper::new(&req);

    let realm = shelper.get_realm()?;
    let (session, cookie) = shelper.get_or_build_session(&req)?;
    if let Some(c) = cookie {
        response.insert_cookie(c)
    }

    let auth = shelper.get_auth_for_session(realm.id(), &session);

    #[derive(serde::Deserialize)]
    struct LoginQuery {
        redirect: Option<String>,
    }

    let query: LoginQuery = req.query().unwrap();

    Ok(shelper.render_login_from_auth(
        response,
        query.redirect.unwrap_or_else(|| "../..".to_string()),
        auth.map(Stored::wrapped),
        None,
    ))
}

async fn v1_login_post(mut req: Request) -> tide::Result<tide::Response> {
    let mut response = tide::Response::builder(200).build();

    #[derive(Deserialize)]
    struct ResponseBody {
        challenge_type: String,
        challenge: String,
        reset: Option<String>,
        redirect: String,
    }

    let body: ResponseBody = req.body_form().await?;

    let shelper = SessionHelper::new(&req);

    let realm = shelper.get_realm()?;
    let (session, cookie) = shelper.get_or_build_session(&req)?;
    if let Some(c) = cookie {
        response.insert_cookie(c)
    }

    let mut auth = shelper.get_auth_for_session(realm.id(), &session);

    // check if a login reset was requested; if so, we start again from the top
    if body.reset.is_some() {
        shelper.destroy_auth(realm.id(), &session)?;
        // TODO: include original redirect URL here
        return Ok(tide::Redirect::new("login").into());
    }

    use schema::AuthChallengeType as ChallengeType;

    let challenge: schema::AuthChallengeType = match body.challenge_type.as_str() {
        "Username" => ChallengeType::Username,
        "Password" => ChallengeType::Password,
        "TOTP" => ChallengeType::Totp,
        _ => Err(tide::Error::from_str(400, "Unknown challenge type"))?,
    };

    let mut error = None;

    // check that the response matches what we're expecting next
    let to_be_presented: Option<schema::AuthChallengeType> = match &auth {
        None => Some(schema::AuthChallengeType::Username),
        Some(auth) => auth.pending_challenges.as_ref().first().copied(),
    };

    if to_be_presented != Some(challenge) {
        Err(tide::Error::from_str(400, "Unexpected challenge type"))?
    }

    match challenge {
        // handle the username challenge specially because this sets up the pending challenges.
        ChallengeType::Username => {
            shelper.destroy_auth(realm.id(), &session)?;

            let user = realm
                .users
                .with(schema::User::Username, &body.challenge)
                .first()
                .get()?;
            if user.is_none() {
                error = Some(format!("No such user {}", body.challenge));
            } else {
                let user = user.unwrap();

                let has_totp = user
                    .auth
                    .with(
                        schema::AuthChallenge::ChallengeType,
                        microrm::schema::Serialized::from(schema::AuthChallengeType::Totp),
                    )
                    .count()?
                    > 0;

                // TODO: support more flows than just username,password[,totp]
                auth = Some(
                    session.auth.insert_and_return(schema::SessionAuth {
                        realm: realm.id(),
                        user: None,
                        pending_user: Some(user.id()),
                        pending_challenges: if has_totp {
                            vec![
                                schema::AuthChallengeType::Password,
                                schema::AuthChallengeType::Totp,
                            ]
                        } else {
                            vec![schema::AuthChallengeType::Password]
                        }
                        .into(),
                    })?,
                );
                // auth = Some(session.auth.with(id, id).first().get()?.expect("can't re-get just-added entity"));
            }
        }
        ct => {
            if let Some(auth) = auth.as_mut() {
                if let Some(user_id) = auth.pending_user {
                    let user = realm
                        .users
                        .with_id(user_id)
                        .get()?
                        .ok_or(UIDCError::Abort("session auth refers to nonexistent user"))?;

                    let verification = user.verify_challenge_by_type(ct, body.challenge.as_bytes());

                    match verification {
                        Ok(true) => {
                            auth.pending_challenges.as_mut().remove(0);

                            // if we're done with the last challenge, mark us as logged in
                            if auth.pending_challenges.as_ref().is_empty() {
                                auth.user = auth.pending_user.take();
                            }

                            auth.sync()?;
                        }
                        Ok(false) => {
                            error = Some("Incorrect response. Please try again".into());
                        }
                        Err(_) => {
                            error = Some("Internal error. Please contact an administrator.".into());
                        }
                    }
                } else {
                    error = Some("User is not configured correctly: either it was deleted or it lacks a required authentication challenge type. Please contact an administrator.".into());
                }
            } else {
                error = Some("Please restart login process.".into());
            }
        }
    };

    Ok(shelper.render_login_from_auth(response, body.redirect, auth.map(Stored::wrapped), error))
}

async fn v1_logout(req: Request) -> tide::Result<tide::Response> {
    let shelper = SessionHelper::new(&req);

    #[derive(serde::Deserialize)]
    struct LogoutQuery {
        redirect: Option<String>,
    }

    let query: LogoutQuery = req.query().unwrap();

    let realm = shelper.get_realm()?;
    shelper
        .get_session(&req)
        .map(|sid| shelper.destroy_auth(realm.id(), &sid));
    Ok(tide::Redirect::new(query.redirect.unwrap_or_else(|| "../..".into())).into())
}

pub(super) fn session_v1_server(mut route: tide::Route<super::ServerStateWrapper>) {
    route.at("login").get(v1_login).post(v1_login_post);
    route.at("logout").get(v1_logout);
}